<commit_before>package spire.random

import spire.algebra._
import spire.implicits._
import spire.math._
import spire.syntax.cfor._

import scala.collection.mutable.{ArrayBuffer, Builder}
import scala.collection.SeqLike
import scala.collection.generic.CanBuildFrom
import scala.annotation.tailrec
import scala.{specialized => spec}
import scala.reflect.ClassTag

sealed trait Op[+A] {

  def flatMap[B](f: A => Op[B]): Op[B] =
    this match {
      case FlatMap(a, g) => FlatMap(a, (x: Any) => g(x).flatMap(f))
      case o => FlatMap(o, f)
    }

  def map[B](f: A => B): Op[B] =
    flatMap(a => Const(f(a)))

  def resume(gen: Generator): Either[() => Op[A], A] =
    this match {
      case Const(a) =>
        Right(a)
      case More(k) =>
        Left(k)
      case Next(f) =>
        Right(f(gen))
      case FlatMap(a, f) =>
        a match {
          case Const(x) => f(x).resume(gen)
          case More(k) => Left(() => FlatMap(k(), f))
          case Next(g) => f(g(gen)).resume(gen)
          case FlatMap(b, g) => FlatMap(b, (x: Any) => FlatMap(g(x), f)).resume(gen)
        }
    }

  def run(gen: Generator): A = {
    def loop(e: Either[() => Op[A], A]): A = e match {
      case Right(a) => a
      case Left(k) => loop(k().resume(gen))
    }
    loop(resume(gen))
  }
}

case class Const[+A](a: A) extends Op[A]
case class More[+A](k: () => Op[A]) extends Op[A]
case class Next[+A](f: Generator => A) extends Op[A]
case class FlatMap[A, +B](sub: Op[A], k: A => Op[B]) extends Op[B]

object Random extends RandomCompanion[rng.Cmwc5] {
  def initGenerator() = rng.Cmwc5.fromTime()

  def spawn[B](op: Op[B]) = new RandomCmwc5(op)
}

trait RandomCompanion[G <: Generator] { self =>
  type R[X] = Random[X, G]

  def initGenerator(): G //IO

  def generatorFromSeed(seed: Seed): G = {
    val gen = initGenerator()
    gen.setSeedBytes(seed.bytes)
    gen
  }

  def spawn[B](op: Op[B]): R[B]

  def next[B](f: Generator => B): R[B] = spawn(Next(f))

  def constant[B](b: B) = spawn(Const(b))

  def unit = constant(Unit)
  def boolean = next(_.nextBoolean)
  def byte = next(_.nextInt.toByte)
  def short = next(_.nextInt.toShort)
  def char = next(_.nextInt.toChar)

  def int = next(_.nextInt)
  def int(n: Int) = next(_.nextInt(n))
  def int(n1: Int, n2: Int) = next(_.nextInt(n1, n2))

  def float = next(_.nextFloat)
  def long = next(_.nextLong)
  def double = next(_.nextDouble)

  def string(size: Size): R[String] =
    size.random(this).flatMap(stringOfSize)

  def stringOfSize(n: Int): Random[String, G] =
    char.foldLeftOfSize(n)(new StringBuilder) { (sb, c) => sb.append(c); sb }.map(_.toString)

