Source file gen.ml

module Random = PRNG.Splitmix.Pure
module Forest = Tree.Forest

(* Intuitively, a generator produces a single value from a
   pseudo-random generator, commonly represented by the type
   [Random.Pure.t -> 'a]. In this library, we extend this concept in
   two significant ways.

   1) Instead of generating a singular value, we generate a tree of
   values, modifying the type to [Random.Pure.t -> 'a Tree.t]. The
   tree's root represents the initially generated value, with its
   children representing smaller, "shrunk" candidates. These
   candidates are utilized by the shrinking algorithm to find smaller
   values that still meet certain criteria if the initially generated
   value does not (see [Tree.shrink]).

   The choice of [Random.Pure] for number generation is deliberate. It
   ensures that the [bind] function remains pure, maintaining the
   independence of values produced by a function [f] in [bind x f]
   from those generated by [x]. Consequently, re-executing [f] with
   smaller values of [x] will yield the same values as the original
   [x]. For instance, if [x] generates a positive number [n], and [f]
   subsequently generates [n] numbers, then calling [f] with [n' < n]
   will result in [f] producing the same first [n'] numbers.

   2) The evolution to [Random.Pure,t -> 'a Forest.t] facilitates more
   refined shrinkers, especially for recursive data types like
   lists. When shrinking recursive generators, it might be beneficial
   to bypass certain values. For example, when shrinking a list of
   size 3 (e.g., [1 2 3]), a shrinker might opt to test a size 2
   list. But which sublist should it choose? The default might be [1
   2], but employing [Forest.t] instead of [Tree.t] enables the
   shrinker to consider [1 3] or [2 3] as well. This approach implies
   that during shrinking, the generator might yield multiple values
   instead of just one. Continuing with the list example, after
   deciding to test a list of size 2, the first list value could be
   either 1 or 2, leading to the conceptualization of [Forest.t] as a
   sequence of [Tree.t]. Note, however, that a [Forest.t] is never
   empty.

   It's essential to maintain that the sequence's size should always
   be one upon value generation (refer to [run]). However, if the
   generator is invoked recursively during shrinking, the sequence
   size might increase. This behavior is not enforced by the type
   system, placing the onus on the generator's developer to uphold
   this invariant. *)
type 'a t = Random.t -> 'a Forest.t

let return : 'a -> 'a t = fun value _rs -> Forest.return value

let make : 'a -> ('a -> 'a Seq.t) -> 'a t =
 fun root make_children _rs -> Forest.make root make_children

let bind : 'a t -> ('a -> 'b t) -> 'b t =
 fun gen f rs ->
  (* Split guarantees that [rs_left] and [rs_right] are independent
     allowing [f] to be called multiple times with different values
     for [a] and still produces the same values. *)
  let rs_left, rs_right = Random.split rs in
  let forest = gen rs_left in
  Forest.bind forest (fun a -> f a rs_right)

let map : ('a -> 'b) -> 'a t -> 'b t = fun f gen rs -> Forest.map f (gen rs)

(* This function does a lookup on the generator given and always
   returns the generated value. It can be used to implement shrinkers.

   This lookup is correct only if the generator is given as the left
   parameter of a bind. *)
let root (gen : 'a t) f rs =
  (* The split call mimic what bind is doing so that the value given
     to the function [f] is indeed the one that would be produced with
     a bind. *)
  let rs_left, _ = Random.split rs in
  Forest.first (gen rs_left) |> Tree.root |> Fun.flip f rs

module Syntax = struct
  let ( let* ) x f = bind x f

  let ( let*! ) = root

  let return = return
end

(* When implementing shrinkers, the generator may produce a sequence
   of values. This function make a single generator out a sequence of
   generators. *)
let sequence : 'a t -> 'a t Seq.t -> 'a t =
 fun gen seq rs ->
  let gen = gen rs in
  let seq = Seq.map (fun gen -> gen rs) seq in
  Forest.sequence gen seq

(* This module can be used to define better shrinkers by defining a
   merging strategy. *)
module Merge = struct
  (* For abstraction purpose, we hide the merge function over
     trees. In the future, this module could be extended with other
     strategies. *)
  type 'a t = 'a Tree.t Seq.t -> 'a Tree.t Seq.t -> 'a Tree.t Seq.t

  let default = Seq.append

  let drop_left _ y = y

  let drop_right x _ = x

  let of_compare : compare:('a -> 'a -> int) -> 'a t =
   fun ~compare ->
    let compare left right = compare (Tree.root left) (Tree.root right) in
    Seq.sorted_merge compare
end

(* Set the merging behavior for all the trees defined.

   Since [sequence] does not change the merge strategy, if this
   function is called before calling [sequence], each tree of the
   forest may have there own merge strategy. *)
let with_merge : 'a Merge.t -> 'a t -> 'a t =
 fun merge gen rs ->
  Forest.map_tree (fun tree -> Tree.with_merge ~merge tree) (gen rs)

let z_range : ?origin:Z.t -> min:Z.t -> max:Z.t -> unit -> Z.t t =
 fun ?origin ~min ~max () rs ->
  let open Z.Compare in
  if max <= min then Forest.return min
  else
    let start =
      let upper_bound = Z.succ (Z.sub max min) in
      let rs = Obj.magic rs in
      Z.random_int_gen
        ~fill:(fun bytes pos len -> PRNG.Splitmix.State.bytes rs bytes pos len)
        upper_bound
    in
    let initial = Z.add min start in
    let origin =
      Option.value origin
        ~default:(if min <= Z.zero && Z.zero <= max then Z.zero else min)
    in
    Tree.binary_search ~initial ~origin () |> Forest.lift

let float_range :
       ?exhaustive_search_digits:int
    -> ?precision_digits:int
    -> ?origin:float
    -> min:float
    -> max:float
    -> unit
    -> float t =
 fun ?exhaustive_search_digits ?precision_digits ?origin ~min ~max () rs ->
  let origin =
    Option.value origin ~default:(if min <= 0. && 0. <= max then 0. else min)
  in
  if min >= max then return min rs
  else if max -. min <= 1. then
    let initial, _ = Random.float (max -. min) rs in
    Tree.fractional_search ?exhaustive_search_digits ?precision_digits ~initial
      ~origin ()
    |> Forest.lift
    |> Forest.map (fun x -> x +. min)
  else
    let rs, rs' = Random.split rs in
    let _, mini = Float.modf min in
    let _, maxi = Float.modf max in
    let originf, origini = Float.modf origin in
    let shift = Z.of_float mini in
    let forest =
      z_range
        ~origin:(Z.sub (Z.of_float origini) shift)
        ~min:Z.zero
        ~max:(Z.sub (Z.of_float maxi) shift)
        () rs
    in
    let fractional = Random.float 1. rs' |> fst in
    let ff, fi = Float.modf fractional in
    let fractional_forest =
      Tree.fractional_search ?exhaustive_search_digits ?precision_digits
        ~initial:ff ~origin:originf ()
      |> Forest.lift
    in
    Forest.bind forest (fun x ->
        let value = Z.add x (Z.of_float fi) |> Z.to_float in
        Forest.map
          (fun fractional ->
            Float.max min (value +. fractional +. min) |> Float.min max )
          fractional_forest )

let crunch i (gen : 'a t) : 'a t =
 fun rs ->
  let forest = gen rs in
  Forest.crunch i forest

let shrink = Tree.shrink

(* [t] is a runnable monad. Hence to run a generator, one needs to
   provide an initial state for the random generator. We check the
   validity of the generator by checking whether the [Forest] is a
   singleton. Having multiple trees in the forest is allowed only
   during shrinking, not while generated values. *)
let run ?(on_failure = failwith) gen state =
  let message =
    "[Gen.run] was called with an erroneous generator. The generator is \
     expected to return a single value. Instead: multiple values were \
     returned. You should probably fix your generator or provide a \
     [on_failure] argument to [Gen.run]."
  in
  let forest = gen state in
  if Forest.is_singleton forest then Forest.first forest else on_failure message