Stdlib: Add OneOrBoth and zipLongest on lists and streams (#1003)

Co-authored-by: Jiří Beneš <[email protected]>

Author: marzipankaiser

examples/stdlib/list/zip.check

@@ -1,4 +1,8 @@
 Cons((1, 100), Cons((2, 200), Cons((3, 300), Nil())))
 Nil()
 Nil()
-Cons((1, 42), Nil())
+Cons((1, 42), Nil())
+Cons(Both(1, 100), Cons(Both(2, 200), Cons(Both(3, 300), Nil())))
+Cons(This(1), Cons(This(2), Cons(This(3), Nil())))
+Cons(That(1), Cons(That(2), Cons(That(3), Nil())))
+Cons(Both(1, 42), Cons(This(2), Cons(This(3), Nil())))

examples/stdlib/list/zip.effekt

@@ -10,10 +10,25 @@ def printListOfPairs(list: List[Tuple2[Int, Int]]): Unit = {
   println(strings)
 }
 
+// used to circumvent missing universal `show` & `println` on some backends
+def printListOfOneOrBoth(list: List[OneOrBoth[Int, Int]]): Unit = {
+  val strings: List[String] = list.map {
+    case Both(fst, snd) => "Both(" ++ show(fst) ++ ", " ++ show(snd) ++ ")"
+    case This(fst) => "This(" ++ show(fst) ++ ")"
+    case That(snd) => "That(" ++ show(snd) ++ ")"
+  }
+  println(strings)
+}
+
 def main() = {
   // synchronized with doctest in `zip`
   printListOfPairs(zip([1, 2, 3], [100, 200, 300]))
   printListOfPairs(zip([1, 2, 3], Nil[Int]()))
   printListOfPairs(zip(Nil[Int](), [1, 2, 3]))
   printListOfPairs(zip([1, 2, 3], [42]))
+  // synchronized with doctest in `zipLongest`
+  printListOfOneOrBoth(zipLongest([1, 2, 3], [100, 200, 300]))
+  printListOfOneOrBoth(zipLongest([1, 2, 3], Nil[Int]()))
+  printListOfOneOrBoth(zipLongest(Nil[Int](), [1, 2, 3]))
+  printListOfOneOrBoth(zipLongest([1, 2, 3], [42]))
 }

examples/stdlib/stream/zip_longest.check

@@ -0,0 +1,10 @@
+1, H
+2, e
+3, l
+4, l
+5, o
+6
+7
+8
+9
+10

examples/stdlib/stream/zip_longest.effekt

@@ -0,0 +1,13 @@
+import stream
+
+def main() = {
+  def stream1() = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10].each
+  // Ideally this would be an array to demonstrate the capabilities. 
+  def stream2() = ['H', 'e', 'l', 'l', 'o'].each
+
+  zipLongest[Int, Char]{stream1}{stream2} {
+    case Both(a, b) => println(show(a) ++ ", " ++ show(b))
+    case This(a) => println(show(a))
+    case That(b) => println("?, " ++ show(b))
+  }
+}

libraries/common/effekt.effekt

@@ -715,6 +715,13 @@ record Tuple4[A, B, C, D](first: A, second: B, third: C, fourth: D)
 record Tuple5[A, B, C, D, E](first: A, second: B, third: C, fourth: D, fifth: E)
 record Tuple6[A, B, C, D, E, F](first: A, second: B, third: C, fourth: D, fifth: E, sixth: F)
 
+// OneOrBoth
+// =========
+type OneOrBoth[A, B] {
+  This(a: A)
+  That(b: B)
+  Both(a: A, b: B)
+}
 
 // Control Flow
 // ============

libraries/common/list.effekt

@@ -601,6 +601,40 @@ def zip[A, B](left: List[A], right: List[B]): List[(A, B)] = {
   go(Nil(), left, right)
 }
 
+/// Produce a list of OneOrBoth from two lists.
+/// The length is the maximum of lenghts of the two lists.
+///
+/// Examples:
+/// ```
+/// > zipLongest([1, 2, 3], [100, 200, 300])
+/// [Both(1, 100), Both(2, 200), Both(3, 300)]
+///
+/// > zipLongest([1, 2, 3], Nil[Int]())
+/// [This(1), This(2), This(3)]
+///
+/// > zipLongest(Nil[Int](), [1, 2, 3])
+/// [That(1), That(2), That(3)]
+///
+/// > zipLongest([1, 2, 3], [42])
+/// [These(1, 42), This(2), This(3)]
+/// ```
+///
+/// O(N)
+def zipLongest[A, B](left: List[A], right: List[B]): List[OneOrBoth[A, B]] = {
+  def go(acc: List[OneOrBoth[A, B]], left: List[A], right: List[B]): List[OneOrBoth[A,B]] = {
+    (left, right) match {
+      case (Cons(a, as), Cons(b, bs)) =>
+        go(Cons(Both(a, b), acc), as, bs)
+      case (Cons(a, as), Nil()) =>
+        go(Cons(This(a), acc), as, Nil())
+      case (Nil(), Cons(b, bs)) =>
+        go(Cons(That(b), acc), Nil(), bs)
+      case _ => acc.reverse
+    }
+  }
+  go(Nil(), left, right)
+}
+
 /// Combine two lists with the given function.
 /// The length of the result is the minimum of lengths of the two lists.
 ///

libraries/common/stream.effekt

@@ -286,7 +286,7 @@ def feed[R](bytes: ByteArray) { reader: () => R / read[Byte] }: R = {
 def source[A] { stream: () => Unit / emit[A] } { reader: () => Unit / read[A] }: Unit =
   returning::source[A, Unit]{stream}{reader}
 
-/// Combines two streams together producing a stream of pairs in lockstep.
+/// Combines two streams together calling action on pairs in lockstep.
 /// Terminates when either of them terminates.
 def zip[A, B] { stream1: () => Unit / emit[A] } { stream2: () => Unit / emit[B] } { action: (A, B) => Unit }: Unit = {
   with source[A] { stream1 }
@@ -297,6 +297,28 @@ def zip[A, B] { stream1: () => Unit / emit[A] } { stream2: () => Unit / emit[B]
   }
 }
 
+/// Combines two streams together calling action on OneOrBoth in lockstep.
+/// Terminates when both of them terminated.
+def zipLongest[A, B] { stream1: () => Unit / emit[A] } { stream2: () => Unit / emit[B] } { action: OneOrBoth[A, B] => Unit }: Unit = {
+  with source[A] { stream1 }
+  with source[B] { stream2 }
+
+  exhaustively {
+     try {
+      val a = do read[A]()
+      try {
+        val b = do read[B]()
+        action(Both(a, b))
+      } with stop { () =>
+        action(This(a))
+      }
+    } with stop{ () =>
+      val b = do read[B]()
+      action(That(b))
+    }
+  }
+}
+
 def writeFile[R](path: String) { stream: () => R / emit[Byte] }: R / Exception[IOError] = {
 
   val file = openForWriting(path);