diff --git a/geny/src/geny/Generator.scala b/geny/src/geny/Generator.scala index 02f0e4d..068219e 100644 --- a/geny/src/geny/Generator.scala +++ b/geny/src/geny/Generator.scala @@ -111,6 +111,11 @@ trait Generator[+A]{ def flatten[V](implicit f: A => Generator[V]) = this.flatMap[V]((x: A) => f(x)) def slice(start: Int, end: Int): Generator[A] = new Generator.Sliced(this, start, end) + def grouped(size: Int): Generator[Iterable[A]] = sliding(size, step = size) + def sliding(size: Int, step: Int): Generator[Iterable[A]] = { + require(size >= 1 && step >= 1, f"size=$size%d and step=$step%d, but both must be positive") + new Generator.Sliding(this, size, step) + } def take(n: Int) = slice(0, n) def drop(n: Int) = slice(n, Int.MaxValue) def takeWhile(pred: A => Boolean): Generator[A] = new Generator.TakeWhile(this, pred) @@ -297,6 +302,38 @@ object Generator{ override def toString = s"$inner.map($func)" } + private class Sliding[+T](inner: Generator[T], size: Int, step: Int) extends Generator[Iterable[T]] { + def generate(f: Iterable[T] => Generator.Action) = { + val builder = Iterable.newBuilder[T] + val capacity = Math.max(size, step) + var n = 0 + builder.sizeHint(capacity) + var action: Generator.Action = Generator.Continue + inner.generate { + t => + if (n < capacity) { + builder += t + n += 1 + action + } + else { + action = f(builder.result.take(size)) + val remainder = builder.result.drop(step) + builder.clear + builder.sizeHint(capacity) + builder ++= remainder + builder += t + n = remainder.size + 1 + action + } + } + if (n > 0 && action == Generator.Continue) { + f(builder.result.take(size)) + } + else action + } + } + private class Sliced[+T](inner: Generator[T], start: Int, end: Int) extends Generator[T]{ def generate(f: T => Generator.Action) = { var count = 0 diff --git a/geny/test/src/geny/TestGenerator.scala b/geny/test/src/geny/TestGenerator.scala index 1a7a478..309e468 100644 --- a/geny/test/src/geny/TestGenerator.scala +++ b/geny/test/src/geny/TestGenerator.scala @@ -133,6 +133,20 @@ object TestGenerator extends TestSuite{ check(Generator.from[IndexedSeq, Int](0 until 10).drop(3), 3 until 10) check(Generator.from[IndexedSeq, Int](0 until 10).drop(-1), 0 until 10) } + test("sliding") { + check(Generator.from[IndexedSeq, Int](IndexedSeq.empty).grouped(2), Seq.empty) + check(Generator.from[IndexedSeq, Int](0 until 1).grouped(2), Seq(Seq(0))) + check(Generator.from[IndexedSeq, Int](0 until 10).grouped(1), (0 until 10).map(Seq.apply(_))) + check(Generator.from[IndexedSeq, Int](0 until 10).grouped(3), Seq(0 until 3, 3 until 6, 6 until 9, 9 until 10)) + check(Generator.from[IndexedSeq, Int](0 until 10).grouped(5), Seq(0 until 5, 5 until 10)) + check(Generator.from[IndexedSeq, Int](0 until 10).grouped(10), Seq(0 until 10)) + check(Generator.from[IndexedSeq, Int](0 until 10).grouped(15), Seq(0 until 10)) + + check(Generator.from[IndexedSeq, Int](0 until 10).sliding(2, 3), Seq(0 until 2, 3 until 5, 6 until 8, 9 until 10)) + check(Generator.from[IndexedSeq, Int](0 until 10).sliding(2, 15), Seq(0 until 2)) + check(Generator.from[IndexedSeq, Int](0 until 10).sliding(3, 2), Seq(0 until 3, 2 until 5, 4 until 7, 6 until 9, 8 until 10)) + check(Generator.from[IndexedSeq, Int](0 until 10).sliding(15, 2), Seq(0 until 10)) + } test("takeWhile"){ check(Generator.from[IndexedSeq, Int](0 until 10).takeWhile(_ < 5), 0 until 5) }