Stack safety and constant factor improvements of Catenable

benchmark/src/main/scala/fs2/benchmark/CatenableBenchmark.scala

@@ -0,0 +1,36 @@
+package fs2
+package benchmark
+
+import org.openjdk.jmh.annotations.{Benchmark, State, Scope}
+
+import fs2.util.Catenable
+
+@State(Scope.Thread)
+class CatenableBenchmark extends BenchmarkUtils {
+
+  val smallCatenable = Catenable(1, 2, 3, 4, 5)
+  val smallVector = Vector(1, 2, 3, 4, 5)
+
+  val largeCatenable = Catenable.fromSeq(0 to 1000000)
+  val largeVector = (0 to 1000000).toVector
+
+  @Benchmark def mapSmallCatenable = smallCatenable.map(_ + 1)
+  @Benchmark def mapSmallVector = smallVector.map(_ + 1)
+  @Benchmark def mapLargeCatenable = largeCatenable.map(_ + 1)
+  @Benchmark def mapLargeVector = largeVector.map(_ + 1)
+
+  @Benchmark def foldLeftSmallCatenable = smallCatenable.foldLeft(0)(_ + _)
+  @Benchmark def foldLeftSmallVector = smallVector.foldLeft(0)(_ + _)
+  @Benchmark def foldLeftLargeCatenable = largeCatenable.foldLeft(0)(_ + _)
+  @Benchmark def foldLeftLargeVector = largeVector.foldLeft(0)(_ + _)
+
+  @Benchmark def consSmallCatenable = 0 +: smallCatenable
+  @Benchmark def consSmallVector = 0 +: smallVector
+  @Benchmark def consLargeCatenable = 0 +: largeCatenable
+  @Benchmark def consLargeVector = 0 +: largeVector
+
+  @Benchmark def createTinyCatenable = Catenable(1)
+  @Benchmark def createTinyVector = Vector(1)
+  @Benchmark def createSmallCatenable = Catenable(1, 2, 3, 4, 5)
+  @Benchmark def createSmallVector = Vector(1, 2, 3, 4, 5)
+}

core/shared/src/main/scala/fs2/StreamCore.scala

