<commit_before>package org.scalameter



import collection._
import org.scalameter.picklers.Pickler
import org.scalameter.picklers.Implicits._



/** Base class for ScalaMeter generators.
 *
 *  Generators create warmup and benchmark datasets for running benchmarks.
 *  They support neat combinator-style syntax for composing more complex
 *  generators out of simpler ones.
 */
abstract class Gen[T] extends Serializable {
  self =>

  def map[S](f: T => S): Gen[S] = new Gen[S] {
    def warmupset = for (x <- self.warmupset) yield f(x)
    def dataset = for (params <- self.dataset) yield params
    def generate(params: Parameters) = f(self.generate(params))
  }

  def flatMap[S](f: T => Gen[S]): Gen[S] = new Gen[S] {
    def warmupset = for {
      x <- self.warmupset
      y <- f(x).warmupset
    } yield y
    def dataset = for {
      selfparams <- self.dataset
      x = self.generate(selfparams)
      thatparams <- f(x).dataset
    } yield selfparams ++ thatparams
    def generate(params: Parameters) = {
      val x = self.generate(params)
      val mapped = f(x)
      mapped.generate(params)
    }
  }

  def zip[S](that: Gen[S]): Gen[(T, S)] = for {
    x <- self
    y <- that
  } yield (x, y)

  def rename(mapping: (String, String)*): Gen[T] = new Gen[T] {
    val reverseMapping = mapping.map(kv => (kv._2, kv._1))
    def warmupset = self.warmupset
    def dataset = self.dataset.map(params => params map {
      case (k, v) => (mapping.toMap.applyOrElse(k, (k: String) => k), v)
    })
    def generate(params: Parameters) = self.generate(params map {
      case (k, v) => (reverseMapping.toMap.applyOrElse(k, (k: String) => k), v)
    })
  }

  def warmupset: Iterator[T]

  def dataset: Iterator[Parameters]

  def generate(params: Parameters): T

  def cached: Gen[T] = new Gen[T] {
    @transient lazy val cachedWarmupset = self.warmupset.toSeq
    @transient lazy val cachedDataset = self.dataset.map(p => (p, self.generate(p))).toMap

    def axes = cachedDataset.head._1.axisData.keys.toSet
    def warmupset = cachedWarmupset.iterator
    def dataset = cachedDataset.keys.iterator
    def generate(params: Parameters) = {
      val desiredParams = Parameters(params.axisData.filterKeys(k => axes.contains(k)).toSeq: _*)
      cachedDataset(desiredParams)
    }
  }

}


object Gen {

  def unit(axisName: String): Gen[Unit] = new Gen[Unit] {
    def warmupset = Iterator.single(())
    def dataset = Iterator.single(Parameters((Parameter[Unit](axisName), ())))
    def generate(params: Parameters) = params[Unit](axisName)
  }

  def single[T: Pickler](axisName: String)(v: T): Gen[T] = enumeration(axisName)(v)
<commit_msg>Use original dataset keys in cached generators<commit_after>package org.scalameter



import collection._
import org.scalameter.picklers.Pickler
import org.scalameter.picklers.Implicits._



/** Base class for ScalaMeter generators.
 *
 *  Generators create warmup and benchmark datasets for running benchmarks.
 *  They support neat combinator-style syntax for composing more complex
 *  generators out of simpler ones.
 */
abstract class Gen[T] extends Serializable {
  self =>

  def map[S](f: T => S): Gen[S] = new Gen[S] {
    def warmupset = for (x <- self.warmupset) yield f(x)
    def dataset = for (params <- self.dataset) yield params
    def generate(params: Parameters) = f(self.generate(params))
  }

  def flatMap[S](f: T => Gen[S]): Gen[S] = new Gen[S] {
    def warmupset = for {
      x <- self.warmupset
      y <- f(x).warmupset
    } yield y
    def dataset = for {
      selfparams <- self.dataset
      x = self.generate(selfparams)
      thatparams <- f(x).dataset
    } yield selfparams ++ thatparams
    def generate(params: Parameters) = {
      val x = self.generate(params)
      val mapped = f(x)
      mapped.generate(params)
    }
  }

  def zip[S](that: Gen[S]): Gen[(T, S)] = for {
    x <- self
    y <- that
  } yield (x, y)

  def rename(mapping: (String, String)*): Gen[T] = new Gen[T] {
    val reverseMapping = mapping.map(kv => (kv._2, kv._1))
    def warmupset = self.warmupset
    def dataset = self.dataset.map(params => params map {
      case (k, v) => (mapping.toMap.applyOrElse(k, (k: String) => k), v)
    })
    def generate(params: Parameters) = self.generate(params map {
      case (k, v) => (reverseMapping.toMap.applyOrElse(k, (k: String) => k), v)
    })
  }

  def warmupset: Iterator[T]

  def dataset: Iterator[Parameters]

  def generate(params: Parameters): T

  def cached: Gen[T] = new Gen[T] {
    @transient lazy val cachedWarmupset = self.warmupset.toSeq
    @transient lazy val cachedDataset = self.dataset.map(p => (p, self.generate(p))).toMap

    def axes = cachedDataset.head._1.axisData.keys.toSet
    def warmupset = cachedWarmupset.iterator
    def dataset = self.dataset
    def generate(params: Parameters) = {
      val desiredParams = Parameters(params.axisData.filterKeys(k => axes.contains(k)).toSeq: _*)
      cachedDataset(desiredParams)
    }
  }

}


object Gen {

  def unit(axisName: String): Gen[Unit] = new Gen[Unit] {
    def warmupset = Iterator.single(())
    def dataset = Iterator.single(Parameters((Parameter[Unit](axisName), ())))
    def generate(params: Parameters) = params[Unit](axisName)
  }

  def single[T: Pickler](axisName: String)(v: T): Gen[T] = enumeration(axisName)(v)

{'fraction_non_alphanumeric': 0.13629546306711662, 'fraction_numerical': 0.0011248593925759281, 'mean_word_length': 3.2885852090032155, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 25, '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': '918315', 'n_tokens_mistral': 1823, 'n_tokens_neox': 1736, 'n_words': 462}