Add Source.mapConcat operator (#99)

Author: kciesielski

core/src/main/scala/ox/channels/SourceOps.scala

@@ -538,6 +538,52 @@ trait SourceOps[+T] { outer: Source[T] =>
     }
     c
 
+  /** Applies the given mapping function `f`, to each element received from this source, transforming it into an Iterable of results, then
+    * sends the results one by one to the returned channel. Can be used to unfold incoming sequences of elements into single elements.
+    *
+    * @param f
+    *   A function that transforms the element from this source into a pair of the next state into an [[scala.collection.IterableOnce]] of
+    *   results which are sent one by one to the returned channel. If the result of `f` is empty, nothing is sent to the returned channel.
+    * @return
+    *   A source to which the results of applying `f` to the elements from this source would be sent.
+    * @example
+    *   {{{
+    *   scala>
+    *   import ox.*
+    *   import ox.channels.Source
+    *
+    *   supervised {
+    *     val s = Source.fromValues(List(1, 2, 3), List(4, 5, 6), List(7, 8, 9))
+    *     s.mapConcat(identity)
+    *   }
+    *
+    *   scala> val res0: List[Int] = List(1, 2, 3, 4, 5, 6, 7, 8, 9)
+    *   }}}
+    */
+  def mapConcat[U](f: T => IterableOnce[U])(using Ox, StageCapacity): Source[U] =
+    val c = StageCapacity.newChannel[U]
+    fork {
+      repeatWhile {
+        receiveSafe() match
+          case ChannelClosed.Done =>
+            c.doneSafe()
+            false
+          case ChannelClosed.Error(r) =>
+            c.errorSafe(r)
+            false
+          case t: T @unchecked =>
+            try
+              val results: IterableOnce[U] = f(t)
+              results.iterator.foreach(c.send)
+              true
+            catch
+              case t: Throwable =>
+                c.errorSafe(t)
+                false
+      }
+    }
+    c
+
   /** Returns the first element from this source wrapped in [[Some]] or [[None]] when this source is empty. Note that `headOption` is not an
     * idempotent operation on source as it receives elements from it.
     *
@@ -565,7 +611,7 @@ trait SourceOps[+T] { outer: Source[T] =>
         case e: ChannelClosed.Error => throw e.toThrowable
         case t: T @unchecked        => Some(t)
     }
-    
+
   /** Sends elements to the returned channel limiting the throughput to specific number of elements (evenly spaced) per time unit. Note that
     * the element's `receive()` time is included in the resulting throughput. For instance having `throttle(1, 1.second)` and `receive()`
     * taking `Xms` means that resulting channel will receive elements every `1s + Xms` time. Throttling is not applied to the empty source.

core/src/test/scala/ox/channels/SourceOpsMapConcatTest.scala

@@ -0,0 +1,51 @@
+package ox.channels
+
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+import ox.*
+
+class SourceOpsMapConcatTest extends AnyFlatSpec with Matchers {
+
+  behavior of "Source.mapConcat"
+
+  it should "unfold iterables" in supervised {
+    val c = Source.fromValues(List("a", "b", "c"), List("d", "e"), List("f"))
+    val s = c.mapConcat(identity)
+    s.toList shouldBe List("a", "b", "c", "d", "e", "f")
+  }
+
+  it should "transform elements" in supervised {
+    val c = Source.fromValues("ab", "cd")
+    val s = c.mapConcat { str => str.toList }
+
+    s.toList shouldBe List('a', 'b', 'c', 'd')
+  }
+
+  it should "handle empty lists" in supervised {
+    val c = Source.fromValues(List.empty, List("a"), List.empty, List("b", "c"))
+    val s = c.mapConcat(identity)
+
+    s.toList shouldBe List("a", "b", "c")
+  }
+
+  it should "propagate errors in the mapping function" in supervised {
+    // given
+    given StageCapacity = StageCapacity(0) // so that the error isn't created too early
+    val c = Source.fromValues(List("a"), List("b", "c"), List("error here"))
+
+    // when
+    val s = c.mapConcat { element =>
+      if (element != List("error here"))
+        element
+      else throw new RuntimeException("boom")
+    }
+
+    // then
+    s.receive() shouldBe "a"
+    s.receive() shouldBe "b"
+    s.receive() shouldBe "c"
+    s.receiveSafe() should matchPattern {
+      case ChannelClosed.Error(reason) if reason.getMessage == "boom" =>
+    }
+  }
+}