Add Iterator

typelevel · Aug 9, 2018 · 2854651 · 2854651
1 parent 0c4cc5d
commit 2854651
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 5 deletions.
diff --git a/bench/src/main/scala/cats/bench/CatenableBench.scala b/bench/src/main/scala/cats/bench/CatenableBench.scala
@@ -4,7 +4,7 @@ import cats.data.Catenable
 import org.openjdk.jmh.annotations.{Benchmark, Scope, State}
 
 @State(Scope.Thread)
-class CatenableBenchmark {
+class CatenableBench {
 
   private val smallCatenable = Catenable(1, 2, 3, 4, 5)
   private val smallVector = Vector(1, 2, 3, 4, 5)

diff --git a/core/src/main/scala/cats/data/Catenable.scala b/core/src/main/scala/cats/data/Catenable.scala
@@ -148,14 +148,13 @@ sealed abstract class Catenable[+A] {
   }
 
   /** Applies the supplied function to each element, left to right. */
-  private final def foreach(f: A => Unit): Unit =
-    foreachUntil { a => f(a); false }
+  private final def foreach(f: A => Unit): Unit = foreachUntil { a => f(a); false }
 
   /** Applies the supplied function to each element, left to right, but stops when true is returned */
   private final def foreachUntil(f: A => Boolean): Unit = {
     var c: Catenable[A] = this
     val rights = new collection.mutable.ArrayBuffer[Catenable[A]]
-    // scalastyle:off null
+    // scalastyle:off null return
     while (c ne null) {
       c match {
         case Empty =>
@@ -175,9 +174,12 @@ sealed abstract class Catenable[+A] {
         case Append(l, r) => c = l; rights += r
       }
     }
-    // scalastyle:on null
+    // scalastyle:on null return
   }
 
+
+  final def iterator: Iterator[A] = new CatenableIterator[A](this)
+
   /** Returns the number of elements in this structure */
   final def length: Int = {
     var i: Int = 0
@@ -276,6 +278,31 @@ object Catenable extends CatenableInstances {
         }
       case _ => fromSeq(as)
     }
+
+  class CatenableIterator[A](self: Catenable[A]) extends Iterator[A] {
+    var c: Catenable[A] = if (self.isEmpty) null else self
+    val rights = new collection.mutable.ArrayBuffer[Catenable[A]]
+
+    override def hasNext: Boolean = c ne null
+
+    override def next(): A = {
+      @tailrec def go: A = c match {
+        case Empty =>
+          go // This can't happen
+        case Singleton(a) =>
+          c =
+            if (rights.isEmpty) null
+            else rights.reduceLeft((x, y) => Append(y, x))
+          rights.clear()
+          a
+        case Append(l, r) =>
+          c = l
+          rights += r
+          go
+      }
+      go
+    }
+  }
 }
 
 private[data] sealed abstract class CatenableInstances {

diff --git a/tests/src/test/scala/cats/tests/CatenableSuite.scala b/tests/src/test/scala/cats/tests/CatenableSuite.scala
@@ -56,4 +56,10 @@ class CatenableSuite extends CatsSuite {
       Catenable.fromSeq(ci.toVector) === ci
     }
   }
+
+  test("fromSeq . toList . iterator is id") {
+    forAll { (ci: Catenable[Int]) =>
+      Catenable.fromSeq(ci.iterator.toList) === ci
+    }
+  }
 }