# MappedListDecorator.java

    package org.thewonderlemming.c4plantuml.graphml.model;

    import java.util.AbstractMap.SimpleEntry;
    import java.util.ArrayList;
    import java.util.Collection;
    import java.util.HashMap;
    import java.util.Iterator;
    import java.util.List;
    import java.util.ListIterator;
    import java.util.Map;
    import java.util.Map.Entry;
    import java.util.Set;
    import java.util.function.Function;
    import java.util.stream.Collectors;

    /**
     * A decorator to the {@link List} interface, that relies on an internal {@link Map} to track existing data by the key
     * {@code K}.
     * <p>
     * The "raison d'être" of that class is to simplify the JAXB data model by putting properties that appear to be lists
     * and managing the add/replace feature internally and at the same location.
     *
     * @author thewonderlemming
     *
     * @param <M> the type of data to manage.
     * @param <K> the type of the data key.
     */
    public class MappedListDecorator<M, K> extends ArrayList<M> implements List<M> {

        private static final long serialVersionUID = -7298751381485516154L;

        /**
         * The lookup map that maps the data to its key. It is protected for testing purposes.
         */
        protected final transient Map<K, M> dataLookup = new HashMap<>();

        private final transient List<M> decoratedList;

        private final transient Function<M, K> getModelId;

        private final Class<M> modelType;

        /**
         * Default constructor.
         *
         * @param modelType the type of the model to manage.
         * @param model the decorated model list.
         * @param getModelId a function that provides a model key of type {@code K} given a model of type {@code M}.
         */
        public MappedListDecorator(final Class<M> modelType, final List<M> model,
            final Function<M, K> getModelId) {

            this.decoratedList = model;
            this.modelType = modelType;
            this.getModelId = getModelId;
            this.dataLookup
                .putAll(
                    model
                        .stream()
                            .map(d -> new SimpleEntry<K, M>(getModelId.apply(d), d))
                            .collect(Collectors.toMap(Entry::getKey, Entry::getValue)));
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public void add(final int index, final M element) {

            this.decoratedList.add(index, element);
            this.dataLookup.put(this.getModelId.apply(element), element);
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public boolean add(final M e) {

            final boolean result = this.decoratedList.add(e);

            if (result) {
                this.dataLookup.put(this.getModelId.apply(e), e);
            }

            return result;
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public boolean addAll(final Collection<? extends M> c) {

            final boolean result = this.decoratedList.addAll(c);

            if (result) {
                c.forEach(item -> {
                    final K key = this.getModelId.apply(item);
                    this.dataLookup.put(key, item);
                });
            }

            return result;
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public boolean addAll(final int index, final Collection<? extends M> c) {

            final boolean result = this.decoratedList.addAll(index, c);

            if (result) {
                c.forEach(item -> {
                    final K key = this.getModelId.apply(item);
                    this.dataLookup.put(key, item);
                });
            }

            return result;
        }

        /**
         * Adds or replaces the given {@code model}, given that its key {@code K} is not yet registered in the current
         * mapping or not.
         *
         * @param model the model to add or update.
         */
        public void addOrReplaceData(final M model) {

            final K key = this.getModelId.apply(model);

            if (this.dataLookup.containsKey(key)) {

                this.decoratedList.remove(this.dataLookup.get(key));
                this.dataLookup.remove(key);
            }

            add(model);
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public void clear() {

            this.decoratedList.clear();
            this.dataLookup.clear();
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public boolean contains(final Object o) {
            return this.decoratedList.contains(o);
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public boolean containsAll(final Collection<?> c) {
            return this.decoratedList.containsAll(c);
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public boolean equals(final Object o) {
            return this.decoratedList.equals(o);
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public M get(final int index) {
            return this.decoratedList.get(index);
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public int hashCode() {
            return this.decoratedList.hashCode();
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public int indexOf(final Object o) {
            return this.decoratedList.indexOf(o);
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public boolean isEmpty() {
            return this.decoratedList.isEmpty();
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public Iterator<M> iterator() {
            return this.decoratedList.iterator();
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public int lastIndexOf(final Object o) {
            return this.decoratedList.lastIndexOf(o);
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public ListIterator<M> listIterator() {
            return this.decoratedList.listIterator();
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public ListIterator<M> listIterator(final int index) {
            return this.decoratedList.listIterator(index);
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public M remove(final int index) {

            final M results = this.decoratedList.remove(index);
            final K key = this.getModelId.apply(results);
            this.dataLookup.remove(key);

            return results;
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public boolean remove(final Object o) {

            final boolean results = this.decoratedList.remove(o);

            if (results && this.modelType.isAssignableFrom(o.getClass())) {

                @SuppressWarnings("unchecked")
                final K key = this.getModelId.apply((M) o);
                this.dataLookup.remove(key);
            }

            return results;
        }

        /**
         * {@inheritDoc}
         */
        @SuppressWarnings("unchecked")
        @Override
        public boolean removeAll(final Collection<?> c) {

            final boolean results = this.decoratedList.removeAll(c);

            if (results) {
                c
                    .stream()
                        .filter(item -> this.modelType.isAssignableFrom(item.getClass()))
                        .map(item -> (M) item)
                        .forEach(model -> {

                            final K key = this.getModelId.apply(model);
                            this.dataLookup.remove(key);
                        });
            }

            return results;
        }

        /**
         * {@inheritDoc}
         */
        @SuppressWarnings("unchecked")
        @Override
        public boolean retainAll(final Collection<?> c) {

            final boolean results = this.decoratedList.retainAll(c);

            if (results) {
                final Set<K> keys = c
                    .stream()
                        .filter(item -> this.modelType.isAssignableFrom(item.getClass()))
                        .map(item -> (M) item)
                        .map(this.getModelId::apply)
                        .collect(Collectors.toSet());

                final Set<K> toBeRemoved = this.dataLookup
                    .keySet()
                        .stream()
                        .filter(key -> !keys.contains(key))
                        .collect(Collectors.toSet());

                toBeRemoved.forEach(this.dataLookup::remove);
            }

            return results;
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public M set(final int index, final M element) {

            final M results = this.decoratedList.set(index, element);
            final K key = this.getModelId.apply(element);
            this.dataLookup.put(key, element);

            return results;
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public int size() {
            return this.decoratedList.size();
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public List<M> subList(final int fromIndex, final int toIndex) {
            return this.decoratedList.subList(fromIndex, toIndex);
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public Object[] toArray() {
            return this.decoratedList.toArray();
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public <T> T[] toArray(final T[] a) {
            return this.decoratedList.toArray(a);
        }
    }