@@ -247,7 +247,7 @@ private[fs2] object StreamCore {
     type O0 = O
     def push[G[_],O2](u: NT[F,G], stack: Stack[G,O,O2]) =
       Scope.pure { stack pushSegments (NT.convert(s)(u)) }
-    def render = "Segments(" + s.toStream.toList.mkString(", ") + ")"
+    def render = "Segments(" + s.toList.mkString(", ") + ")"
   }
 
   def scope[F[_],O](s: StreamCore[F,O]): StreamCore[F,O] = StreamCore.evalScope(Scope.snapshot).flatMap { tokens =>
@@ -374,11 +374,11 @@ private[fs2] object StreamCore {
     private final case class Segments[F[_],O](segments: Catenable[Segment[F,O]]) extends Stack[F,O,O] {
 
       def render = {
-        val segmentsList = segments.toStream.toList
+        val segmentsList = segments.toList
         List(s"Segments (${segmentsList.size})\n"+segmentsList.zipWithIndex.map { case (s, idx) => s"    s$idx: $s" }.mkString("\n"))
       }
       def pushNonEmptySegment(s: Segment[F,O]): Stack[F,O,O] =
-        Segments(s :: segments)
+        Segments(s +: segments)
 
       def pushNonEmptySegments(s: Catenable[Segment[F,O]]): Stack[F,O,O] =
         Segments(s ++ segments)
@@ -410,7 +410,7 @@ private[fs2] object StreamCore {
       def render = "Map" :: stack.render
 
       def pushNonEmptySegment(s: Segment[F,O1]): Stack[F,O1,O2] =
-        Map(s :: segments, f, stack)
+        Map(s +: segments, f, stack)
 
       def pushNonEmptySegments(s: Catenable[Segment[F,O1]]): Stack[F,O1,O2] =
         Map(s ++ segments, f, stack)
@@ -452,7 +452,7 @@ private[fs2] object StreamCore {
       def render = "Bind" :: stack.render
 
       def pushNonEmptySegment(s: Segment[F,O1]): Stack[F,O1,O2] =
-        Bind(s :: segments, f, stack)
+        Bind(s +: segments, f, stack)
 
       def pushNonEmptySegments(s: Catenable[Segment[F,O1]]): Stack[F,O1,O2] =
         Bind(s ++ segments, f, stack)
@@ -466,7 +466,7 @@ private[fs2] object StreamCore {
                 case None => stack.pushBind(f).pushSegments(segments).step
                 case Some((hd,tl)) =>
                   val segs2 =
-                    (if (tl.isEmpty) segments else segments.push(Segment.Emit(tl))).map(_.interpretBind(f))
+                    (if (tl.isEmpty) segments else segments.cons(Segment.Emit(tl))).map(_.interpretBind(f))
                   val stack2 = stack.pushSegments(segs2)
                   (try stack2.pushAppend(f(hd)) catch { case NonFatal(t) => stack2.pushFail(t) }).step
               }

core/shared/src/main/scala/fs2/util/Catenable.scala

@@ -5,70 +5,134 @@ import Catenable._
 /**
  * Trivial catenable sequence. Supports O(1) append, and (amortized)
  * O(1) `uncons`, such that walking the sequence via N successive `uncons`
- * steps takes O(N). Like a difference list, conversion to a `Stream[A]`
+ * steps takes O(N). Like a difference list, conversion to a `Seq[A]`
  * takes linear time, regardless of how the sequence is built up.
  */
 sealed abstract class Catenable[+A] {
 
   /** Returns the head and tail of this catenable if non empty, none otherwise. Amortized O(1). */
-  def uncons: Option[(A, Catenable[A])] = {
-    @annotation.tailrec
-    def go(c: Catenable[A], rights: List[Catenable[A]]): Option[(A,Catenable[A])] = c match {
-      case Empty => rights match {
-        case Nil => None
-        case c :: rights => go(c, rights)
+  final def uncons: Option[(A, Catenable[A])] = {
+    var c: Catenable[A] = this
+    var rights: List[Catenable[A]] = Nil
+    var result: Option[(A, Catenable[A])] = null
+    while (result eq null) {
+      c match {
+        case Empty =>
+          rights match {
+            case Nil => result = None
+            case h :: t => c = h; rights = t
+          }
+        case Single(a) =>
+          val next = if (rights.isEmpty) empty else rights.reverse.reduceLeft((x, y) => Append(y,x))
+          result = Some(a -> next)
+        case Append(l, r) => c = l; rights = r :: rights
       }
-      case Single(a) => Some(a -> (if (rights.isEmpty) empty else rights.reduceRight(Append(_,_))))
-      case Append(l,r) => go(l, r :: rights)
     }
-    go(this, List())
+    result
   }
 
-  def isEmpty: Boolean = this match {
-    case Empty => true
-    case _ => false // okay since `append` smart constructor guarantees each branch non-empty
-  }
+  /** Returns true if there are no elements in this collection. */
+  def isEmpty: Boolean
 
   /** Concatenates this with `c` in O(1) runtime. */
-  def ++[A2>:A](c: Catenable[A2])(implicit T: RealSupertype[A,A2]): Catenable[A2] =
+  final def ++[A2>:A](c: Catenable[A2])(implicit T: RealSupertype[A,A2]): Catenable[A2] =
     append(this, c)
 
-  /** Alias for [[push]]. */
-  def ::[A2>:A](a: A2)(implicit T: RealSupertype[A,A2]): Catenable[A2] =
-    this.push(a)
-
   /** Returns a new catenable consisting of `a` followed by this. O(1) runtime. */
-  def push[A2>:A](a: A2)(implicit T: RealSupertype[A,A2]): Catenable[A2] =
+  final def cons[A2>:A](a: A2)(implicit T: RealSupertype[A,A2]): Catenable[A2] =
     append(single(a), this)
 
-  def toStream: Stream[A] = uncons match {
-    case None => Stream.empty
-    case Some((hd, tl)) => hd #:: tl.toStream
+  /** Alias for [[cons]]. */
+  final def +:[A2>:A](a: A2)(implicit T: RealSupertype[A,A2]): Catenable[A2] =
+    cons(a)
+
+  /** Returns a new catenable consisting of this followed by `a`. O(1) runtime. */
+  final def snoc[A2>:A](a: A2)(implicit T: RealSupertype[A,A2]): Catenable[A2] =
+    append(this, single(a))
+
+  /** Alias for [[snoc]]. */
+  final def :+[A2>:A](a: A2)(implicit T: RealSupertype[A,A2]): Catenable[A2] =
+    snoc(a)
+
+  /** Applies the supplied function to each element and returns a new catenable. */
+  final def map[B](f: A => B): Catenable[B] =
+    foldLeft(empty: Catenable[B])((acc, a) => acc :+ f(a))
+
+  /** Folds over the elements from left to right using the supplied initial value and function. */
+  final def foldLeft[B](z: B)(f: (B, A) => B): B = {
+    var result = z
+    foreach(a => result = f(result, a))
+    result
+  }
+
+  /** Applies the supplied function to each element, left to right. */
+  final def foreach(f: A => Unit): Unit = {
+    var c: Catenable[A] = this
+    var rights: List[Catenable[A]] = Nil
+    while (c ne null) {
+      c match {
+        case Empty =>
+          rights match {
+            case Nil => c = null
+            case h :: t => c = h; rights = t
+          }
+        case Single(a) =>
+          f(a)
+          c = if (rights.isEmpty) Empty else rights.reverse.reduceLeft((x, y) => Append(y,x))
+          rights = Nil
+        case Append(l, r) => c = l; rights = r :: rights
+      }
+    }
   }
 
-  def map[B](f: A => B): Catenable[B] = Catenable.fromSeq(toStream.map(f))
+  /** Converts to a list. */
+  final def toList: List[A] = {
+    val builder = List.newBuilder[A]
+    foreach { a => builder += a; () }
+    builder.result
+  }
+
+  override def toString = "Catenable(..)"
 }
 
 object Catenable {
-  private case object Empty extends Catenable[Nothing]
-  private final case class Single[A](a: A) extends Catenable[A]
-  private final case class Append[A](left: Catenable[A], right: Catenable[A]) extends Catenable[A]
+  private final case object Empty extends Catenable[Nothing] {
+    def isEmpty: Boolean = true
+  }
+  private final case class Single[A](a: A) extends Catenable[A] {
+    def isEmpty: Boolean = false
+  }
+  private final case class Append[A](left: Catenable[A], right: Catenable[A]) extends Catenable[A] {
+    def isEmpty: Boolean = false // b/c `append` constructor doesn't allow either branch to be empty
+  }
 
+  /** Empty catenable. */
   val empty: Catenable[Nothing] = Empty
 
+  /** Creates a catenable of 1 element. */
   def single[A](a: A): Catenable[A] = Single(a)
 
-  def append[A](c: Catenable[A], c2: Catenable[A]): Catenable[A] = c match {
-    case Empty => c2
-    case _ => c2 match {
-      case Empty => c
-      case _ => Append(c,c2)
-    }
-  }
+  /** Appends two catenables. */
+  def append[A](c: Catenable[A], c2: Catenable[A]): Catenable[A] =
+    if (c.isEmpty) c2
+    else if (c2.isEmpty) c
+    else Append(c, c2)
 
+  /** Creates a catenable from the specified sequence. */
   def fromSeq[A](s: Seq[A]): Catenable[A] =
     if (s.isEmpty) empty
-    else s.view.map(single).reduceRight(Append(_,_))
+    else s.view.reverse.map(single).reduceLeft((x, y) => Append(y, x))
 
-  def apply[A](as: A*): Catenable[A] = fromSeq(as)
+  /** Creates a catenable from the specified elements. */
+  def apply[A](as: A*): Catenable[A] = {
+    // Assumption: `as` is small enough that calling size doesn't outweigh benefit of calling empty/single
+    as.size match {
+      case 0 => empty
+      case 1 => single(as.head)
+      case n if n <= 1024 =>
+        // nb: avoid cost of reversing input if colection is small
+        as.view.map(single).reduceRight(Append(_, _))
+      case n => as.view.reverse.map(single).reduceLeft((x, y) => Append(y, x))
+    }
+  }
 }