Stream: tap and tapChunk (#988)

Adds methods for performing side effects on streams without altering them.
getkyo · Jan 9, 2025 · 7a775bf · 7a775bf
1 parent c4dc59c
commit 7a775bf
Show file tree

Hide file tree

Showing 3 changed files with 73 additions and 0 deletions.
diff --git a/kyo-combinators/shared/src/test/scala/kyo/EmitCombinatorTest.scala b/kyo-combinators/shared/src/test/scala/kyo/EmitCombinatorTest.scala
@@ -103,4 +103,9 @@ class EmitCombinatorTest extends Test:
         }
     }
 
+    "stream.tap should not conflict with effect tap combinator" in {
+        val stream = Stream.init(Seq(1, 2, 3)).tap(i => Var.update[Int](_ + i).unit)
+        assert(Var.runTuple(0)(stream.run).eval == (6, Seq(1, 2, 3)))
+    }
+
 end EmitCombinatorTest
diff --git a/kyo-prelude/shared/src/main/scala/kyo/Stream.scala b/kyo-prelude/shared/src/main/scala/kyo/Stream.scala
@@ -117,6 +117,44 @@ sealed abstract class Stream[V, -S]:
                         f(input).map(_.emit).map(ack => ((), cont(ack)))
         ))
 
+    /** Applies a side-effecting function to each element in the stream without altering them.
+      *
+      * @param f
+      *   The function to apply to each value
+      * @return
+      *   A new stream runs f while emitting values
+      */
+    def tap[S1](f: V => Unit < S1)(
+        using
+        tag: Tag[Emit[Chunk[V]]],
+        frame: Frame
+    ): Stream[V, S & S1] =
+        Stream:
+            ArrowEffect.handleState(tag, (), emit: Ack < (Emit[Chunk[V]] & S & S1)):
+                [C] =>
+                    (input, _, cont) =>
+                        Kyo.foreachDiscard(input)(f).andThen:
+                            Emit.andMap(input)(ack => ((), cont(ack)))
+
+    /** Applies a side-effecting function to each chunk in the stream without altering them.
+      *
+      * @param f
+      *   The function to apply to each chunk
+      * @return
+      *   A new stream runs f while emitting chunks
+      */
+    def tapChunk[S1](f: Chunk[V] => Unit < S1)(
+        using
+        tag: Tag[Emit[Chunk[V]]],
+        frame: Frame
+    ): Stream[V, S & S1] =
+        Stream:
+            ArrowEffect.handleState(tag, (), emit: Ack < (Emit[Chunk[V]] & S & S1)):
+                [C] =>
+                    (input, _, cont) =>
+                        f(input).andThen:
+                            Emit.andMap(input)(ack => ((), cont(ack)))
+
     private def discard(using tag: Tag[Emit[Chunk[V]]], frame: Frame): Stream[V, S] =
         Stream(ArrowEffect.handle(tag, emit)(
             [C] => (input, cont) => cont(Stop)

diff --git a/kyo-prelude/shared/src/test/scala/kyo/StreamTest.scala b/kyo-prelude/shared/src/test/scala/kyo/StreamTest.scala
@@ -556,6 +556,36 @@ class StreamTest extends Test:
         }
     }
 
+    "tap" - {
+        "non-empty stream" in {
+            val stream = Stream
+                .init(Seq(1, 2, 3))
+                .tap(i => Var.update[Int](_ + i).unit)
+            assert(Var.runTuple(0)(stream.run).eval == (6, Seq(1, 2, 3)))
+        }
+        "empty stream" in {
+            val stream = Stream
+                .empty[Int]
+                .tap(i => Var.update[Int](_ + i).unit)
+            assert(Var.runTuple(0)(stream.run).eval == (0, Seq()))
+        }
+    }
+
+    "tapChunk" - {
+        "non-empty stream" in {
+            val stream = Stream
+                .apply(Emit.andMap(Chunk(1, 2, 3))(_ => Emit(Chunk(4, 5, 6))))
+                .tapChunk(c => Var.update[Int](_ + c.sum).unit)
+            assert(Var.runTuple(0)(stream.run).eval == (21, Seq(1, 2, 3, 4, 5, 6)))
+        }
+        "empty stream" in {
+            val stream = Stream
+                .empty[Int]
+                .tapChunk(c => Var.update[Int](_ + c.sum).unit)
+            assert(Var.runTuple(0)(stream.run).eval == (0, Seq()))
+        }
+    }
+
     "concat" - {
         "non-empty streams" in {
             assert(