• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /*
2  * Copyright (C) 2016 The Guava Authors
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.google.common.collect;
18 
19 import static com.google.common.base.Preconditions.checkNotNull;
20 import static java.util.stream.Collectors.collectingAndThen;
21 
22 import com.google.common.annotations.GwtCompatible;
23 import com.google.common.annotations.GwtIncompatible;
24 import com.google.common.base.Preconditions;
25 import java.util.Collection;
26 import java.util.Comparator;
27 import java.util.EnumMap;
28 import java.util.EnumSet;
29 import java.util.LinkedHashMap;
30 import java.util.TreeMap;
31 import java.util.function.BinaryOperator;
32 import java.util.function.Function;
33 import java.util.function.Supplier;
34 import java.util.function.ToIntFunction;
35 import java.util.stream.Collector;
36 import java.util.stream.Collectors;
37 import java.util.stream.Stream;
38 import javax.annotation.CheckForNull;
39 import org.checkerframework.checker.nullness.qual.Nullable;
40 
41 /** Collectors utilities for {@code common.collect} internals. */
42 @GwtCompatible
43 @ElementTypesAreNonnullByDefault
44 final class CollectCollectors {
45 
46   private static final Collector<Object, ?, ImmutableList<Object>> TO_IMMUTABLE_LIST =
47       Collector.of(
48           ImmutableList::builder,
49           ImmutableList.Builder::add,
50           ImmutableList.Builder::combine,
51           ImmutableList.Builder::build);
52 
53   private static final Collector<Object, ?, ImmutableSet<Object>> TO_IMMUTABLE_SET =
54       Collector.of(
55           ImmutableSet::builder,
56           ImmutableSet.Builder::add,
57           ImmutableSet.Builder::combine,
58           ImmutableSet.Builder::build);
59 
60   @GwtIncompatible
61   private static final Collector<Range<Comparable<?>>, ?, ImmutableRangeSet<Comparable<?>>>
62       TO_IMMUTABLE_RANGE_SET =
63           Collector.of(
64               ImmutableRangeSet::builder,
65               ImmutableRangeSet.Builder::add,
66               ImmutableRangeSet.Builder::combine,
67               ImmutableRangeSet.Builder::build);
68 
69   // Lists
70 
71   @SuppressWarnings({"rawtypes", "unchecked"})
toImmutableList()72   static <E> Collector<E, ?, ImmutableList<E>> toImmutableList() {
73     return (Collector) TO_IMMUTABLE_LIST;
74   }
75 
76   // Sets
77 
78   @SuppressWarnings({"rawtypes", "unchecked"})
toImmutableSet()79   static <E> Collector<E, ?, ImmutableSet<E>> toImmutableSet() {
80     return (Collector) TO_IMMUTABLE_SET;
81   }
82 
toImmutableSortedSet( Comparator<? super E> comparator)83   static <E> Collector<E, ?, ImmutableSortedSet<E>> toImmutableSortedSet(
84       Comparator<? super E> comparator) {
85     checkNotNull(comparator);
86     return Collector.of(
87         () -> new ImmutableSortedSet.Builder<E>(comparator),
88         ImmutableSortedSet.Builder::add,
89         ImmutableSortedSet.Builder::combine,
90         ImmutableSortedSet.Builder::build);
91   }
92 
93   @SuppressWarnings({"rawtypes", "unchecked"})
toImmutableEnumSet()94   static <E extends Enum<E>> Collector<E, ?, ImmutableSet<E>> toImmutableEnumSet() {
95     return (Collector) EnumSetAccumulator.TO_IMMUTABLE_ENUM_SET;
96   }
97 
98   private static final class EnumSetAccumulator<E extends Enum<E>> {
99     @SuppressWarnings({"rawtypes", "unchecked"})
100     static final Collector<Enum<?>, ?, ImmutableSet<? extends Enum<?>>> TO_IMMUTABLE_ENUM_SET =
101         (Collector)
102             Collector.<Enum, EnumSetAccumulator, ImmutableSet<?>>of(
103                 EnumSetAccumulator::new,
104                 EnumSetAccumulator::add,
105                 EnumSetAccumulator::combine,
106                 EnumSetAccumulator::toImmutableSet,
107                 Collector.Characteristics.UNORDERED);
108 
109     @CheckForNull private EnumSet<E> set;
110 
add(E e)111     void add(E e) {
112       if (set == null) {
113         set = EnumSet.of(e);
114       } else {
115         set.add(e);
116       }
117     }
118 
combine(EnumSetAccumulator<E> other)119     EnumSetAccumulator<E> combine(EnumSetAccumulator<E> other) {
120       if (this.set == null) {
121         return other;
122       } else if (other.set == null) {
123         return this;
124       } else {
125         this.set.addAll(other.set);
126         return this;
127       }
128     }
129 
toImmutableSet()130     ImmutableSet<E> toImmutableSet() {
131       return (set == null) ? ImmutableSet.<E>of() : ImmutableEnumSet.asImmutable(set);
132     }
133   }
134 
135   @GwtIncompatible
136   @SuppressWarnings({"rawtypes", "unchecked"})
137   static <E extends Comparable<? super E>>
toImmutableRangeSet()138       Collector<Range<E>, ?, ImmutableRangeSet<E>> toImmutableRangeSet() {
139     return (Collector) TO_IMMUTABLE_RANGE_SET;
140   }
141 
142   // Multisets
143 
toImmutableMultiset( Function<? super T, ? extends E> elementFunction, ToIntFunction<? super T> countFunction)144   static <T extends @Nullable Object, E> Collector<T, ?, ImmutableMultiset<E>> toImmutableMultiset(
145       Function<? super T, ? extends E> elementFunction, ToIntFunction<? super T> countFunction) {
146     checkNotNull(elementFunction);
147     checkNotNull(countFunction);
148     return Collector.of(
149         LinkedHashMultiset::create,
150         (multiset, t) ->
151             multiset.add(checkNotNull(elementFunction.apply(t)), countFunction.applyAsInt(t)),
152         (multiset1, multiset2) -> {
153           multiset1.addAll(multiset2);
154           return multiset1;
155         },
156         (Multiset<E> multiset) -> ImmutableMultiset.copyFromEntries(multiset.entrySet()));
157   }
158 
159   static <T extends @Nullable Object, E extends @Nullable Object, M extends Multiset<E>>
toMultiset( Function<? super T, E> elementFunction, ToIntFunction<? super T> countFunction, Supplier<M> multisetSupplier)160       Collector<T, ?, M> toMultiset(
161           Function<? super T, E> elementFunction,
162           ToIntFunction<? super T> countFunction,
163           Supplier<M> multisetSupplier) {
164     checkNotNull(elementFunction);
165     checkNotNull(countFunction);
166     checkNotNull(multisetSupplier);
167     return Collector.of(
168         multisetSupplier,
169         (ms, t) -> ms.add(elementFunction.apply(t), countFunction.applyAsInt(t)),
170         (ms1, ms2) -> {
171           ms1.addAll(ms2);
172           return ms1;
173         });
174   }
175 
176   // Maps
177 
178   static <T extends @Nullable Object, K, V> Collector<T, ?, ImmutableMap<K, V>> toImmutableMap(
179       Function<? super T, ? extends K> keyFunction,
180       Function<? super T, ? extends V> valueFunction) {
181     checkNotNull(keyFunction);
182     checkNotNull(valueFunction);
183     return Collector.of(
184         ImmutableMap.Builder<K, V>::new,
185         (builder, input) -> builder.put(keyFunction.apply(input), valueFunction.apply(input)),
186         ImmutableMap.Builder::combine,
187         ImmutableMap.Builder::build);
188   }
189 
190   public static <T extends @Nullable Object, K, V>
191       Collector<T, ?, ImmutableMap<K, V>> toImmutableMap(
192           Function<? super T, ? extends K> keyFunction,
193           Function<? super T, ? extends V> valueFunction,
194           BinaryOperator<V> mergeFunction) {
195     checkNotNull(keyFunction);
196     checkNotNull(valueFunction);
197     checkNotNull(mergeFunction);
198     return collectingAndThen(
199         Collectors.toMap(keyFunction, valueFunction, mergeFunction, LinkedHashMap::new),
200         ImmutableMap::copyOf);
201   }
202 
203   static <T extends @Nullable Object, K, V>
204       Collector<T, ?, ImmutableSortedMap<K, V>> toImmutableSortedMap(
205           Comparator<? super K> comparator,
206           Function<? super T, ? extends K> keyFunction,
207           Function<? super T, ? extends V> valueFunction) {
208     checkNotNull(comparator);
209     checkNotNull(keyFunction);
210     checkNotNull(valueFunction);
211     /*
212      * We will always fail if there are duplicate keys, and the keys are always sorted by
213      * the Comparator, so the entries can come in an arbitrary order -- so we report UNORDERED.
214      */
215     return Collector.of(
216         () -> new ImmutableSortedMap.Builder<K, V>(comparator),
217         (builder, input) -> builder.put(keyFunction.apply(input), valueFunction.apply(input)),
218         ImmutableSortedMap.Builder::combine,
219         ImmutableSortedMap.Builder::build,
220         Collector.Characteristics.UNORDERED);
221   }
222 
223   static <T extends @Nullable Object, K, V>
224       Collector<T, ?, ImmutableSortedMap<K, V>> toImmutableSortedMap(
225           Comparator<? super K> comparator,
226           Function<? super T, ? extends K> keyFunction,
227           Function<? super T, ? extends V> valueFunction,
228           BinaryOperator<V> mergeFunction) {
229     checkNotNull(comparator);
230     checkNotNull(keyFunction);
231     checkNotNull(valueFunction);
232     checkNotNull(mergeFunction);
233     return collectingAndThen(
234         Collectors.toMap(
235             keyFunction, valueFunction, mergeFunction, () -> new TreeMap<K, V>(comparator)),
236         ImmutableSortedMap::copyOfSorted);
237   }
238 
239   static <T extends @Nullable Object, K, V> Collector<T, ?, ImmutableBiMap<K, V>> toImmutableBiMap(
240       Function<? super T, ? extends K> keyFunction,
241       Function<? super T, ? extends V> valueFunction) {
242     checkNotNull(keyFunction);
243     checkNotNull(valueFunction);
244     return Collector.of(
245         ImmutableBiMap.Builder<K, V>::new,
246         (builder, input) -> builder.put(keyFunction.apply(input), valueFunction.apply(input)),
247         ImmutableBiMap.Builder::combine,
248         ImmutableBiMap.Builder::build,
249         new Collector.Characteristics[0]);
250   }
251 
252   static <T extends @Nullable Object, K extends Enum<K>, V>
253       Collector<T, ?, ImmutableMap<K, V>> toImmutableEnumMap(
254           Function<? super T, ? extends K> keyFunction,
255           Function<? super T, ? extends V> valueFunction) {
256     checkNotNull(keyFunction);
257     checkNotNull(valueFunction);
258     return Collector.of(
259         () ->
260             new EnumMapAccumulator<K, V>(
261                 (v1, v2) -> {
262                   throw new IllegalArgumentException("Multiple values for key: " + v1 + ", " + v2);
263                 }),
264         (accum, t) -> {
265           /*
266            * We assign these to variables before calling checkNotNull to work around a bug in our
267            * nullness checker.
268            */
269           K key = keyFunction.apply(t);
270           V newValue = valueFunction.apply(t);
271           accum.put(
272               checkNotNull(key, "Null key for input %s", t),
273               checkNotNull(newValue, "Null value for input %s", t));
274         },
275         EnumMapAccumulator::combine,
276         EnumMapAccumulator::toImmutableMap,
277         Collector.Characteristics.UNORDERED);
278   }
279 
280   static <T extends @Nullable Object, K extends Enum<K>, V>
281       Collector<T, ?, ImmutableMap<K, V>> toImmutableEnumMap(
282           Function<? super T, ? extends K> keyFunction,
283           Function<? super T, ? extends V> valueFunction,
284           BinaryOperator<V> mergeFunction) {
285     checkNotNull(keyFunction);
286     checkNotNull(valueFunction);
287     checkNotNull(mergeFunction);
288     // not UNORDERED because we don't know if mergeFunction is commutative
289     return Collector.of(
290         () -> new EnumMapAccumulator<K, V>(mergeFunction),
291         (accum, t) -> {
292           /*
293            * We assign these to variables before calling checkNotNull to work around a bug in our
294            * nullness checker.
295            */
296           K key = keyFunction.apply(t);
297           V newValue = valueFunction.apply(t);
298           accum.put(
299               checkNotNull(key, "Null key for input %s", t),
300               checkNotNull(newValue, "Null value for input %s", t));
301         },
302         EnumMapAccumulator::combine,
303         EnumMapAccumulator::toImmutableMap);
304   }
305 
306   private static class EnumMapAccumulator<K extends Enum<K>, V> {
307     private final BinaryOperator<V> mergeFunction;
308     @CheckForNull private EnumMap<K, V> map = null;
309 
310     EnumMapAccumulator(BinaryOperator<V> mergeFunction) {
311       this.mergeFunction = mergeFunction;
312     }
313 
314     void put(K key, V value) {
315       if (map == null) {
316         map = new EnumMap<>(key.getDeclaringClass());
317       }
318       map.merge(key, value, mergeFunction);
319     }
320 
321     EnumMapAccumulator<K, V> combine(EnumMapAccumulator<K, V> other) {
322       if (this.map == null) {
323         return other;
324       } else if (other.map == null) {
325         return this;
326       } else {
327         other.map.forEach(this::put);
328         return this;
329       }
330     }
331 
332     ImmutableMap<K, V> toImmutableMap() {
333       return (map == null) ? ImmutableMap.<K, V>of() : ImmutableEnumMap.asImmutable(map);
334     }
335   }
336 
337   @GwtIncompatible
338   static <T extends @Nullable Object, K extends Comparable<? super K>, V>
339       Collector<T, ?, ImmutableRangeMap<K, V>> toImmutableRangeMap(
340           Function<? super T, Range<K>> keyFunction,
341           Function<? super T, ? extends V> valueFunction) {
342     checkNotNull(keyFunction);
343     checkNotNull(valueFunction);
344     return Collector.of(
345         ImmutableRangeMap::<K, V>builder,
346         (builder, input) -> builder.put(keyFunction.apply(input), valueFunction.apply(input)),
347         ImmutableRangeMap.Builder::combine,
348         ImmutableRangeMap.Builder::build);
349   }
350 
351   // Multimaps
352 
353   static <T extends @Nullable Object, K, V>
354       Collector<T, ?, ImmutableListMultimap<K, V>> toImmutableListMultimap(
355           Function<? super T, ? extends K> keyFunction,
356           Function<? super T, ? extends V> valueFunction) {
357     checkNotNull(keyFunction, "keyFunction");
358     checkNotNull(valueFunction, "valueFunction");
359     return Collector.of(
360         ImmutableListMultimap::<K, V>builder,
361         (builder, t) -> builder.put(keyFunction.apply(t), valueFunction.apply(t)),
362         ImmutableListMultimap.Builder::combine,
363         ImmutableListMultimap.Builder::build);
364   }
365 
366   static <T extends @Nullable Object, K, V>
367       Collector<T, ?, ImmutableListMultimap<K, V>> flatteningToImmutableListMultimap(
368           Function<? super T, ? extends K> keyFunction,
369           Function<? super T, ? extends Stream<? extends V>> valuesFunction) {
370     checkNotNull(keyFunction);
371     checkNotNull(valuesFunction);
372     return collectingAndThen(
373         flatteningToMultimap(
374             input -> checkNotNull(keyFunction.apply(input)),
375             input -> valuesFunction.apply(input).peek(Preconditions::checkNotNull),
376             MultimapBuilder.linkedHashKeys().arrayListValues()::<K, V>build),
377         ImmutableListMultimap::copyOf);
378   }
379 
380   static <T extends @Nullable Object, K, V>
381       Collector<T, ?, ImmutableSetMultimap<K, V>> toImmutableSetMultimap(
382           Function<? super T, ? extends K> keyFunction,
383           Function<? super T, ? extends V> valueFunction) {
384     checkNotNull(keyFunction, "keyFunction");
385     checkNotNull(valueFunction, "valueFunction");
386     return Collector.of(
387         ImmutableSetMultimap::<K, V>builder,
388         (builder, t) -> builder.put(keyFunction.apply(t), valueFunction.apply(t)),
389         ImmutableSetMultimap.Builder::combine,
390         ImmutableSetMultimap.Builder::build);
391   }
392 
393   static <T extends @Nullable Object, K, V>
394       Collector<T, ?, ImmutableSetMultimap<K, V>> flatteningToImmutableSetMultimap(
395           Function<? super T, ? extends K> keyFunction,
396           Function<? super T, ? extends Stream<? extends V>> valuesFunction) {
397     checkNotNull(keyFunction);
398     checkNotNull(valuesFunction);
399     return collectingAndThen(
400         flatteningToMultimap(
401             input -> checkNotNull(keyFunction.apply(input)),
402             input -> valuesFunction.apply(input).peek(Preconditions::checkNotNull),
403             MultimapBuilder.linkedHashKeys().linkedHashSetValues()::<K, V>build),
404         ImmutableSetMultimap::copyOf);
405   }
406 
407   static <
408           T extends @Nullable Object,
409           K extends @Nullable Object,
410           V extends @Nullable Object,
411           M extends Multimap<K, V>>
412       Collector<T, ?, M> toMultimap(
413           Function<? super T, ? extends K> keyFunction,
414           Function<? super T, ? extends V> valueFunction,
415           Supplier<M> multimapSupplier) {
416     checkNotNull(keyFunction);
417     checkNotNull(valueFunction);
418     checkNotNull(multimapSupplier);
419     return Collector.of(
420         multimapSupplier,
421         (multimap, input) -> multimap.put(keyFunction.apply(input), valueFunction.apply(input)),
422         (multimap1, multimap2) -> {
423           multimap1.putAll(multimap2);
424           return multimap1;
425         });
426   }
427 
428   static <
429           T extends @Nullable Object,
430           K extends @Nullable Object,
431           V extends @Nullable Object,
432           M extends Multimap<K, V>>
433       Collector<T, ?, M> flatteningToMultimap(
434           Function<? super T, ? extends K> keyFunction,
435           Function<? super T, ? extends Stream<? extends V>> valueFunction,
436           Supplier<M> multimapSupplier) {
437     checkNotNull(keyFunction);
438     checkNotNull(valueFunction);
439     checkNotNull(multimapSupplier);
440     return Collector.of(
441         multimapSupplier,
442         (multimap, input) -> {
443           K key = keyFunction.apply(input);
444           Collection<V> valuesForKey = multimap.get(key);
445           valueFunction.apply(input).forEachOrdered(valuesForKey::add);
446         },
447         (multimap1, multimap2) -> {
448           multimap1.putAll(multimap2);
449           return multimap1;
450         });
451   }
452 }
453