  implicit class RandomOps[A](lhs: R[A]) {
    def collection[CC[_]](size: Size)(implicit cbf: CanBuildFrom[CC[A], A, CC[A]]): Random[CC[A], G] =
      size.random(self).flatMap(collectionOfSize(_))

    def collectionOfSize[CC[_]](n: Int)(implicit cbf: CanBuildFrom[CC[A], A, CC[A]]): Random[CC[A], G] =
      foldLeftOfSize(n)(cbf()) { (b, a) => b += a; b }.map(_.result)

    def foldLeftOfSize[B](n: Int)(init: => B)(f: (B, A) => B): Random[B, G] = {
      def loop(n: Int, ma: Op[A]): Op[B] =<commit_msg>add fromDist method<commit_after>package spire.random

import spire.algebra._
import spire.implicits._
import spire.math._
import spire.syntax.cfor._

import scala.collection.mutable.{ArrayBuffer, Builder}
import scala.collection.SeqLike
import scala.collection.generic.CanBuildFrom
import scala.annotation.tailrec
import scala.{specialized => spec}
import scala.reflect.ClassTag

sealed trait Op[+A] {

  def flatMap[B](f: A => Op[B]): Op[B] =
    this match {
      case FlatMap(a, g) => FlatMap(a, (x: Any) => g(x).flatMap(f))
      case o => FlatMap(o, f)
    }

  def map[B](f: A => B): Op[B] =
    flatMap(a => Const(f(a)))

  def resume(gen: Generator): Either[() => Op[A], A] =
    this match {
      case Const(a) =>
        Right(a)
      case More(k) =>
        Left(k)
      case Next(f) =>
        Right(f(gen))
      case FlatMap(a, f) =>
        a match {
          case Const(x) => f(x).resume(gen)
          case More(k) => Left(() => FlatMap(k(), f))
          case Next(g) => f(g(gen)).resume(gen)
          case FlatMap(b, g) => FlatMap(b, (x: Any) => FlatMap(g(x), f)).resume(gen)
        }
    }

  def run(gen: Generator): A = {
    def loop(e: Either[() => Op[A], A]): A = e match {
      case Right(a) => a
      case Left(k) => loop(k().resume(gen))
    }
    loop(resume(gen))
  }
}

case class Const[+A](a: A) extends Op[A]
case class More[+A](k: () => Op[A]) extends Op[A]
case class Next[+A](f: Generator => A) extends Op[A]
case class FlatMap[A, +B](sub: Op[A], k: A => Op[B]) extends Op[B]

object Random extends RandomCompanion[rng.Cmwc5] {
  def initGenerator() = rng.Cmwc5.fromTime()

  def spawn[B](op: Op[B]) = new RandomCmwc5(op)
}

trait RandomCompanion[G <: Generator] { self =>
  type R[X] = Random[X, G]

  def initGenerator(): G //IO

  def generatorFromSeed(seed: Seed): G = {
    val gen = initGenerator()
    gen.setSeedBytes(seed.bytes)
    gen
  }

  def spawn[B](op: Op[B]): R[B]

  def next[B](f: Generator => B): R[B] = spawn(Next(f))

  def fromDist[B](dist: Dist[B]): R[B] = spawn(Next(g => dist(g)))

  def constant[B](b: B) = spawn(Const(b))

  def unit = constant(Unit)
  def boolean = next(_.nextBoolean)
  def byte = next(_.nextInt.toByte)
  def short = next(_.nextInt.toShort)
  def char = next(_.nextInt.toChar)

  def int = next(_.nextInt)
  def int(n: Int) = next(_.nextInt(n))
  def int(n1: Int, n2: Int) = next(_.nextInt(n1, n2))

  def float = next(_.nextFloat)
  def long = next(_.nextLong)
  def double = next(_.nextDouble)

  def string(size: Size): R[String] =
    size.random(this).flatMap(stringOfSize)

  def stringOfSize(n: Int): Random[String, G] =
    char.foldLeftOfSize(n)(new StringBuilder) { (sb, c) => sb.append(c); sb }.map(_.toString)

  implicit class RandomOps[A](lhs: R[A]) {
    def collection[CC[_]](size: Size)(implicit cbf: CanBuildFrom[CC[A], A, CC[A]]): Random[CC[A], G] =
      size.random(self).flatMap(collectionOfSize(_))

    def collectionOfSize[CC[_]](n: Int)(implicit cbf: CanBuildFrom[CC[A], A, CC[A]]): Random[CC[A], G] =
      foldLeftOfSize(n)(cbf()) { (b, a) => b += a; b }.map(_.result)

    def foldLeftOfSize[B](n: Int)(init: => B)(f: (B, A) => B): Random[B, G] = {
      def loop(n: Int, ma: Op[A]): Op[B] =
{'fraction_non_alphanumeric': 0.19021397105097546, 'fraction_numerical': 0.0022026431718061676, 'mean_word_length': 3.384137931034483, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 64, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4254763', 'n_tokens_mistral': 2535, 'n_tokens_neox': 2476, 'n_words': 584}