Java 8 Collector for Google Guava SetMultimap

gistfile1.java

package com.github.gist.example.collectors;

import java.util.Collections;
import java.util.EnumSet;
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.stream.Collector;

import com.google.common.collect.LinkedHashMultimap;
import com.google.common.collect.SetMultimap;

public class SetMultimapCollector<T, K, U> {

	public static <T, K, U, M extends SetMultimap<K, U>> Collector<T, ?, M> toSetMultiMap(
			Function<? super T, ? extends K> keyMapper,
			Function<? super T, ? extends Set<U>> valueMapper) {
		Supplier<M> mapSupplier = new Supplier<M>() {
			@SuppressWarnings("unchecked")
			@Override
			public M get() {
				return (M) LinkedHashMultimap.create();
			}
		};
		BiConsumer<M, T> accumulator = (map, element) -> map.putAll(
				keyMapper.apply(element), valueMapper.apply(element));
		
		BinaryOperator<M> combiner = (m1, m2) -> {
			m1.putAll(m2);
			return m1;
		};
		return new CollectorImpl<>(mapSupplier, accumulator, combiner, CH_ID);
	}

	static final Set<Collector.Characteristics> CH_ID = Collections
			.unmodifiableSet(EnumSet
					.of(Collector.Characteristics.IDENTITY_FINISH));

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

	private static class CollectorImpl<T, A, R> implements Collector<T, A, R> {
		private final Supplier<A> supplier;
		private final BiConsumer<A, T> accumulator;
		private final BinaryOperator<A> combiner;
		private final Function<A, R> finisher;
		private final Set<Characteristics> characteristics;

		private CollectorImpl(Supplier<A> supplier, BiConsumer<A, T> accumulator,
				BinaryOperator<A> combiner, Function<A, R> finisher,
				Set<Characteristics> characteristics) {
			this.supplier = supplier;
			this.accumulator = accumulator;
			this.combiner = combiner;
			this.finisher = finisher;
			this.characteristics = characteristics;
		}

		private CollectorImpl(Supplier<A> supplier, BiConsumer<A, T> accumulator,
				BinaryOperator<A> combiner, Set<Characteristics> characteristics) {
			this(supplier, accumulator, combiner, castingIdentity(),
					characteristics);
		}

		@Override
		public BiConsumer<A, T> accumulator() {
			return accumulator;
		}

		@Override
		public Supplier<A> supplier() {
			return supplier;
		}

		@Override
		public BinaryOperator<A> combiner() {
			return combiner;
		}

		@Override
		public Function<A, R> finisher() {
			return finisher;
		}

		@Override
		public Set<Characteristics> characteristics() {
			return characteristics;
		}
	}

}

I'm sure there's a more elegant way - but this was my first custom collector under Java 8. Makes me like Scala even more :-)

A wee test to show how to use it:

package com.alstom.miso.mss.copy.fileutils;
import java.util.List;
import org.junit.Test;
import com.google.common.collect.Lists;
import com.google.common.collect.SetMultimap;
import com.google.common.collect.Sets;

public class SetMultimapCollectorsTest {
    @Test
    public void test() {
        List<String> list = Lists.newArrayList("one", "two", "three");
        SetMultimap<String, Integer> map = list.stream().collect(
                SetMultimapCollector.toSetMultiMap(ea -> ea,
                        ea -> Sets.newHashSet(ea.length(), 1, 3)));
        System.out.println(map);
    }
}