Take deforests

jnape · jnape · commit 9d243788025d · 2018-02-24T18:21:12.000-06:00
diff --git a/src/main/java/com/jnape/palatable/lambda/functions/builtin/fn2/Take.java b/src/main/java/com/jnape/palatable/lambda/functions/builtin/fn2/Take.java
@@ -2,7 +2,7 @@
 
 import com.jnape.palatable.lambda.functions.Fn1;
 import com.jnape.palatable.lambda.functions.Fn2;
-import com.jnape.palatable.lambda.iteration.TakingIterator;
+import com.jnape.palatable.lambda.iteration.TakingIterable;
 
 /**
  * Lazily limit the <code>Iterable</code> to <code>n</code> elements by returning an <code>Iterable</code> that stops
@@ -22,7 +22,7 @@ private Take() {
 
     @Override
     public Iterable<A> apply(Integer n, Iterable<A> as) {
-        return () -> new TakingIterator<>(n, as.iterator());
+        return new TakingIterable<>(n, as);
     }
 
     @SuppressWarnings("unchecked")
diff --git a/src/main/java/com/jnape/palatable/lambda/iteration/TakingIterable.java b/src/main/java/com/jnape/palatable/lambda/iteration/TakingIterable.java
@@ -0,0 +1,25 @@
+package com.jnape.palatable.lambda.iteration;
+
+import java.util.Iterator;
+
+import static java.lang.Math.min;
+
+public final class TakingIterable<A> implements Iterable<A> {
+    private final int         n;
+    private final Iterable<A> as;
+
+    public TakingIterable(int n, Iterable<A> as) {
+        while (as instanceof TakingIterable) {
+            TakingIterable<A> nested = (TakingIterable<A>) as;
+            n = min(n, nested.n);
+            as = nested.as;
+        }
+        this.n = n;
+        this.as = as;
+    }
+
+    @Override
+    public Iterator<A> iterator() {
+        return new TakingIterator<>(n, as.iterator());
+    }
+}
diff --git a/src/test/java/com/jnape/palatable/lambda/iteration/TakingIterableTest.java b/src/test/java/com/jnape/palatable/lambda/iteration/TakingIterableTest.java
@@ -0,0 +1,18 @@
+package com.jnape.palatable.lambda.iteration;
+
+import com.jnape.palatable.lambda.functions.Fn1;
+import com.jnape.palatable.traitor.annotations.TestTraits;
+import com.jnape.palatable.traitor.runners.Traits;
+import org.junit.runner.RunWith;
+import testsupport.traits.Deforesting;
+
+import static com.jnape.palatable.lambda.functions.builtin.fn2.Take.take;
+
+@RunWith(Traits.class)
+public class TakingIterableTest {
+
+    @TestTraits({Deforesting.class})
+    public Fn1<Iterable<Integer>, Iterable<Integer>> testSubject() {
+        return take(1);
+    }
+}
