[#61] limitSorted implementation via headTail()

amaembo · Jan 28, 2016 · 0fe7c2b · 0fe7c2b
1 parent d09eec9
commit 0fe7c2b
Showing 1 changed file with 24 additions and 0 deletions.
diff --git a/src/test/java/one/util/streamex/StreamExHeadTailTest.java b/src/test/java/one/util/streamex/StreamExHeadTailTest.java
@@ -23,6 +23,7 @@
 import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.Deque;
 import java.util.HashSet;
 import java.util.List;
@@ -34,15 +35,21 @@
 import java.util.Spliterator;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
 import java.util.function.BiPredicate;
 import java.util.function.BinaryOperator;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.function.UnaryOperator;
+import java.util.stream.Collector;
 import java.util.stream.Stream;
 
+
+
+
+
 import org.junit.Test;
 
 /**
@@ -309,6 +316,20 @@ private static <T> StreamEx<T> skipLast(StreamEx<T> input, int n, Deque<T> buf)
         });
     }
 
+    static <T> UnaryOperator<StreamEx<T>> limitSorted(Comparator<T> comparator, int n) {
+        @SuppressWarnings("unchecked")
+        Collector<T, Object, List<T>> collector = (Collector<T, Object, List<T>>) MoreCollectors.least(comparator, n);
+        return stream -> collectAndStream(stream, collector.supplier().get(), collector.accumulator(), collector
+                .finisher().andThen(StreamEx::of));
+    }
+
+    private static <T, A, R> StreamEx<R> collectAndStream(StreamEx<T> input, A buf, BiConsumer<A, T> accumulator, Function<A, StreamEx<R>> finisher) {
+        return input.headTail((head, tail) -> {
+            accumulator.accept(buf, head);
+            return collectAndStream(tail, buf, accumulator, finisher);
+        }, () -> finisher.apply(buf));
+    }
+
     // ///////////////////////
     // Terminal ops
 
@@ -409,6 +430,9 @@ public void testHeadTailRecursive() {
 
         assertEquals(asList(5, 10, 1, 6, 7), skipLast(Stream.of(5, 10, 1, 6, 7, 15, -1, 10), 3).toList());
         assertEquals(asList(0, 3, 6, 9, 12, 15, 18), every3(IntStreamEx.range(20).boxed()).toList());
+
+        assertEquals(asList(0, 1, 2, 3, 3), StreamEx.of(0, 1, 4, 2, 10, 3, 5, 10, 3, 15).chain(
+            limitSorted(Comparator.<Integer> naturalOrder(), 5)).toList());
     }
 
     @Test