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