Fix DropWhile fusion bug (palatable#120)

7h3kk1d · web-flow · commit 418cfa108db5 · 2022-03-13T19:09:40.000-05:00
diff --git a/src/main/java/com/jnape/palatable/lambda/internal/iteration/PredicatedDroppingIterable.java b/src/main/java/com/jnape/palatable/lambda/internal/iteration/PredicatedDroppingIterable.java
@@ -1,32 +1,27 @@
 package com.jnape.palatable.lambda.internal.iteration;
 
 import com.jnape.palatable.lambda.functions.Fn1;
+import com.jnape.palatable.lambda.internal.ImmutableQueue;
 
-import java.util.ArrayList;
 import java.util.Iterator;
-import java.util.List;
-
-import static com.jnape.palatable.lambda.functions.builtin.fn2.Any.any;
-import static java.util.Collections.singletonList;
 
 public final class PredicatedDroppingIterable<A> implements Iterable<A> {
-    private final List<Fn1<? super A, ? extends Boolean>> predicates;
-    private final Iterable<A>                             as;
+    private final ImmutableQueue<Fn1<? super A, ? extends Boolean>> predicates;
+    private final Iterable<A>                                       as;
 
     public PredicatedDroppingIterable(Fn1<? super A, ? extends Boolean> predicate, Iterable<A> as) {
-        List<Fn1<? super A, ? extends Boolean>> predicates = new ArrayList<>(singletonList(predicate));
+        ImmutableQueue<Fn1<? super A, ? extends Boolean>> predicates = ImmutableQueue.singleton(predicate);
         while (as instanceof PredicatedDroppingIterable) {
             PredicatedDroppingIterable<A> nested = (PredicatedDroppingIterable<A>) as;
             as = nested.as;
-            predicates.addAll(0, nested.predicates);
+            predicates = nested.predicates.concat(predicates);
         }
         this.predicates = predicates;
         this.as = as;
     }
 
     @Override
     public Iterator<A> iterator() {
-        Fn1<? super A, ? extends Boolean> metaPredicate = a -> any(p -> p.apply(a), predicates);
-        return new PredicatedDroppingIterator<>(metaPredicate, as.iterator());
+        return new PredicatedDroppingIterator<>(predicates, as.iterator());
     }
 }
diff --git a/src/main/java/com/jnape/palatable/lambda/internal/iteration/PredicatedDroppingIterator.java b/src/main/java/com/jnape/palatable/lambda/internal/iteration/PredicatedDroppingIterator.java
@@ -1,19 +1,18 @@
 package com.jnape.palatable.lambda.internal.iteration;
 
 import com.jnape.palatable.lambda.functions.Fn1;
+import com.jnape.palatable.lambda.internal.ImmutableQueue;
 
 import java.util.Iterator;
 import java.util.NoSuchElementException;
 
 public final class PredicatedDroppingIterator<A> extends ImmutableIterator<A> {
-    private final Fn1<? super A, ? extends Boolean> predicate;
-    private final RewindableIterator<A>             rewindableIterator;
-    private       boolean                           finishedDropping;
+    private final Iterator<Fn1<? super A, ? extends Boolean>> predicates;
+    private final RewindableIterator<A>                       rewindableIterator;
 
-    public PredicatedDroppingIterator(Fn1<? super A, ? extends Boolean> predicate, Iterator<A> asIterator) {
-        this.predicate = predicate;
+    public PredicatedDroppingIterator(ImmutableQueue<Fn1<? super A, ? extends Boolean>> predicates, Iterator<A> asIterator) {
+        this.predicates = predicates.iterator();
         rewindableIterator = new RewindableIterator<>(asIterator);
-        finishedDropping = false;
     }
 
     @Override
@@ -31,11 +30,17 @@ public A next() {
     }
 
     private void dropElementsIfNecessary() {
-        while (rewindableIterator.hasNext() && !finishedDropping) {
-            if (!predicate.apply(rewindableIterator.next())) {
-                rewindableIterator.rewind();
-                finishedDropping = true;
+        while (predicates.hasNext() && rewindableIterator.hasNext()) {
+            Fn1<? super A, ? extends Boolean> predicate = predicates.next();
+            boolean predicateDone = false;
+
+            while (rewindableIterator.hasNext() && !predicateDone) {
+                if (!predicate.apply(rewindableIterator.next())) {
+                    rewindableIterator.rewind();
+                    predicateDone = true;
+                }
             }
+
         }
     }
 }
diff --git a/src/main/java/com/jnape/palatable/lambda/internal/iteration/RewindableIterator.java b/src/main/java/com/jnape/palatable/lambda/internal/iteration/RewindableIterator.java
@@ -54,9 +54,7 @@ public A retrieve() {
             if (cache == null)
                 throw new NoSuchElementException("Cache is empty.");
 
-            A cache = this.cache;
-            this.cache = null;
-            return cache;
+            return this.cache;
         }
 
         public boolean isEmpty() {
diff --git a/src/test/java/com/jnape/palatable/lambda/functions/builtin/fn2/DropWhileTest.java b/src/test/java/com/jnape/palatable/lambda/functions/builtin/fn2/DropWhileTest.java
@@ -56,7 +56,13 @@ public void deforestingExecutesPredicatesInOrder() {
             innerInvocations.add(x);
             return x > 2;
         }, asList(1, 2, 3))).forEach(__ -> {});
-        assertThat(innerInvocations, iterates(1, 2, 3));
-        assertThat(outerInvocations, iterates(1, 2));
+        assertThat(innerInvocations, iterates(1));
+        assertThat(outerInvocations, iterates(1, 2, 3));
+    }
+
+    @Test
+    public void eachLayerIsAppliedOnce() {
+        assertThat(dropWhile(i -> i % 2 == 0, dropWhile(i -> i % 2 == 1, asList(1, 2, 3))),
+                   iterates(3));
     }
 }
diff --git a/src/test/java/com/jnape/palatable/lambda/internal/iteration/PredicatedDroppingIteratorTest.java b/src/test/java/com/jnape/palatable/lambda/internal/iteration/PredicatedDroppingIteratorTest.java
@@ -1,15 +1,18 @@
 package com.jnape.palatable.lambda.internal.iteration;
 
 import com.jnape.palatable.lambda.functions.Fn1;
+import com.jnape.palatable.lambda.internal.ImmutableQueue;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnitRunner;
 
+import java.util.Collections;
 import java.util.Iterator;
 import java.util.NoSuchElementException;
 
+import static java.util.Collections.singletonList;
 import static org.hamcrest.core.Is.is;
 import static org.junit.Assert.assertThat;
 import static testsupport.Mocking.mockIteratorToHaveValues;
@@ -25,7 +28,7 @@ public class PredicatedDroppingIteratorTest {
 
     @Before
     public void setUp() {
-        predicatedDroppingIterator = new PredicatedDroppingIterator<>(EVEN, iterator);
+        predicatedDroppingIterator = new PredicatedDroppingIterator<>(ImmutableQueue.singleton(EVEN), iterator);
     }
 
     @Test
diff --git a/src/test/java/com/jnape/palatable/lambda/internal/iteration/RewindableIteratorTest.java b/src/test/java/com/jnape/palatable/lambda/internal/iteration/RewindableIteratorTest.java
@@ -10,6 +10,7 @@
 import java.util.NoSuchElementException;
 
 import static org.hamcrest.core.Is.is;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertThat;
 import static testsupport.Mocking.mockIteratorToHaveValues;
 
@@ -51,13 +52,14 @@ public void cannotRewindIfNoValuesIterated() {
         rewindableIterator.rewind();
     }
 
-    @Test(expected = NoSuchElementException.class)
-    public void cannotRewindTheSameElementTwice() {
+    @Test
+    public void canRewindTheSameElementTwice() {
         mockIteratorToHaveValues(iterator, 1, 2, 3);
         rewindableIterator.next();
         rewindableIterator.rewind();
         rewindableIterator.next();
         rewindableIterator.rewind();
+        assertEquals(1, rewindableIterator.next());
     }
 
     @Test

Original file line number	Diff line number	Diff line change
`@@ -54,9 +54,7 @@ public A retrieve() {`
`54`	`54`	`if (cache == null)`
`55`	`55`	`throw new NoSuchElementException("Cache is empty.");`
`56`	`56`
`57`		`- A cache = this.cache;`
`58`		`- this.cache = null;`
`59`		`- return cache;`
	`57`	`+ return this.cache;`
`60`	`58`	`}`
`61`	`59`
`62`	`60`	`public boolean isEmpty() {`