All Downloads are FREE. Search and download functionalities are using the official Maven repository.

com.github.phantomthief.util.MoreCollectors Maven / Gradle / Ivy

There is a newer version: 0.1.55
Show newest version
package com.github.phantomthief.util;

import static java.util.stream.Collectors.toList;
import static java.util.stream.Stream.of;

import java.util.Collections;
import java.util.EnumSet;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.function.BiConsumer;
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 com.carrotsearch.hppc.IntArrayList;
import com.carrotsearch.hppc.IntHashSet;
import com.carrotsearch.hppc.IntObjectHashMap;
import com.carrotsearch.hppc.LongArrayList;
import com.carrotsearch.hppc.LongHashSet;
import com.google.common.collect.HashMultiset;
import com.google.common.collect.Multimap;

/**
 * @author w.vela
 */
public final class MoreCollectors {

    public static final Set CH_ID = Collections
            .unmodifiableSet(EnumSet.of(Collector.Characteristics.IDENTITY_FINISH));
    public static final Set CH_NOID = Collections.emptySet();

    private MoreCollectors() {
        throw new UnsupportedOperationException();
    }

    public static Collector toIntList() {
        return new CollectorImpl<>(IntArrayList::new, IntArrayList::add, (left, right) -> {
            left.addAll(right);
            return left;
        }, CH_ID);
    }

    public static Collector toLongList() {
        return new CollectorImpl<>(LongArrayList::new, LongArrayList::add, (left, right) -> {
            left.addAll(right);
            return left;
        }, CH_ID);
    }

    public static Collector toIntSet() {
        return new CollectorImpl<>(IntHashSet::new, IntHashSet::add, (left, right) -> {
            left.addAll(right);
            return left;
        }, CH_ID);
    }

    public static Collector toLongSet() {
        return new CollectorImpl<>(LongHashSet::new, LongHashSet::add, (left, right) -> {
            left.addAll(right);
            return left;
        }, CH_ID);
    }

    public static  Collector, IntObjectHashMap> toIntMap(
            ToIntFunction keyMapper, Function valueMapper) {
        BiConsumer, T> accumulator = (map, element) -> map
                .put(keyMapper.applyAsInt(element), valueMapper.apply(element));
        return new CollectorImpl<>(IntObjectHashMap::new, accumulator, (m1, m2) -> {
            m1.putAll(m2);
            return m1;
        }, CH_ID);
    }

    public static  Collector, ?, Map> toMap() {
        return Collectors.toMap(Entry::getKey, Entry::getValue);
    }

    @Deprecated
    public static  Collector> toMultiset(
            Function elementMapper, ToIntFunction countMapper) {
        BiConsumer, T> accumulator = (set, element) -> set
                .add(elementMapper.apply(element), countMapper.applyAsInt(element));
        BinaryOperator> finisher = (m1, m2) -> {
            m1.addAll(m2);
            return m1;
        };
        return new CollectorImpl<>(HashMultiset::create, accumulator, finisher, CH_ID);
    }

    /**
     * use {@link com.google.common.collect.Multimaps#toMultimap(Function, Function, Supplier)} if using guava 21.0+
     */
    @Deprecated
    public static > Collector toMultimap(
            Function keyMapper,
            Function valueMapper, Supplier supplier) {
        BiConsumer accumulator = (multimap, element) -> multimap.put(keyMapper.apply(element),
                valueMapper.apply(element));
        BinaryOperator finisher = (m1, m2) -> {
            m1.putAll(m2);
            return m1;
        };
        return new CollectorImpl<>(supplier, accumulator, finisher, CH_ID);
    }

    public static > Collector toMap(
            Function keyMapper,
            Function valueMapper, Supplier supplier) {
        return toMap(keyMapper, valueMapper, throwingMerger(), supplier);
    }

    private static > Collector toMap(
            Function keyMapper,
            Function valueMapper, BinaryOperator mergeFunction,
            Supplier mapSupplier) {
        BiConsumer accumulator = (map, element) -> map.merge(keyMapper.apply(element),
                valueMapper.apply(element), mergeFunction);
        return new CollectorImpl<>(mapSupplier, accumulator, mapMerger(mergeFunction), CH_ID);
    }

