Add Gen to split seq into n partitions
japgolly opened this issue · comments
Add Gen to split seq into n partitions
Ahh now two similar things outside of Nyaya
// HashLogicTest.scala
def sizedGroups[C[X] <: SeqLike[X, C[X]], A](input: C[A])(implicit ss: SizeSpec): Gen[Vector[C[A]]] = { // TODO Add to Nyaya
val size = ss match {
case x@ SizeSpec.Default => x.gen
case x: SizeSpec.Exactly => x.gen
case SizeSpec.OneOf(ns) => Gen.chooseIndexed_!(ns)
}
Gen { ctx =>
val groups = Vector.newBuilder[C[A]]
@tailrec def go(as: C[A]): Unit =
if (as.nonEmpty) {
val n = size.run(ctx)
groups += as.take(n)
go(as.drop(n))
}
go(input)
groups.result()
}
}
// ProjectStateTest.scala
// TODO Move into Nyaya
def Gen_batches[A](as: Vector[A], partitionSize: Range.Inclusive, keepRemainder: Boolean = true): Gen[Vector[Vector[A]]] = {
val genSize = Gen.chooseInt(partitionSize.min, partitionSize.max)
Gen { ctx =>
val b = Vector.newBuilder[Vector[A]]
@tailrec def go(rem: Vector[A]): Unit =
if (rem.isEmpty)
()
else if (rem.length >= partitionSize.min) {
val n = genSize.run(ctx)
b += rem.take(n)
go(rem.drop(n))
} else if (keepRemainder)
b += rem
go(as)
b.result()
}
}