Skip to content

Commit

Permalink
Add Source.interleave combinator
Browse files Browse the repository at this point in the history
  • Loading branch information
rucek committed Oct 5, 2023
1 parent 3b6f687 commit 151b52b
Show file tree
Hide file tree
Showing 2 changed files with 106 additions and 0 deletions.
61 changes: 61 additions & 0 deletions core/src/main/scala/ox/channels/SourceOps.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
45 changes: 45 additions & 0 deletions core/src/test/scala/ox/channels/SourceOpsTest.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

0 comments on commit 151b52b

Please sign in to comment.