Segmented repartition

core/jvm/src/test/scala/fs2/StreamSpec.scala

@@ -107,6 +107,27 @@ class StreamSpec extends Fs2Spec with Inside {
         IndexedSeq.range(0, 100)
     }
 
+    "repartition" in {
+      Stream("Lore", "m ip", "sum dolo", "r sit amet").repartition(s => Chunk.array(s.split(" "))).toList shouldBe
+        List("Lorem", "ipsum", "dolor", "sit", "amet")
+      Stream("hel", "l", "o Wor", "ld").repartition(s => Chunk.indexedSeq(s.grouped(2).toVector)).toList shouldBe
+        List("he", "ll", "o ", "Wo", "rl", "d")
+      Stream.empty.covaryOutput[String].repartition(_ => Chunk.empty).toList shouldBe List()
+      Stream("hello").repartition(_ => Chunk.empty).toList shouldBe List()
+
+      def input = Stream("ab").repeat
+      def ones(s: String) = Chunk vector s.grouped(1).toVector
+      input.take(2).repartition(ones).toVector shouldBe Vector("a", "b", "a", "b")
+      input.take(4).repartition(ones).toVector shouldBe Vector("a", "b", "a", "b", "a", "b", "a", "b")
+      input.repartition(ones).take(2).toVector shouldBe Vector("a", "b")
+      input.repartition(ones).take(4).toVector shouldBe Vector("a", "b", "a", "b")
+      Stream.emits(input.take(4).toVector).repartition(ones).toVector shouldBe Vector("a", "b", "a", "b", "a", "b", "a", "b")
+
+      Stream(1, 2, 3, 4, 5).repartition(i => Chunk(i, i)).toList shouldBe List(1, 3, 6, 10, 15, 15)
+
+      Stream(1, 10, 100).repartition(i => Segment.from(i).map(_.toInt).take(1000).toChunk).take(4).toList shouldBe List(1, 2, 3, 4)
+    }
+
     "translate" in forAll { (s: PureStream[Int]) =>
       runLog(s.get.covary[IO].flatMap(i => Stream.eval(IO.pure(i))).translate(cats.arrow.FunctionK.id[IO])) shouldBe
       runLog(s.get)

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

@@ -2013,6 +2013,29 @@ object Stream {
     def reduceSemigroup(implicit S: Semigroup[O]): Stream[F, O] =
       self.reduce(S.combine(_, _))
 
+    /**
+      * Repartitions the input with the function `f`. On each step `f` is applied
+      * to the input and all elements but the last of the resulting sequence
+      * are emitted. The last element is then appended to the next input using the
+      * Semigroup `S`.
+      *
+      * @example {{{
+      * scala> import cats.implicits._
+      * scala> Stream("Hel", "l", "o Wor", "ld").repartition(s => Chunk.array(s.split(" "))).toList
+      * res0: List[String] = List(Hello, World)
+      * }}}
+      */
+    def repartition(f: O => Chunk[O])(implicit S: Semigroup[O]): Stream[F,O] = {
+      pull.scanSegments(Option.empty[O]) { (carry, segment) =>
+        segment.scan((Segment.empty[O], carry)) { case ((_, carry), o) =>
+          val o2: O = carry.fold(o)(S.combine(_, o))
+          val partitions: Chunk[O] = f(o2)
+          if (partitions.isEmpty) partitions -> None
+          else partitions.take(partitions.size - 1).voidResult -> Some(partitions.strict.last)
+        }.flatMap { case (out, carry) => out }.mapResult { case ((out, carry), unit) => carry }
+      }.flatMap { case Some(carry) => Pull.output1(carry); case None => Pull.done }.stream
+    }
+
     /**
      * Repeatedly invokes `using`, running the resultant `Pull` each time, halting when a pull
      * returns `None` instead of `Some(nextStream)`.