diff --git a/core/src/main/scala/ox/channels/SourceOps.scala b/core/src/main/scala/ox/channels/SourceOps.scala index 1fa5b4fd..7d113f64 100644 --- a/core/src/main/scala/ox/channels/SourceOps.scala +++ b/core/src/main/scala/ox/channels/SourceOps.scala @@ -168,6 +168,75 @@ 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. The order of elements in both sources is preserved. + * + * If one of the sources is closed before the other, the behavior depends on the `eagerCancel` flag. When set to `true`, the other source + * is cancelled immediately, otherwise the remaining elements from the other source are sent to the returned channel. + * + * 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. + * @param eagerComplete + * If `true`, the returned channel is completed as soon as either of the sources completes. If 'false`, the remaining elements of the + * non-completed source are sent downstream. + * @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, 5, 6, 7) + * 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, 5, 6, 7) + * }}} + */ + def interleave[U >: T](other: Source[U], segmentSize: Int = 1, eagerComplete: Boolean = false)(using Ox, StageCapacity): Source[U] = + val c = StageCapacity.newChannel[U] + + forkDaemon { + var source: Source[U] = this + var counter = 0 + var neitherCompleted = true + + def switchSource(): Unit = { + if (source == this) source = other else source = this + counter = 0 + } + + repeatWhile { + source.receive() match + case ChannelClosed.Done => + // if one source has completed, either complete the resulting source immediately if eagerComplete is set, or: + // - continue with the other source if it hasn't completed yet, or + // - complete the resulting source if both input sources have completed + if (neitherCompleted && !eagerComplete) { + 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/SourceOpsInterleaveTest.scala b/core/src/test/scala/ox/channels/SourceOpsInterleaveTest.scala new file mode 100644 index 00000000..8120b3c5 --- /dev/null +++ b/core/src/test/scala/ox/channels/SourceOpsInterleaveTest.scala @@ -0,0 +1,82 @@ +package ox.channels + +import org.scalatest.flatspec.AnyFlatSpec +import org.scalatest.matchers.should.Matchers +import ox.* + +class SourceOpsInterleaveTest extends AnyFlatSpec with Matchers { + + behavior of "Source.interleave" + + 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 "interleave two sources with different lengths and complete eagerly" in scoped { + val c1 = Source.fromValues(1, 3, 5) + val c2 = Source.fromValues(2, 4, 6, 8, 10, 12) + + val s = c1.interleave(c2, eagerComplete = true) + + s.toList shouldBe List(1, 2, 3, 4, 5, 6) + } + + it should "when empty, interleave with a non-empty source and complete eagerly" in scoped { + val c1 = Source.fromValues() + val c2 = Source.fromValues(1, 2, 3) + + val s1 = c1.interleave(c2, eagerComplete = true) + + s1.toList shouldBe empty + } + + it should "interleave with an empty source and complete eagerly" in scoped { + val c1 = Source.fromValues(1, 2, 3) + val c2 = Source.fromValues() + + val s1 = c1.interleave(c2, eagerComplete = true) + + s1.toList shouldBe List(1) + } +}