com.google.common.collect.CollectCollectors Maven / Gradle / Ivy

Go to download
/*
 * Copyright (C) 2016 The Guava Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.google.common.collect;

import static com.google.common.base.Preconditions.checkNotNull;

import com.google.common.annotations.Beta;
import com.google.common.annotations.GwtCompatible;
import com.google.common.annotations.GwtIncompatible;
import com.google.common.base.Preconditions;
import java.util.Collection;
import java.util.Comparator;
import java.util.EnumMap;
import java.util.EnumSet;
import java.util.LinkedHashMap;
import java.util.TreeMap;
import java.util.function.BinaryOperator;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.function.ToIntFunction;
import java.util.stream.Collector;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.checkerframework.checker.nullness.qual.Nullable;

/** Collectors utilities for {@code common.collect} internals. */
@GwtCompatible
final class CollectCollectors {

  private static final Collector> TO_IMMUTABLE_LIST =
      Collector.of(
          ImmutableList::builder,
          ImmutableList.Builder::add,
          ImmutableList.Builder::combine,
          ImmutableList.Builder::build);

  private static final Collector> TO_IMMUTABLE_SET =
      Collector.of(
          ImmutableSet::builder,
          ImmutableSet.Builder::add,
          ImmutableSet.Builder::combine,
          ImmutableSet.Builder::build);

  @GwtIncompatible
  private static final Collector>, ?, ImmutableRangeSet>>
      TO_IMMUTABLE_RANGE_SET =
          Collector.of(
              ImmutableRangeSet::builder,
              ImmutableRangeSet.Builder::add,
              ImmutableRangeSet.Builder::combine,
              ImmutableRangeSet.Builder::build);

  // Lists

  @SuppressWarnings({"rawtypes", "unchecked"})
  static  Collector> toImmutableList() {
    return (Collector) TO_IMMUTABLE_LIST;
  }

  // Sets

  @SuppressWarnings({"rawtypes", "unchecked"})
  static  Collector> toImmutableSet() {
    return (Collector) TO_IMMUTABLE_SET;
  }

  static  Collector> toImmutableSortedSet(
      Comparator comparator) {
    checkNotNull(comparator);
    return Collector.of(
        () -> new ImmutableSortedSet.Builder(comparator),
        ImmutableSortedSet.Builder::add,
        ImmutableSortedSet.Builder::combine,
        ImmutableSortedSet.Builder::build);
  }

  @SuppressWarnings({"rawtypes", "unchecked"})
  static > Collector> toImmutableEnumSet() {
    return (Collector) EnumSetAccumulator.TO_IMMUTABLE_ENUM_SET;
  }

  private static final class EnumSetAccumulator> {
    @SuppressWarnings({"rawtypes", "unchecked"})
    static final Collector, ?, ImmutableSet>> TO_IMMUTABLE_ENUM_SET =
        (Collector)
            Collector.>of(
                EnumSetAccumulator::new,
                EnumSetAccumulator::add,
                EnumSetAccumulator::combine,
                EnumSetAccumulator::toImmutableSet,
                Collector.Characteristics.UNORDERED);

    private @Nullable EnumSet set;

    void add(E e) {
      if (set == null) {
        set = EnumSet.of(e);
      } else {
        set.add(e);
      }
    }

    EnumSetAccumulator combine(EnumSetAccumulator other) {
      if (this.set == null) {
        return other;
      } else if (other.set == null) {
        return this;
      } else {
        this.set.addAll(other.set);
        return this;
      }
    }

    ImmutableSet toImmutableSet() {
      return (set == null) ? ImmutableSet.of() : ImmutableEnumSet.asImmutable(set);
    }
  }

  @GwtIncompatible
  @SuppressWarnings({"rawtypes", "unchecked"})
  static >
      Collector, ?, ImmutableRangeSet> toImmutableRangeSet() {
    return (Collector) TO_IMMUTABLE_RANGE_SET;
  }

  // Multisets

  static  Collector> toImmutableMultiset(
      Function elementFunction, ToIntFunction countFunction) {
    checkNotNull(elementFunction);
    checkNotNull(countFunction);
    return Collector.of(
        LinkedHashMultiset::create,
        (multiset, t) ->
            multiset.add(checkNotNull(elementFunction.apply(t)), countFunction.applyAsInt(t)),
        (multiset1, multiset2) -> {
          multiset1.addAll(multiset2);
          return multiset1;
        },
        (Multiset multiset) -> ImmutableMultiset.copyFromEntries(multiset.entrySet()));
  }

