classlib: fix bug in Stream.takeWhile

Fix #913

classlib/src/main/java/org/teavm/classlib/java/util/stream/doubleimpl/TTakeWhileDoubleStream.java

@@ -17,30 +17,33 @@ package org.teavm.classlib.java.util.stream.doubleimpl;
 
 import java.util.function.DoublePredicate;
 
-public class TTakeWhileDoubleStream extends TWrappingDoubleStreamImpl {
+public class TTakeWhileDoubleStream extends TSimpleDoubleStreamImpl {
+    private TSimpleDoubleStreamImpl sourceStream;
     private DoublePredicate predicate;
 
     /* set to `true` as soon as we see a value `v` in the source stream for which `predicate.test(v)` is false */
     private boolean isStopped;
 
-    TTakeWhileDoubleStream(TSimpleDoubleStreamImpl innerStream, DoublePredicate predicate) {
+    TTakeWhileDoubleStream(TSimpleDoubleStreamImpl sourceStream, DoublePredicate predicate) {
-        super(innerStream);
+        this.sourceStream = sourceStream;
         this.predicate = predicate;
     }
 
     @Override
-    protected DoublePredicate wrap(DoublePredicate consumer) {
+    public boolean next(DoublePredicate consumer) {
-        return t -> {
         if (isStopped) {
             return false;
         }
-
+        var result = sourceStream.next(e -> {
-            if (predicate.test(t)) {
+            if (!predicate.test(e)) {
-                return consumer.test(t);
-            } else {
                 isStopped = true;
                 return false;
             }
-        };
+            return consumer.test(e);
+        });
+        if (!result) {
+            isStopped = true;
+        }
+        return result;
     }
 }

classlib/src/main/java/org/teavm/classlib/java/util/stream/impl/TTakeWhileStream.java

@@ -17,30 +17,33 @@ package org.teavm.classlib.java.util.stream.impl;
 
 import java.util.function.Predicate;
 
-public class TTakeWhileStream<T> extends TWrappingStreamImpl<T, T> {
+public class TTakeWhileStream<T> extends TSimpleStreamImpl<T> {
+    private TSimpleStreamImpl<T> sourceStream;
     private Predicate<? super T> predicate;
 
     /* set to `true` as soon as we see a value `v` in the source stream for which `predicate.test(v)` is false */
     private boolean isStopped;
 
-    TTakeWhileStream(TSimpleStreamImpl<T> innerStream, Predicate<? super T> predicate) {
+    TTakeWhileStream(TSimpleStreamImpl<T> sourceStream, Predicate<? super T> predicate) {
-        super(innerStream);
+        this.sourceStream = sourceStream;
         this.predicate = predicate;
     }
 
     @Override
-    protected Predicate<T> wrap(Predicate<? super T> consumer) {
+    public boolean next(Predicate<? super T> consumer) {
-        return t -> {
         if (isStopped) {
             return false;
         }
-
+        var result = sourceStream.next(e -> {
-            if (predicate.test(t)) {
+            if (!predicate.test(e)) {
-                return consumer.test(t);
-            } else {
                 isStopped = true;
                 return false;
             }
-        };
+            return consumer.test(e);
+        });
+        if (!result) {
+            isStopped = true;
+        }
+        return result;
     }
 }

classlib/src/main/java/org/teavm/classlib/java/util/stream/intimpl/TTakeWhileIntStream.java

@@ -17,30 +17,33 @@ package org.teavm.classlib.java.util.stream.intimpl;
 
 import java.util.function.IntPredicate;
 
-public class TTakeWhileIntStream extends TWrappingIntStreamImpl {
+public class TTakeWhileIntStream extends TSimpleIntStreamImpl {
+    private TSimpleIntStreamImpl sourceStream;
     private IntPredicate predicate;
 
     /* set to `true` as soon as we see a value `v` in the source stream for which `predicate.test(v)` is false */
     private boolean isStopped;
 
-    TTakeWhileIntStream(TSimpleIntStreamImpl innerStream, IntPredicate predicate) {
+    TTakeWhileIntStream(TSimpleIntStreamImpl sourceStream, IntPredicate predicate) {
-        super(innerStream);
+        this.sourceStream = sourceStream;
         this.predicate = predicate;
     }
 
     @Override
-    protected IntPredicate wrap(IntPredicate consumer) {
+    public boolean next(IntPredicate consumer) {
-        return t -> {
         if (isStopped) {
             return false;
         }
-
+        var result = sourceStream.next(e -> {
-            if (predicate.test(t)) {
+            if (!predicate.test(e)) {
-                return consumer.test(t);
-            } else {
                 isStopped = true;
                 return false;
             }
-        };
+            return consumer.test(e);
+        });
+        if (!result) {
+            isStopped = true;
+        }
+        return result;
     }
 }

classlib/src/main/java/org/teavm/classlib/java/util/stream/longimpl/TTakeWhileLongStream.java

@@ -17,30 +17,33 @@ package org.teavm.classlib.java.util.stream.longimpl;
 
 import java.util.function.LongPredicate;
 
-public class TTakeWhileLongStream extends TWrappingLongStreamImpl {
+public class TTakeWhileLongStream extends TSimpleLongStreamImpl {
+    private TSimpleLongStreamImpl sourceStream;
     private LongPredicate predicate;
 
     /* set to `true` as soon as we see a value `v` in the source stream for which `predicate.test(v)` is false */
     private boolean isStopped;
 
     TTakeWhileLongStream(TSimpleLongStreamImpl innerStream, LongPredicate predicate) {
-        super(innerStream);
+        this.sourceStream = innerStream;
         this.predicate = predicate;
     }
 
     @Override
-    protected LongPredicate wrap(LongPredicate consumer) {
+    public boolean next(LongPredicate consumer) {
-        return t -> {
         if (isStopped) {
             return false;
         }
-
+        var result = sourceStream.next(e -> {
-            if (predicate.test(t)) {
+            if (!predicate.test(e)) {
-                return consumer.test(t);
-            } else {
                 isStopped = true;
                 return false;
             }
-        };
+            return consumer.test(e);
+        });
+        if (!result) {
+            isStopped = true;
+        }
+        return result;
     }
 }

tests/src/test/java/org/teavm/classlib/java/util/stream/IntStreamTest.java

@@ -27,6 +27,7 @@ import static org.teavm.classlib.java.util.stream.Helper.testLongStream;
 import java.util.IntSummaryStatistics;
 import java.util.PrimitiveIterator;
 import java.util.Spliterator;
+import java.util.function.IntSupplier;
 import java.util.stream.IntStream;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -112,6 +113,23 @@ public class IntStreamTest {
     public void takeWhileWorks() {
         assertArrayEquals(new int[] { 1, 2, 3 },
                 IntStream.of(1, 2, 3, 4, 0, 5, 6).takeWhile(n -> n < 4).toArray());
+
+        class ForeverIncreasingSupplier implements IntSupplier {
+            int value = 1;
+
+            @Override
+            public int getAsInt() {
+                return value++;
+            }
+        }
+
+        testIntStream(
+                () -> IntStream.concat(
+                        IntStream.generate(new ForeverIncreasingSupplier()).takeWhile(n -> n < 4),
+                        IntStream.generate(new ForeverIncreasingSupplier()).takeWhile(n -> n < 3)
+                ),
+                1, 2, 3, 1, 2
+        );
     }
 
     @Test

tests/src/test/java/org/teavm/classlib/java/util/stream/StreamTest.java

@@ -33,6 +33,7 @@ import java.util.Iterator;
 import java.util.List;
 import java.util.Spliterator;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import java.util.stream.DoubleStream;
 import java.util.stream.IntStream;
 import java.util.stream.LongStream;
@@ -394,6 +395,23 @@ public class StreamTest {
         var sb = new StringBuilder();
         Stream.of(1, 2, 3, 4, 0, 5, 6).takeWhile(n -> n < 4).forEach(sb::append);
         assertEquals("123", sb.toString());
+
+        class ForeverIncreasingSupplier implements Supplier<Integer> {
+            int value = 1;
+
+            @Override
+            public Integer get() {
+                return value++;
+            }
+        }
+
+        testIntegerStream(
+                () -> Stream.concat(
+                        Stream.generate(new ForeverIncreasingSupplier()).takeWhile(n -> n < 4),
+                        Stream.generate(new ForeverIncreasingSupplier()).takeWhile(n -> n < 3)
+                ),
+                1, 2, 3, 1, 2
+        );
     }
 
     @Test
@@ -488,4 +506,5 @@ public class StreamTest {
         List<String> repetitions = Stream.iterate("", s -> s.length() < 5, s -> s + "a").toList();
         assertEquals(List.of("", "a", "aa", "aaa", "aaaa"), repetitions);
     }
+
 }