Add com.yahoo.stream.CustomCollectors.

1 files changed, 85 insertions, 0 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/stream/CustomCollectors.java b/vespajlib/src/main/java/com/yahoo/stream/CustomCollectors.java
new file mode 100644
index 00000000000..8198c519dfd
--- /dev/null
+++ b/vespajlib/src/main/java/com/yahoo/stream/CustomCollectors.java
@@ -0,0 +1,85 @@
+// Copyright 2017 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.stream;
+
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.function.BinaryOperator;
+import java.util.function.Function;
+import java.util.function.Supplier;
+import java.util.stream.Collector;
+import java.util.stream.Collectors;
+
+/**
+ * The purpose of this class is to fill gaps in the Java {@link Collectors} api
+ * by offering convenient ways to retrieve implementations of {@link Collector}.
+ *
+ * <p>For example, to get a collector that accumulates elements into an insertion
+ * order map:
+ * <pre>{@code
+ *
+ *     Map<String, Person> idToPerson =
+ *         persons.stream().collect(toLinkedMap(Person::id, Functions.identity());
+ * }</pre>
+ *
+ * @author gjoranv
+ */
+public class CustomCollectors {
+
+    private CustomCollectors() { }
+
+    /**
+     * Returns a {@code Collector} that accumulates elements into a {@code Map}
+     * with predictable iteration order. This a convenience that can be used
+     * instead of calling {@link java.util.stream.Collectors#toMap(Function, Function, BinaryOperator, Supplier)}.
+     * with a merger that throws upon duplicate keys.
+     *
+     * @param keyMapper Mapping function to produce keys.
+     * @param valueMapper Mapping function to produce values.
+     * @param <T> Type of the input elements.
+     * @param <K> Output type of the key mapping function.
+     * @param <U> Output type of the value mapping function.
+     * @return A collector which collects elements into a map with predictable iteration order.
+     * @throws DuplicateKeyException If two elements map to the same key.
+     */
+    public static <T, K, U>
+    Collector<T, ?, Map<K,U>> toLinkedMap(Function<? super T, ? extends K> keyMapper,
+                                          Function<? super T, ? extends U> valueMapper) {
+        return Collectors.toMap(keyMapper, valueMapper, throwingMerger(), LinkedHashMap::new);
+    }
+
+    /**
+     * Returns a {@code Collector} that accumulates elements into a {@code Map}
+     * created by the given supplier. This a convenience that can be used
+     * instead of calling {@link java.util.stream.Collectors#toMap(Function, Function, BinaryOperator, Supplier)}.
+     * with a merger that throws upon duplicate keys.
+     *
+     * @param keyMapper Mapping function to produce keys.
+     * @param valueMapper Mapping function to produce values.
+     * @param mapSupplier Supplier of a new map.
+     * @param <T> Type of the input elements.
+     * @param <K> Output type of the key mapping function.
+     * @param <U> Output type of the value mapping function.
+     * @param <M> Type of the resulting map.
+     * @return A collector which collects elements into a map created by the given supplier.
+     * @throws DuplicateKeyException If two elements map to the same key.
+     */
+    public static <T, K, U, M extends Map<K,U>>
+    Collector<T, ?, M> toCustomMap(Function<? super T, ? extends K> keyMapper,
+                                   Function<? super T, ? extends U> valueMapper,
+                                   Supplier<M> mapSupplier) {
+        return Collectors.toMap(keyMapper, valueMapper, throwingMerger(), mapSupplier);
+    }
+
+
+    private static <T> BinaryOperator<T> throwingMerger() {
+        return (u,v) -> { throw new DuplicateKeyException(u); };
+    }
+
+    public static class DuplicateKeyException extends RuntimeException {
+        private static final long serialVersionUID = 1L;
+
+        DuplicateKeyException(Object key) {
+            super(String.format("Duplicate keys: %s", key));
+        }
+    }
+}