  static > Collector toMultiset(
      Function elementFunction,
      ToIntFunction countFunction,
      Supplier multisetSupplier) {
    checkNotNull(elementFunction);
    checkNotNull(countFunction);
    checkNotNull(multisetSupplier);
    return Collector.of(
        multisetSupplier,
        (ms, t) -> ms.add(elementFunction.apply(t), countFunction.applyAsInt(t)),
        (ms1, ms2) -> {
          ms1.addAll(ms2);
          return ms1;
        });
  }

  // Maps

  static  Collector> toImmutableMap(
      Function keyFunction,
      Function valueFunction) {
    checkNotNull(keyFunction);
    checkNotNull(valueFunction);
    return Collector.of(
        ImmutableMap.Builder::new,
        (builder, input) -> builder.put(keyFunction.apply(input), valueFunction.apply(input)),
        ImmutableMap.Builder::combine,
        ImmutableMap.Builder::build);
  }

  public static  Collector> toImmutableMap(
      Function keyFunction,
      Function valueFunction,
      BinaryOperator mergeFunction) {
    checkNotNull(keyFunction);
    checkNotNull(valueFunction);
    checkNotNull(mergeFunction);
    return Collectors.collectingAndThen(
        Collectors.toMap(keyFunction, valueFunction, mergeFunction, LinkedHashMap::new),
        ImmutableMap::copyOf);
  }

  static  Collector> toImmutableSortedMap(
      Comparator comparator,
      Function keyFunction,
      Function valueFunction) {
    checkNotNull(comparator);
    checkNotNull(keyFunction);
    checkNotNull(valueFunction);
    /*
     * We will always fail if there are duplicate keys, and the keys are always sorted by
     * the Comparator, so the entries can come in in arbitrary order -- so we report UNORDERED.
     */
    return Collector.of(
        () -> new ImmutableSortedMap.Builder(comparator),
        (builder, input) -> builder.put(keyFunction.apply(input), valueFunction.apply(input)),
        ImmutableSortedMap.Builder::combine,
        ImmutableSortedMap.Builder::build,
        Collector.Characteristics.UNORDERED);
  }

  static  Collector> toImmutableSortedMap(
      Comparator comparator,
      Function keyFunction,
      Function valueFunction,
      BinaryOperator mergeFunction) {
    checkNotNull(comparator);
    checkNotNull(keyFunction);
    checkNotNull(valueFunction);
    checkNotNull(mergeFunction);
    return Collectors.collectingAndThen(
        Collectors.toMap(
            keyFunction, valueFunction, mergeFunction, () -> new TreeMap(comparator)),
        ImmutableSortedMap::copyOfSorted);
  }

  static  Collector> toImmutableBiMap(
      Function keyFunction,
      Function valueFunction) {
    checkNotNull(keyFunction);
    checkNotNull(valueFunction);
    return Collector.of(
        ImmutableBiMap.Builder::new,
        (builder, input) -> builder.put(keyFunction.apply(input), valueFunction.apply(input)),
        ImmutableBiMap.Builder::combine,
        ImmutableBiMap.Builder::build,
        new Collector.Characteristics[0]);
  }

  static , V> Collector> toImmutableEnumMap(
      Function keyFunction,
      Function valueFunction) {
    checkNotNull(keyFunction);
    checkNotNull(valueFunction);
    return Collector.of(
        () ->
            new EnumMapAccumulator(
                (v1, v2) -> {
                  throw new IllegalArgumentException("Multiple values for key: " + v1 + ", " + v2);
                }),
        (accum, t) -> {
          K key = checkNotNull(keyFunction.apply(t), "Null key for input %s", t);
          V newValue = checkNotNull(valueFunction.apply(t), "Null value for input %s", t);
          accum.put(key, newValue);
        },
        EnumMapAccumulator::combine,
        EnumMapAccumulator::toImmutableMap,
        Collector.Characteristics.UNORDERED);
  }

  static , V> Collector> toImmutableEnumMap(
      Function keyFunction,
      Function valueFunction,
      BinaryOperator mergeFunction) {
    checkNotNull(keyFunction);
    checkNotNull(valueFunction);
    checkNotNull(mergeFunction);
    // not UNORDERED because we don't know if mergeFunction is commutative
    return Collector.of(
        () -> new EnumMapAccumulator(mergeFunction),
        (accum, t) -> {
          K key = checkNotNull(keyFunction.apply(t), "Null key for input %s", t);
          V newValue = checkNotNull(valueFunction.apply(t), "Null value for input %s", t);
          accum.put(key, newValue);
        },
        EnumMapAccumulator::combine,
        EnumMapAccumulator::toImmutableMap);
  }