    public static  Collector>> groupingByAllowNullKey(
            Function classifier) {
        return groupingByAllowNullKey(classifier, toList());
    }

    public static  Collector> groupingByAllowNullKey(
            Function classifier,
            Collector downstream) {
        return groupingByAllowNullKey(classifier, HashMap::new, downstream);
    }

    /**
     * allow null key for mapFactory, nothing else changes
     */
    public static > Collector groupingByAllowNullKey(
            Function classifier,
            Supplier mapFactory,
            Collector downstream) {
        Supplier downstreamSupplier = downstream.supplier();
        BiConsumer downstreamAccumulator = downstream.accumulator();
        BiConsumer, T> accumulator = (m, t) -> {
            K key = classifier.apply(t);
            A container = m.computeIfAbsent(key, k -> downstreamSupplier.get());
            downstreamAccumulator.accept(container, t);
        };
        BinaryOperator> merger = mapMerger(downstream.combiner());
        @SuppressWarnings("unchecked")
        Supplier> mangledFactory = (Supplier>) mapFactory;

        if (downstream.characteristics().contains(Collector.Characteristics.IDENTITY_FINISH)) {
            return new CollectorImpl<>(mangledFactory, accumulator, merger, CH_ID);
        } else {
            @SuppressWarnings("unchecked")
            Function downstreamFinisher = (Function) downstream.finisher();
            Function, M> finisher = intermediate -> {
                intermediate.replaceAll((k, v) -> downstreamFinisher.apply(v));
                @SuppressWarnings("unchecked")
                M castResult = (M) intermediate;
                return castResult;
            };
            return new CollectorImpl<>(mangledFactory, accumulator, merger, finisher, CH_NOID);
        }
    }

    @SuppressWarnings("unchecked")
    private static  Function castingIdentity() {
        return i -> (R) i;
    }

    public static  BinaryOperator throwingMerger() {
        return (u, v) -> {
            throw new IllegalStateException(String.format("Duplicate key %s", u));
        };
    }

    private static > BinaryOperator
            mapMerger(BinaryOperator mergeFunction) {
        return (m1, m2) -> {
            for (Map.Entry e : m2.entrySet()) {
                m1.merge(e.getKey(), e.getValue(), mergeFunction);
            }
            return m1;
        };
    }

    private static  Collector combine(Collector collector,
            Function function) {
        return Collector.of(collector.supplier(), collector.accumulator(), collector.combiner(),
                collector.finisher().andThen(function));
    }

    public static  Collector> concat(Stream other) {
        return combine(toList(), list -> Stream.concat(list.stream(), other));
    }

    public static  Collector> concat(T element) {
        return concat(of(element));
    }

    /**
     * Simple implementation class for {@code Collector}.
     *
     * @param  the type of elements to be collected
     * @param  the type of the result
     */
    public static class CollectorImpl implements Collector {

        private final Supplier supplier;

        private final BiConsumer accumulator;

        private final BinaryOperator combiner;

        private final Function finisher;

        private final Set characteristics;

        public CollectorImpl(Supplier supplier, BiConsumer accumulator,
                BinaryOperator combiner, Function finisher,
                Set characteristics) {
            this.supplier = supplier;
            this.accumulator = accumulator;
            this.combiner = combiner;
            this.finisher = finisher;
            this.characteristics = characteristics;
        }

        public CollectorImpl(Supplier supplier, BiConsumer accumulator,
                BinaryOperator combiner, Set characteristics) {
            this(supplier, accumulator, combiner, castingIdentity(), characteristics);
        }

        @Override
        public BiConsumer accumulator() {
            return accumulator;
        }

        @Override
        public Supplier supplier() {
            return supplier;
        }

        @Override
        public BinaryOperator combiner() {
            return combiner;
        }

        @Override
        public Function finisher() {
            return finisher;
        }

        @Override
        public Set characteristics() {
            return characteristics;
        }
    }
}