From eba1d1e574ee534f8de315c095f0dab9e0146c34 Mon Sep 17 00:00:00 2001
From: Ulugbek Abdullaev <ulugbekna@gmail.com>
Date: Fri, 10 Jun 2022 12:41:16 +0200
Subject: [PATCH] stdlib: implement `Stream.takeWhile` (#588)

Co-authored-by: Ulugbek <ulugbek.abdullaev@jetbrains.com>
---
 .../classlib/java/util/stream/TStream.java    |  2 +
 .../util/stream/impl/TSimpleStreamImpl.java   |  5 ++
 .../util/stream/impl/TTakeWhileStream.java    | 46 +++++++++++++++++++
 .../classlib/java/util/stream/StreamTest.java | 14 ++++++
 4 files changed, 67 insertions(+)
 create mode 100644 classlib/src/main/java/org/teavm/classlib/java/util/stream/impl/TTakeWhileStream.java

diff --git a/classlib/src/main/java/org/teavm/classlib/java/util/stream/TStream.java b/classlib/src/main/java/org/teavm/classlib/java/util/stream/TStream.java
index 38ed37c98..5e682638b 100644
--- a/classlib/src/main/java/org/teavm/classlib/java/util/stream/TStream.java
+++ b/classlib/src/main/java/org/teavm/classlib/java/util/stream/TStream.java
@@ -79,6 +79,8 @@ public interface TStream<T> extends TBaseStream<T, TStream<T>> {
 
     TStream<T> limit(long maxSize);
 
+    TStream<T> takeWhile(Predicate<? super T> predicate);
+
     TStream<T> skip(long n);
 
     void forEach(Consumer<? super T> action);
diff --git a/classlib/src/main/java/org/teavm/classlib/java/util/stream/impl/TSimpleStreamImpl.java b/classlib/src/main/java/org/teavm/classlib/java/util/stream/impl/TSimpleStreamImpl.java
index 07ed91145..10f839089 100644
--- a/classlib/src/main/java/org/teavm/classlib/java/util/stream/impl/TSimpleStreamImpl.java
+++ b/classlib/src/main/java/org/teavm/classlib/java/util/stream/impl/TSimpleStreamImpl.java
@@ -111,6 +111,11 @@ public abstract class TSimpleStreamImpl<T> implements TStream<T> {
         return new TLimitingStreamImpl<>(this, (int) maxSize);
     }
 
+    @Override
+    public TStream<T> takeWhile(Predicate<? super T> predicate) {
+        return new TTakeWhileStream<>(this, predicate);
+    }
+
     @Override
     public TStream<T> skip(long n) {
         return new TSkippingStreamImpl<>(this, (int) n);
diff --git a/classlib/src/main/java/org/teavm/classlib/java/util/stream/impl/TTakeWhileStream.java b/classlib/src/main/java/org/teavm/classlib/java/util/stream/impl/TTakeWhileStream.java
new file mode 100644
index 000000000..18fb24164
--- /dev/null
+++ b/classlib/src/main/java/org/teavm/classlib/java/util/stream/impl/TTakeWhileStream.java
@@ -0,0 +1,46 @@
+/*
+ *  Copyright 2022 ulugbek.abdullaev.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+package org.teavm.classlib.java.util.stream.impl;
+
+import java.util.function.Predicate;
+
+public class TTakeWhileStream<T> extends TWrappingStreamImpl<T, T> {
+    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) {
+        super(innerStream);
+        this.predicate = predicate;
+    }
+
+    @Override
+    protected Predicate<T> wrap(Predicate<? super T> consumer) {
+        return t -> {
+            if (isStopped) {
+                return false;
+            }
+
+            if (predicate.test(t)) {
+                return consumer.test(t);
+            } else {
+                isStopped = true;
+                return false;
+            }
+        };
+    }
+}
diff --git a/tests/src/test/java/org/teavm/classlib/java/util/stream/StreamTest.java b/tests/src/test/java/org/teavm/classlib/java/util/stream/StreamTest.java
index 16d5dd24e..5f88205db 100644
--- a/tests/src/test/java/org/teavm/classlib/java/util/stream/StreamTest.java
+++ b/tests/src/test/java/org/teavm/classlib/java/util/stream/StreamTest.java
@@ -386,4 +386,18 @@ public class StreamTest {
         spliterator.forEachRemaining(appendNumbersTo(sb));
         assertEquals("1;2;3;", sb.toString());
     }
+
+    @Test
+    public void takeWhileWorks() {
+        var sb = new StringBuilder();
+        Stream.of(1, 2, 3, 4, 0, 5, 6).takeWhile(n -> n < 4).forEach(sb::append);
+        assertEquals("123", sb.toString());
+    }
+
+    @Test
+    public void takeWhileWithOtherStreamOps() {
+        var sb = new StringBuilder();
+        Stream.of(1, 2, 3, 4, 0, 5, 6).takeWhile(i -> i < 4).filter(i -> i % 2 != 0).forEach(sb::append);
+        assertEquals("13", sb.toString());
+    }
 }
\ No newline at end of file