  private static class EnumMapAccumulator, V> {
    private final BinaryOperator mergeFunction;
    private EnumMap map = null;

    EnumMapAccumulator(BinaryOperator mergeFunction) {
      this.mergeFunction = mergeFunction;
    }

    void put(K key, V value) {
      if (map == null) {
        map = new EnumMap<>(key.getDeclaringClass());
      }
      map.merge(key, value, mergeFunction);
    }

    EnumMapAccumulator combine(EnumMapAccumulator other) {
      if (this.map == null) {
        return other;
      } else if (other.map == null) {
        return this;
      } else {
        other.map.forEach(this::put);
        return this;
      }
    }

    ImmutableMap toImmutableMap() {
      return (map == null) ? ImmutableMap.of() : ImmutableEnumMap.asImmutable(map);
    }
  }

  @GwtIncompatible
  static , V>
      Collector> toImmutableRangeMap(
          Function> keyFunction,
          Function valueFunction) {
    checkNotNull(keyFunction);
    checkNotNull(valueFunction);
    return Collector.of(
        ImmutableRangeMap::builder,
        (builder, input) -> builder.put(keyFunction.apply(input), valueFunction.apply(input)),
        ImmutableRangeMap.Builder::combine,
        ImmutableRangeMap.Builder::build);
  }

  // Multimaps

  static  Collector> toImmutableListMultimap(
      Function keyFunction,
      Function valueFunction) {
    checkNotNull(keyFunction, "keyFunction");
    checkNotNull(valueFunction, "valueFunction");
    return Collector.of(
        ImmutableListMultimap::builder,
        (builder, t) -> builder.put(keyFunction.apply(t), valueFunction.apply(t)),
        ImmutableListMultimap.Builder::combine,
        ImmutableListMultimap.Builder::build);
  }

  static  Collector> flatteningToImmutableListMultimap(
      Function keyFunction,
      Function> valuesFunction) {
    checkNotNull(keyFunction);
    checkNotNull(valuesFunction);
    return Collectors.collectingAndThen(
        flatteningToMultimap(
            input -> checkNotNull(keyFunction.apply(input)),
            input -> valuesFunction.apply(input).peek(Preconditions::checkNotNull),
            MultimapBuilder.linkedHashKeys().arrayListValues()::build),
        ImmutableListMultimap::copyOf);
  }

  static  Collector> toImmutableSetMultimap(
      Function keyFunction,
      Function valueFunction) {
    checkNotNull(keyFunction, "keyFunction");
    checkNotNull(valueFunction, "valueFunction");
    return Collector.of(
        ImmutableSetMultimap::builder,
        (builder, t) -> builder.put(keyFunction.apply(t), valueFunction.apply(t)),
        ImmutableSetMultimap.Builder::combine,
        ImmutableSetMultimap.Builder::build);
  }

  static  Collector> flatteningToImmutableSetMultimap(
      Function keyFunction,
      Function> valuesFunction) {
    checkNotNull(keyFunction);
    checkNotNull(valuesFunction);
    return Collectors.collectingAndThen(
        flatteningToMultimap(
            input -> checkNotNull(keyFunction.apply(input)),
            input -> valuesFunction.apply(input).peek(Preconditions::checkNotNull),
            MultimapBuilder.linkedHashKeys().linkedHashSetValues()::build),
        ImmutableSetMultimap::copyOf);
  }

  static > Collector toMultimap(
      Function keyFunction,
      Function valueFunction,
      Supplier multimapSupplier) {
    checkNotNull(keyFunction);
    checkNotNull(valueFunction);
    checkNotNull(multimapSupplier);
    return Collector.of(
        multimapSupplier,
        (multimap, input) -> multimap.put(keyFunction.apply(input), valueFunction.apply(input)),
        (multimap1, multimap2) -> {
          multimap1.putAll(multimap2);
          return multimap1;
        });
  }

  @Beta
  static > Collector flatteningToMultimap(
      Function keyFunction,
      Function> valueFunction,
      Supplier multimapSupplier) {
    checkNotNull(keyFunction);
    checkNotNull(valueFunction);
    checkNotNull(multimapSupplier);
    return Collector.of(
        multimapSupplier,
        (multimap, input) -> {
          K key = keyFunction.apply(input);
          Collection valuesForKey = multimap.get(key);
          valueFunction.apply(input).forEachOrdered(valuesForKey::add);
        },
        (multimap1, multimap2) -> {
          multimap1.putAll(multimap2);
          return multimap1;
        });
  }
}