japgolly / nyaya

Random Data Generation and/or Property Testing in Scala & Scala.JS.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

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()
    }
  }