001/**
002 *
003 * Copyright © 2015-2021 Florian Schmaus
004 *
005 * Licensed under the Apache License, Version 2.0 (the "License");
006 * you may not use this file except in compliance with the License.
007 * You may obtain a copy of the License at
008 *
009 *     http://www.apache.org/licenses/LICENSE-2.0
010 *
011 * Unless required by applicable law or agreed to in writing, software
012 * distributed under the License is distributed on an "AS IS" BASIS,
013 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
014 * See the License for the specific language governing permissions and
015 * limitations under the License.
016 */
017package org.jivesoftware.smack.util;
018
019import java.util.ArrayList;
020import java.util.Collections;
021import java.util.LinkedHashMap;
022import java.util.LinkedHashSet;
023import java.util.List;
024import java.util.Map;
025import java.util.Set;
026
027/**
028 * A lightweight implementation of a MultiMap, that is a Map that is able to hold multiple values for every key.
029 * <p>
030 * This MultiMap uses a {@link LinkedHashMap} together with a {@link ArrayList} in order to keep the order of its entries.
031 * </p>
032 *
033 * @param <K> the type of the keys the map uses.
034 * @param <V> the type of the values the map uses.
035 */
036public class MultiMap<K, V> {
037
038    /**
039     * The constant value {@value}.
040     */
041    public static final int DEFAULT_MAP_SIZE = 6;
042
043    private static final int ENTRY_LIST_SIZE = 3;
044
045    private final Map<K, List<V>> map;
046
047    /**
048     * Constructs a new MultiMap with a initial capacity of {@link #DEFAULT_MAP_SIZE}.
049     */
050    public MultiMap() {
051        this(DEFAULT_MAP_SIZE);
052    }
053
054    /**
055     * Constructs a new MultiMap.
056     *
057     * @param size the initial capacity.
058     */
059    public MultiMap(int size) {
060        this(new LinkedHashMap<>(size));
061    }
062
063    private MultiMap(Map<K, List<V>> map) {
064        this.map = map;
065    }
066
067    public int size() {
068        int size = 0;
069        for (List<V> list : map.values()) {
070            size += list.size();
071        }
072        return size;
073    }
074
075    public boolean isEmpty() {
076        return map.isEmpty();
077    }
078
079    public boolean containsKey(K key) {
080        return map.containsKey(key);
081    }
082
083    public boolean containsValue(V value) {
084        for (List<V> list : map.values()) {
085            if (list.contains(value)) {
086                return true;
087            }
088        }
089        return false;
090    }
091
092    /**
093     * Get the first value for the given key, or <code>null</code> if there are no entries.
094     *
095     * @param key TODO javadoc me please
096     * @return the first value or null.
097     */
098    public V getFirst(K key) {
099        List<V> res = getAll(key);
100        if (res.isEmpty()) {
101            return null;
102        } else {
103            return res.iterator().next();
104        }
105    }
106
107    /**
108     * Get all values for the given key. Returns the empty set if there are none.
109     * <p>
110     * Changes to the returned set will update the underlying MultiMap if the return set is not empty.
111     * </p>
112     *
113     * @param key TODO javadoc me please
114     * @return all values for the given key.
115     */
116    public List<V> getAll(K key) {
117        List<V> res = map.get(key);
118        if (res == null) {
119            res = Collections.emptyList();
120        }
121        return res;
122    }
123
124    public boolean put(K key, V value) {
125        return putInternal(key, list -> list.add(value));
126    }
127
128    public boolean putFirst(K key, V value) {
129        return putInternal(key, list -> list.add(0, value));
130    }
131
132    private boolean putInternal(K key, Consumer<List<V>> valueListConsumer) {
133        boolean keyExisted;
134        List<V> list = map.get(key);
135        if (list == null) {
136            list = new ArrayList<>(ENTRY_LIST_SIZE);
137            map.put(key, list);
138            keyExisted = false;
139        } else {
140            keyExisted = true;
141        }
142
143        valueListConsumer.accept(list);
144
145        return keyExisted;
146    }
147
148    /**
149     * Removes all mappings for the given key and returns the first value if there where mappings or <code>null</code> if not.
150     *
151     * @param key TODO javadoc me please
152     * @return the first value of the given key or null.
153     */
154    public V remove(K key) {
155        List<V> res = map.remove(key);
156        if (res == null) {
157            return null;
158        }
159        assert !res.isEmpty();
160        return res.iterator().next();
161    }
162
163    /**
164     * Remove the mapping of the given key to the value.
165     * <p>
166     * Returns <code>true</code> if the mapping was removed and <code>false</code> if the mapping did not exist.
167     * </p>
168     *
169     * @param key TODO javadoc me please
170     * @param value TODO javadoc me please
171     * @return true if the mapping was removed, false otherwise.
172     */
173    public boolean removeOne(K key, V value) {
174        List<V> list = map.get(key);
175        if (list == null) {
176            return false;
177        }
178        boolean res = list.remove(value);
179        if (list.isEmpty()) {
180            // Remove the key also if the value set is now empty
181            map.remove(key);
182        }
183        return res;
184    }
185
186    /**
187     * Remove the given number of values for a given key. May return less values then requested.
188     *
189     * @param key the key to remove from.
190     * @param num the number of values to remove.
191     * @return a list of the removed values.
192     * @since 4.4.0
193     */
194    public List<V> remove(K key, int num) {
195        List<V> values = map.get(key);
196        if (values == null) {
197            return Collections.emptyList();
198        }
199
200        final int resultSize = values.size() > num ? num : values.size();
201        final List<V> result = new ArrayList<>(resultSize);
202        for (int i = 0; i < resultSize; i++) {
203            result.add(values.get(0));
204        }
205
206        if (values.isEmpty()) {
207            map.remove(key);
208        }
209
210        return result;
211    }
212
213    public void putAll(Map<? extends K, ? extends V> map) {
214        for (Map.Entry<? extends K, ? extends V> entry : map.entrySet()) {
215            put(entry.getKey(), entry.getValue());
216        }
217    }
218
219    public void clear() {
220        map.clear();
221    }
222
223    public Set<K> keySet() {
224        return map.keySet();
225    }
226
227    /**
228     * Returns a new list containing all values of this multi map.
229     *
230     * @return a new list with all values.
231     */
232    public List<V> values() {
233        List<V> values = new ArrayList<>(size());
234        for (List<V> list : map.values()) {
235            values.addAll(list);
236        }
237        return values;
238    }
239
240    public Set<Map.Entry<K, V>> entrySet() {
241        Set<Map.Entry<K, V>> entrySet = new LinkedHashSet<>(size());
242        for (Map.Entry<K, List<V>> entries : map.entrySet()) {
243            K key = entries.getKey();
244            for (V value : entries.getValue()) {
245                entrySet.add(new SimpleMapEntry<>(key, value));
246            }
247        }
248        return entrySet;
249    }
250
251    public MultiMap<K, V> asUnmodifiableMultiMap() {
252        LinkedHashMap<K, List<V>> mapCopy = new LinkedHashMap<>(map.size());
253        for (Map.Entry<K, List<V>> entry : map.entrySet()) {
254            K key = entry.getKey();
255            List<V> values = entry.getValue();
256
257            mapCopy.put(key, Collections.unmodifiableList(values));
258        }
259
260        return new MultiMap<K, V>(Collections.unmodifiableMap(mapCopy));
261    }
262
263    @Override
264    public MultiMap<K, V> clone() {
265        Map<K, List<V>> clonedMap = new LinkedHashMap<>(map.size());
266
267        // TODO: Use Map.forEach() once Smack's minimum Android API is 24 or higher.
268        for (Map.Entry<K, List<V>> entry : map.entrySet()) {
269            List<V> clonedList = CollectionUtil.newListWith(entry.getValue());
270            clonedMap.put(entry.getKey(), clonedList);
271        }
272
273        return new MultiMap<>(clonedMap);
274    }
275
276    private static final class SimpleMapEntry<K, V> implements Map.Entry<K, V> {
277
278        private final K key;
279        private V value;
280
281        private SimpleMapEntry(K key, V value) {
282            this.key = key;
283            this.value = value;
284        }
285
286        @Override
287        public K getKey() {
288            return key;
289        }
290
291        @Override
292        public V getValue() {
293            return value;
294        }
295
296        @Override
297        public V setValue(V value) {
298            V tmp = this.value;
299            this.value = value;
300            return tmp;
301        }
302    }
303
304}