From 151b52bcd246ecb577792224a3d31ec2af3dd67e Mon Sep 17 00:00:00 2001
From: Jacek Kunicki <jacek.kunicki@softwaremill.com>
Date: Thu, 5 Oct 2023 09:13:50 +0200
Subject: [PATCH] Add Source.interleave combinator

---
 .../main/scala/ox/channels/SourceOps.scala    | 61 +++++++++++++++++++
 .../scala/ox/channels/SourceOpsTest.scala     | 45 ++++++++++++++
 2 files changed, 106 insertions(+)

diff --git a/core/src/main/scala/ox/channels/SourceOps.scala b/core/src/main/scala/ox/channels/SourceOps.scala
index 1fa5b4fd..8277b55b 100644
--- a/core/src/main/scala/ox/channels/SourceOps.scala
+++ b/core/src/main/scala/ox/channels/SourceOps.scala
@@ -168,6 +168,67 @@ trait SourceOps[+T] { this: Source[T] =>
 
   //
 
+  /** Sends a given number of elements (determined byc `segmentSize`) from this source to the returned channel, then sends the same number
+    * of elements from the `other` source and repeats. If one of the sources is closed before the other, the remaining elements from the
+    * open one are sent to the returned channel. The order of elements in both sources is preserved.
+    *
+    * Must be run within a scope, since a child fork is created which receives from both sources and sends to the resulting channel.
+    *
+    * @param other
+    *   The source whose elements will be interleaved with the elements of this source.
+    * @param segmentSize
+    *   The number of elements sent from each source before switching to the other one. Default is 1.
+    * @return
+    *   A source to which the interleaved elements from both sources would be sent.
+    * @example
+    *   {{{
+    *   scala>
+    *   import ox.*
+    *   import ox.channels.Source
+    *
+    *   scoped {
+    *     val s1 = Source.fromValues(1, 2, 3, 4)
+    *     val s2 = Source.fromValues(10, 20, 30, 40)
+    *     s1.interleave(s2, segmentSize = 2).toList
+    *   }
+    *
+    *   scala> val res0: List[Int] = List(1, 2, 10, 20, 3, 4, 30, 40)
+    *   }}}
+    */
+  def interleave[U >: T](other: Source[U], segmentSize: Int = 1)(using Ox, StageCapacity): Source[U] =
+    val c = StageCapacity.newChannel[U]
+    var source: Source[U] = this
+    var counter = 0
+    var neitherCompleted = true
+
+    def switchSource(): Unit = {
+      if (source == this) source = other else source = this
+      counter = 0
+    }
+
+    forkDaemon {
+      repeatWhile {
+        source.receive() match
+          case ChannelClosed.Done =>
+            // if one source has completed, switch to the other one, otherwise (i.e. if both sources have completed) complete the resulting source
+            if (neitherCompleted) {
+              neitherCompleted = false
+              switchSource()
+              true
+            } else {
+              c.done()
+              false
+            }
+          case ChannelClosed.Error(r) => c.error(r); false
+          case value: U @unchecked =>
+            counter += 1
+            // after reaching segmentSize, only switch to the other source if it hasn't completed yet
+            if (counter == segmentSize && neitherCompleted) switchSource()
+            c.send(value).isValue
+      }
+    }
+    c
+
   /** Invokes the given function for each received element. Blocks until the channel is done.
     * @throws ChannelClosedException
     *   when there is an upstream error.
diff --git a/core/src/test/scala/ox/channels/SourceOpsTest.scala b/core/src/test/scala/ox/channels/SourceOpsTest.scala
index e48d1b6f..85ded9a4 100644
--- a/core/src/test/scala/ox/channels/SourceOpsTest.scala
+++ b/core/src/test/scala/ox/channels/SourceOpsTest.scala
@@ -209,6 +209,51 @@ class SourceOpsTest extends AnyFlatSpec with Matchers with Eventually {
     }
   }
 
+  it should "interleave with an empty source" in scoped {
+    val c1 = Source.fromValues(1, 2, 3)
+    val c2 = Source.fromValues()
+
+    val s1 = c1.interleave(c2)
+
+    s1.toList shouldBe List(1, 2, 3)
+  }
+
+  it should "interleave two sources with default segment size" in scoped {
+    val c1 = Source.fromValues(1, 3, 5)
+    val c2 = Source.fromValues(2, 4, 6)
+
+    val s = c1.interleave(c2)
+
+    s.toList shouldBe List(1, 2, 3, 4, 5, 6)
+  }
+
+  it should "interleave two sources with default segment size and different lengths" in scoped {
+    val c1 = Source.fromValues(1, 3, 5)
+    val c2 = Source.fromValues(2, 4, 6, 8, 10, 12)
+
+    val s = c1.interleave(c2)
+
+    s.toList shouldBe List(1, 2, 3, 4, 5, 6, 8, 10, 12)
+  }
+
+  it should "interleave two sources with custom segment size" in scoped {
+    val c1 = Source.fromValues(1, 2, 3, 4)
+    val c2 = Source.fromValues(10, 20, 30, 40)
+
+    val s = c1.interleave(c2, segmentSize = 2)
+
+    s.toList shouldBe List(1, 2, 10, 20, 3, 4, 30, 40)
+  }
+
+  it should "interleave two sources with custom segment size and different lengths" in scoped {
+    val c1 = Source.fromValues(1, 2, 3, 4, 5, 6, 7)
+    val c2 = Source.fromValues(10, 20, 30, 40)
+
+    val s = c1.interleave(c2, segmentSize = 2)
+
+    s.toList shouldBe List(1, 2, 10, 20, 3, 4, 30, 40, 5, 6, 7)
+  }
+
   it should "merge two sources" in {
     scoped {
       val c1 = Source.fromValues(1, 2, 3)