Merge pull request #4 from jeantil/iterator-group_until_changed

adds splitBy extension method on scala.collection.Iterator

Author: julienrf

src/main/scala/scala/collection/decorators/IterableDecorator.scala

@@ -48,4 +48,28 @@ class IterableDecorator[C, I <: IsIterable[C]](coll: C)(implicit val it: I) {
   def lazyFoldRight[B](z: B)(op: it.A => Either[B, B => B]): B =
     it(coll).iterator.lazyFoldRight(z)(op)
 
+
+  /**
+    * Constructs a collection where consecutive elements are accumulated as
+    * long as the output of f for each element doesn't change.
+    * <pre>
+    * Vector(1,2,2,3,3,3,2,2)
+    * .splitBy(identity)
+    * </pre>
+    * produces
+    * <pre>
+    * Vector(Vector(1),
+    * Vector(2,2),
+    * Vector(3,3,3),
+    * Vector(2,2))
+    * </pre>
+    *
+    * @param f the function to compute a key for an element
+    * @tparam K the type of the computed key
+    * @return a collection of collections of the consecutive elements with the
+    *         same key in the original collection
+    */
+  def splitBy[K, CC1, CC2](f: it.A => K)(implicit bf: BuildFrom[C, it.A, CC1], bff: BuildFrom[C, CC1, CC2]): CC2 = {
+    bff.fromSpecific(coll)(it(coll).iterator.splitBy(f).map(bf.fromSpecific(coll)))
+  }
 }

src/main/scala/scala/collection/decorators/IteratorDecorator.scala

@@ -86,4 +86,68 @@ class IteratorDecorator[A](val `this`: Iterator[A]) extends AnyVal {
     loop(immutable.List.empty)
   }
 
+  /**
+    * Constructs an iterator where consecutive elements are accumulated as
+    * long as the output of f for each element doesn't change.
+    * <pre>
+    * Vector(1,2,2,3,3,3,2,2)
+    * .iterator
+    * .splitBy(identity)
+    * .toList
+    * </pre>
+    * produces
+    * <pre>
+    * List(Seq(1),
+    * Seq(2,2),
+    * Seq(3,3,3),
+    * Seq(2,2))
+    * </pre>
+    *
+    * @param f the function to compute a key for an element
+    * @tparam K the type of the computed key
+    * @return an iterator of sequences of the consecutive elements with the
+    *         same key in the original iterator
+    */
+  def splitBy[K](f: A => K): Iterator[immutable.Seq[A]] =
+    new AbstractIterator[immutable.Seq[A]] {
+      private var hd: A = _
+      private var hdKey: K = _
+      private var hdDefined: Boolean = false
+
+      override def hasNext: Boolean = hdDefined || `this`.hasNext
+
+      override def next(): immutable.Seq[A] = {
+        if (hasNext) {
+          val seq = Vector.newBuilder[A]
+          if (hdDefined) {
+            seq += hd
+          } else {
+            val init = `this`.next()
+            hd = init
+            hdKey = f(init)
+            hdDefined = true
+            seq += init
+          }
+          var hadSameKey = true
+          while (`this`.hasNext && hadSameKey) {
+            val el = `this`.next()
+            hdDefined = true
+            val key = f(el)
+            if (key == hdKey) {
+              seq += el
+            } else {
+              hadSameKey = false
+              hdKey = key
+              hd = el
+            }
+          }
+          if (hadSameKey) {
+            hdDefined = false
+          }
+          seq.result()
+        } else {
+          Iterator.empty.next()
+        }
+      }
+    }
 }

src/test/scala/scala/collection/decorators/IterableDecoratorTest.scala

@@ -2,7 +2,8 @@ package scala.collection
 package decorators
 
 import org.junit.{Assert, Test}
