changeset 7279:57f32115652d

Update partition collector to return Map<Boolean,T>
author briangoetz
date Tue, 05 Feb 2013 11:52:03 -0500
parents 16030712f9eb
children a4f32e60635c c20d8fb25a17
files src/share/classes/java/util/stream/Collectors.java
diffstat 1 files changed, 110 insertions(+), 48 deletions(-) [+]
line wrap: on
line diff
--- a/src/share/classes/java/util/stream/Collectors.java	Tue Feb 05 11:02:28 2013 -0500
+++ b/src/share/classes/java/util/stream/Collectors.java	Tue Feb 05 11:52:03 2013 -0500
@@ -24,12 +24,16 @@
  */
 package java.util.stream;
 
+import java.util.AbstractMap;
+import java.util.AbstractSet;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.NoSuchElementException;
 import java.util.Objects;
 import java.util.OptionalDouble;
 import java.util.Set;
@@ -38,7 +42,6 @@
 import java.util.function.BinaryOperator;
 import java.util.function.Function;
 import java.util.function.Functions;
-import java.util.function.IntFunction;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
 
@@ -212,75 +215,134 @@
         };
     }
 
-    public static<T> Collector<T, Collection<T>[]> partitioningBy(Predicate<T> predicate,
-                                                                  IntFunction<Collection<T>[]> arraySupplier) {
-        return partitioningBy(predicate, arraySupplier, ArrayList::new);
-    }
+    static class Partition<T> extends AbstractMap<Boolean, T> implements Map<Boolean, T> {
+        private T forTrue;
+        private T forFalse;
 
-    public static<T, C extends Collection<T>> Collector<T, C[]> partitioningBy(Predicate<T> predicate,
-                                                                               IntFunction<C[]> arraySupplier,
-                                                                               Supplier<C> rowFactory) {
-        return partitioningBy(predicate, arraySupplier, toCollection(rowFactory));
-    }
+        Partition(T forTrue, T forFalse) {
+            this.forTrue = forTrue;
+            this.forFalse = forFalse;
+        }
 
-    public static<T, D> Collector<T, D[]> partitioningBy(Predicate<T> predicate,
-                                                         IntFunction<D[]> arraySupplier,
-                                                         Collector<T, D> downstream) {
-        return new Collector<T, D[]>() {
-            @Override
-            public D[] makeResult() {
-                D[] array = arraySupplier.apply(2);
-                array[0] = downstream.makeResult();
-                array[1] = downstream.makeResult();
-                return array;
+        class Entry implements Map.Entry<Boolean, T> {
+            private final int state;
+
+            Entry(int state) {
+                this.state = state;
             }
 
             @Override
-            public void accumulate(D[] result, T value) {
-                int index = predicate.test(value) ? 1 : 0;
-                downstream.accumulate(result[index], value);
+            public Boolean getKey() {
+                return state == 1;
             }
 
             @Override
-            public D[] combine(D[] result, D[] other) {
-                result[0] = downstream.combine(result[0], other[0]);
-                result[1] = downstream.combine(result[1], other[1]);
+            public T getValue() {
+                return state == 0 ? forFalse : forTrue;
+            }
+
+            @Override
+            public T setValue(T value) {
+                throw new UnsupportedOperationException();
+            }
+        }
+
+        @Override
+        public Set<Map.Entry<Boolean, T>> entrySet() {
+            return new AbstractSet<Map.Entry<Boolean, T>>() {
+                @Override
+                public Iterator<Map.Entry<Boolean, T>> iterator() {
+                    return new Iterator<Map.Entry<Boolean, T>>() {
+                        int state = 0;
+
+                        @Override
+                        public boolean hasNext() {
+                            return state < 2;
+                        }
+
+                        @Override
+                        public Map.Entry<Boolean, T> next() {
+                            if (state < 2)
+                                return new Entry(++state);
+                            else
+                                throw new NoSuchElementException();
+                        }
+                    };
+                }
+
+                @Override
+                public int size() {
+                    return 2;
+                }
+            };
+        }
+    }
+
+    public static<T> Collector<T, Map<Boolean, Collection<T>>> partitioningBy(Predicate<T> predicate) {
+        return partitioningBy(predicate, ArrayList::new);
+    }
+
+    public static<T, C extends Collection<T>> Collector<T, Map<Boolean, C>> partitioningBy(Predicate<T> predicate,
+                                                                                           Supplier<C> rowFactory) {
+        return partitioningBy(predicate, toCollection(rowFactory));
+    }
+
+    public static<T, D> Collector<T, Map<Boolean, D>> partitioningBy(Predicate<T> predicate,
+                                                                     Collector<T, D> downstream) {
+        return new Collector<T, Map<Boolean, D>>() {
+            @Override
+            public Map<Boolean, D> makeResult() {
+                return new Partition<>(downstream.makeResult(), downstream.makeResult());
+            }
+
+            @Override
+            public void accumulate(Map<Boolean, D> result, T value) {
+                Partition<D> asPartition = ((Partition<D>) result);
+                downstream.accumulate(predicate.test(value) ? asPartition.forTrue : asPartition.forFalse, value);
+            }
+
+            @Override
+            public Map<Boolean, D> combine(Map<Boolean, D> result, Map<Boolean, D> other) {
+                Partition<D> left = (Partition<D>) result;
+                Partition<D> right = (Partition<D>) other;
+                left.forFalse = downstream.combine(left.forFalse, right.forFalse);
+                left.forTrue = downstream.combine(left.forTrue, right.forTrue);
                 return result;
             }
         };
     }
 
-    public static<T> Collector<T, T[]> partitioningReduce(Predicate<T> predicate,
-                                                          IntFunction<T[]> arraySupplier,
-                                                          T identity,
-                                                          BinaryOperator<T> reducer) {
-        return partitioningReduce(predicate, arraySupplier, identity, Functions.identity(), reducer);
+    public static<T> Collector<T, Map<Boolean, T>> partitioningReduce(Predicate<T> predicate,
+                                                                      T identity,
+                                                                      BinaryOperator<T> reducer) {
+        return partitioningReduce(predicate, identity, Functions.identity(), reducer);
     }
 
-    public static<T, U> Collector<T, U[]> partitioningReduce(Predicate<T> predicate,
-                                                             IntFunction<U[]> arraySupplier,
-                                                             U identity,
-                                                             Function<T, U> mapper,
-                                                             BinaryOperator<U> reducer) {
-        return new Collector<T, U[]>() {
+    public static<T, U> Collector<T, Map<Boolean, U>> partitioningReduce(Predicate<T> predicate,
+                                                                         U identity,
+                                                                         Function<T, U> mapper,
+                                                                         BinaryOperator<U> reducer) {
+        return new Collector<T, Map<Boolean, U>>() {
             @Override
-            public U[] makeResult() {
-                U[] array = arraySupplier.apply(2);
-                array[0] = identity;
-                array[1] = identity;
-                return array;
+            public Map<Boolean, U> makeResult() {
+                return new Partition<>(identity, identity);
             }
 
             @Override
-            public void accumulate(U[] result, T value) {
-                int index = predicate.test(value) ? 1 : 0;
-                result[index] = reducer.apply(result[index], mapper.apply(value));
+            public void accumulate(Map<Boolean, U> result, T value) {
+                Partition<U> asPartition = ((Partition<U>) result);
+                if (predicate.test(value))
+                    asPartition.forTrue = reducer.apply(asPartition.forTrue, mapper.apply(value));
+                else
+                    asPartition.forFalse = reducer.apply(asPartition.forFalse, mapper.apply(value));
             }
 
             @Override
-            public U[] combine(U[] result, U[] other) {
-                result[0] = reducer.apply(result[0], other[0]);
-                result[1] = reducer.apply(result[1], other[1]);
+            public Map<Boolean, U> combine(Map<Boolean, U> result, Map<Boolean, U> other) {
+                Partition<U> left = (Partition<U>) result;
+                Partition<U> right = (Partition<U>) other;
+                left.forFalse = reducer.apply(left.forFalse, right.forFalse);
+                left.forTrue = reducer.apply(left.forTrue, right.forTrue);
                 return result;
             }
         };