-import scala.collection.immutable.{LazyList, List, Range, Map}
+
+import scala.collection.immutable.{LazyList, List, Map, Range}
 
 class IterableDecoratorTest {
 
@@ -54,4 +55,44 @@ class IterableDecoratorTest {
     Assert.assertEquals(2, result2)
   }
 
+  @Test
+  def splitByShouldHonorEmptyIterator(): Unit = {
+    val split = Vector.empty[Int].splitBy(identity)
+    Assert.assertEquals(Vector.empty, split)
+  }
+
+  @Test
+  def splitByShouldReturnSingleSeqWhenSingleElement(): Unit = {
+    val value = Vector("1")
+    val split = value.splitBy(identity)
+    Assert.assertEquals(Vector(value), split)
+  }
+
+  @Test
+  def splitByShouldReturnSingleSeqWhenAllElHaveTheSameKey(): Unit = {
+    val value = Vector("1", "1", "1")
+    val split = value.splitBy(identity)
+    Assert.assertEquals(Vector(value), split)
+  }
+
+  @Test
+  def splitByShouldReturnVectorOfVectorOrConsecutiveElementsWithTheSameKey(): Unit = {
+    val value = Vector("1", "2", "2", "3", "3", "3", "2", "2")
+    val split: Vector[Vector[String]] = value.splitBy(identity)
+    Assert.assertEquals(Vector(Vector("1"), Vector("2", "2"), Vector("3", "3", "3"), Vector("2", "2")), split)
+  }
+
+  @Test
+  def splitByShouldReturnListOfListOfConsecutiveElementsWithTheSameKey(): Unit = {
+    val value = List("1", "2", "2", "3", "3", "3", "2", "2")
+    val split: List[List[String]] = value.splitBy(identity)
+    Assert.assertEquals(List(List("1"), List("2", "2"), List("3", "3", "3"), List("2", "2")), split)
+  }
+
+  @Test
+  def splitByShouldReturnSetOfSetOfConsecutiveElementsWithTheSameKey(): Unit = {
+    val value = Set("1", "2", "2", "3", "3", "3", "2", "2")
+    val split: Set[Set[String]] = value.splitBy(identity)
+    Assert.assertEquals(Set(Set("1"), Set("2"), Set("3")), split)
+  }
 }

src/test/scala/scala/collection/decorators/IteratorDecoratorTest.scala

@@ -0,0 +1,41 @@
+package scala.collection
+package decorators
+
+import org.junit.{Assert, Test}
+
+import scala.util.Try
+
+class IteratorDecoratorTest {
+  @Test
+  def splitByShouldHonorEmptyIterator(): Unit = {
+    val groupedIterator = Iterator.empty.splitBy(identity)
+    Assert.assertFalse(groupedIterator.hasNext)
+    Assert.assertEquals(Try(groupedIterator.next).toString, Try(Iterator.empty.next()).toString)
+  }
+
+  @Test
+  def splitByShouldReturnIteratorOfSingleSeqWhenAllElHaveTheSameKey(): Unit = {
+    val value = Vector("1", "1", "1")
+    val groupedIterator = value.iterator.splitBy(identity)
+    Assert.assertTrue(groupedIterator.hasNext)
+    Assert.assertEquals(groupedIterator.next.toVector, value)
+    Assert.assertFalse(groupedIterator.hasNext)
+    Assert.assertEquals(Try(groupedIterator.next).toString, Try(Iterator.empty.next()).toString)
+  }
+
+  @Test
+  def splitByShouldReturnIteratorOfSeqOfConsecutiveElementsWithTheSameKey(): Unit = {
+    val value = Vector("1", "2", "2", "3", "3", "3", "2", "2")
+    val groupedIterator = value.iterator.splitBy(identity)
+    Assert.assertTrue(groupedIterator.hasNext)
+    Assert.assertEquals(groupedIterator.next.toVector, Vector("1"))
+    Assert.assertTrue(groupedIterator.hasNext)
+    Assert.assertEquals(groupedIterator.next.toVector, Vector("2", "2"))
+    Assert.assertTrue(groupedIterator.hasNext)
+    Assert.assertEquals(groupedIterator.next.toVector, Vector("3", "3", "3"))
+    Assert.assertTrue(groupedIterator.hasNext)
+    Assert.assertEquals(groupedIterator.next.toVector, Vector("2", "2"))
+    Assert.assertFalse(groupedIterator.hasNext)
+    Assert.assertEquals(Try(groupedIterator.next).toString, Try(Iterator.empty.next()).toString)
+  }
+}