Source file gen.ml

(*---------------------------------------------------------------------------
   Copyright (c) 2013-2017 Guillaume Bury, Simon Cruanes, Vincent Hugot,
                           Jan Midtgaard
   Copyright (c) 2026 Invariant Systems. All rights reserved.
   SPDX-License-Identifier: BSD-2-Clause

   Generator design and several implementations (nat distribution, float bit
   manipulation, option/result ratios) derived from QCheck2
   (https://github.com/c-cube/qcheck).
  ---------------------------------------------------------------------------*)

type 'a t = Random.State.t -> 'a Tree.t

(* ───── Combinators ───── *)

let pure x _ = Tree.pure x
let map f gen st = Tree.map f (gen st)
let ( >|= ) gen f = map f gen

let ap fgen xgen st =
  let st' = Random.State.split st in
  Tree.ap (fgen st) (xgen st')

let bind gen f st =
  let st' = Random.State.split st in
  Tree.bind (gen st) (fun x -> f x (Random.State.copy st'))

let ( >>= ) = bind
let ( let+ ) gen f = map f gen

let ( and+ ) a b st =
  let st' = Random.State.split st in
  Tree.liftA2 (fun x y -> (x, y)) (a st) (b st')

let ( let* ) = bind

(* ───── Primitives ───── *)

let make_primitive ~gen ~shrink st = Tree.make_primitive shrink (gen st)
let unit _ = Tree.pure ()

let bool st =
  if Random.State.bool st then Tree.Tree (true, Seq.return (Tree.pure false))
  else Tree.pure false

(* Random.State.int bound is 2^30-1 on all platforms, not max_int *)
let max_random_int = (1 lsl 30) - 1

let int_pos_raw st =
  if Sys.word_size = 32 then Random.State.bits st
  else
    (* Build a full-width positive int from 30+30+2 random bits. *)
    let top2_mask = 0b11 in
    let left = (Random.State.bits st land top2_mask) lsl 60 in
    let middle = Random.State.bits st lsl 30 in
    let right = Random.State.bits st in
    left lor middle lor right

let int_pos st =
  let x = int_pos_raw st in
  Tree.make_primitive (fun n -> Shrink.int_towards 0 n) x

let int st =
  if Random.State.bool st then Tree.map (fun n -> -n - 1) (int_pos st)
  else int_pos st

let int_bound upper st =
  if upper < 0 then invalid_arg "Gen.int_bound: upper < 0";
  if upper <= max_random_int - 1 then Random.State.int st (upper + 1)
  else if upper = max_int then int_pos_raw st
  else int_pos_raw st mod (upper + 1)

let pick_origin_in_range ~low ~high ~goal =
  if goal < low then low else if goal > high then high else goal

let resolve_origin ~loc ~low ~high ~origin =
  if origin < low then invalid_arg (loc ^ ": origin < low")
  else if origin > high then invalid_arg (loc ^ ": origin > high")
  else origin

let int_range ?origin low high st =
  if high < low then invalid_arg "Gen.int_range: high < low";
  let n =
    if low = high then low
    else if low >= 0 || high < 0 then
      (* Range is entirely non-negative or entirely negative *)
      low + int_bound (high - low) st
    else
      (* Range spans zero: choose side proportionally, then sample within it *)
      let f_low = float_of_int low in
      let f_high = float_of_int high in
      let ratio = -.f_low /. (1.0 +. f_high -. f_low) in
      if Random.State.float st 1.0 <= ratio then
        if low = min_int then -int_pos_raw st - 1
        else -int_bound (pred (-low)) st - 1
      else int_bound high st
  in
  let default_origin = pick_origin_in_range ~low ~high ~goal:0 in
  let origin =
    resolve_origin ~loc:"Gen.int_range" ~low ~high
      ~origin:(Option.value origin ~default:default_origin)
  in
  Tree.make_primitive (fun x -> Shrink.int_towards origin x) n

let nat st =
  let p = Random.State.float st 1.0 in
  let x =
    if p < 0.5 then Random.State.int st 10
    else if p < 0.75 then Random.State.int st 100
    else if p < 0.95 then Random.State.int st 1_000
    else Random.State.int st 10_000
  in
  Tree.make_primitive (fun n -> Shrink.int_towards 0 n) x

let small_int st =
  if Random.State.bool st then nat st else Tree.map Int.neg (nat st)

let int32 st =
  (* Random.State.bits yields 30 bits; two calls cover the full 32-bit range *)
  let low = Int32.of_int (Random.State.bits st) in
  let high = Int32.of_int (Random.State.bits st land 0x3) in
  let bits = Int32.logor low (Int32.shift_left high 30) in
  Tree.make_primitive (fun n -> Shrink.int32_towards 0l n) bits

let pick_origin_int32 ~low ~high ~goal =
  if goal < low then low else if goal > high then high else goal

let int32_range ?origin low high st =
  if high < low then invalid_arg "Gen.int32_range: high < low";
  (* Compute range in int64 to avoid int32 overflow *)
  let range = Int64.sub (Int64.of_int32 high) (Int64.of_int32 low) in
  let range = Int64.add range 1L in
  let n =
    if Int64.compare range (Int64.of_int max_random_int) <= 0 then
      Int32.add low (Int32.of_int (Random.State.int st (Int64.to_int range)))
    else
      (* Range exceeds single random call capacity; use float approximation *)
      let f_range = Int64.to_float range in
      let offset = Int64.of_float (Random.State.float st f_range) in
      Int32.add low (Int64.to_int32 offset)
  in
  let origin =
    pick_origin_int32 ~low ~high ~goal:(Option.value origin ~default:0l)
  in
  Tree.make_primitive (fun x -> Shrink.int32_towards origin x) n

let int64 st =
  (* Random.State.bits yields 30 bits; three calls cover the full 64-bit range *)
  let low = Int64.of_int (Random.State.bits st) in
  let mid = Int64.shift_left (Int64.of_int (Random.State.bits st)) 30 in
  let high =
    Int64.shift_left (Int64.of_int (Random.State.bits st land 0xF)) 60
  in
  let bits = Int64.(logor high (logor mid low)) in
  Tree.make_primitive (fun n -> Shrink.int64_towards 0L n) bits

let int64_nonneg_raw st =
  (* Random.State.bits yields 30 bits; 30+30+3 gives a full non-negative
     63-bit value. *)
  let low = Int64.of_int (Random.State.bits st) in
  let mid = Int64.shift_left (Int64.of_int (Random.State.bits st)) 30 in
  let high =
    Int64.shift_left (Int64.of_int (Random.State.bits st land 0x7)) 60
  in
  Int64.(logor high (logor mid low))

let int64_bound upper st =
  if upper < 0L then invalid_arg "Gen.int64_bound: upper < 0";
  if Int64.equal upper Int64.max_int then int64_nonneg_raw st
  else
    let bound = Int64.add upper 1L in
    Int64.rem (int64_nonneg_raw st) bound

let pick_origin_int64 ~low ~high ~goal =
  if goal < low then low else if goal > high then high else goal

let int64_range ?origin low high st =
  if high < low then invalid_arg "Gen.int64_range: high < low";
  let n =
    if low >= 0L || high < 0L then
      let offset = int64_bound (Int64.sub high low) st in
      Int64.add low offset
    else
      (* Range spans zero: choose side proportionally, then sample that side
         with integer arithmetic. *)
      let f_low = Int64.to_float low in
      let f_high = Int64.to_float high in
      let ratio = -.f_low /. (1.0 +. f_high -. f_low) in
      if Random.State.float st 1.0 <= ratio then
        if Int64.equal low Int64.min_int then
          Int64.sub (Int64.neg (int64_nonneg_raw st)) 1L
        else
          let offset = int64_bound (Int64.pred (Int64.neg low)) st in
          Int64.neg (Int64.succ offset)
      else int64_bound high st
  in
  let origin =
    pick_origin_int64 ~low ~high ~goal:(Option.value origin ~default:0L)
  in
  Tree.make_primitive (fun x -> Shrink.int64_towards origin x) n

let nativeint st = Tree.map Nativeint.of_int (int st)

let float st =
  (* Assemble 64 random bits and reinterpret as IEEE 754 double to cover the
     full float range including NaN, infinity, and subnormals *)
  let bits =
    let left = Int64.(shift_left (of_int (Random.State.bits st land 0xF)) 60) in
    let middle = Int64.(shift_left (of_int (Random.State.bits st)) 30) in
    let right = Int64.of_int (Random.State.bits st) in
    Int64.(logor left (logor middle right))
  in
  let x = Int64.float_of_bits bits in
  Tree.make_primitive (fun f -> Shrink.float_towards 0.0 f) x

let float_range ?origin low high st =
  if high < low then invalid_arg "Gen.float_range: high < low";
  if high -. low > max_float then
    invalid_arg "Gen.float_range: high -. low > max_float";
  let x = low +. Random.State.float st (high -. low) in
  let default_origin = pick_origin_in_range ~low ~high ~goal:0.0 in
  let origin =
    resolve_origin ~loc:"Gen.float_range" ~low ~high
      ~origin:(Option.value origin ~default:default_origin)
  in
  Tree.make_primitive (fun f -> Shrink.float_towards origin f) x

let char_range ?origin low high =
  let lo = Char.code low in
  let hi = Char.code high in
  let origin = Option.map Char.code origin in
  map Char.chr (int_range ?origin lo hi)

let char st =
  let c = Random.State.int st 256 in
  let shrink i = Shrink.int_towards (Char.code 'a') i in
  Tree.map Char.chr (Tree.make_primitive shrink c)

(* ───── Containers ───── *)

let option ?(ratio = 0.85) gen st =
  let p = Random.State.float st 1.0 in
  if p < 1.0 -. ratio then Tree.pure None else Tree.opt (gen st)

let result ?(ratio = 0.75) ok_gen err_gen st =
  let p = Random.State.float st 1.0 in
  if p < 1.0 -. ratio then Tree.map (fun e -> Error e) (err_gen st)
  else Tree.map (fun o -> Ok o) (ok_gen st)

let either ?(ratio = 0.5) left_gen right_gen st =
  let p = Random.State.float st 1.0 in
  if p < ratio then Tree.map (fun x -> Either.Left x) (left_gen st)
  else Tree.map (fun x -> Either.Right x) (right_gen st)

let list_size size_gen gen st =
  let st' = Random.State.split st in
  Tree.bind (size_gen st) (fun size ->
      let st' = Random.State.copy st' in
      let rec build n acc =
        if n <= 0 then
          List.fold_left
            (fun acc elt_tree -> Tree.liftA2 List.cons elt_tree acc)
            (Tree.pure []) acc
        else build (n - 1) (gen st' :: acc)
      in
      build size [])

(* Like QCheck2 list/list_small: use only size root, then apply a custom
   bisection shrink strategy over list structure. *)
let list_ignore_size_tree size_gen gen st =
  let st' = Random.State.split st in
  let size = Tree.root (size_gen st) in
  let st' = Random.State.copy st' in
  let rec build n acc =
    if n <= 0 then
      let l = List.rev acc in
      Tree.Tree (List.map Tree.root l, Tree.build_list_shrink_tree l)
    else build (n - 1) (gen st' :: acc)
  in
  build size []

let list gen = list_ignore_size_tree nat gen
let array gen st = Tree.map Array.of_list (list gen st)

let pair a b =
  let+ x = a and+ y = b in
  (x, y)

let triple a b c =
  let+ x, y = pair a b and+ z = c in
  (x, y, z)

let quad a b c d =
  let+ x, y, z = triple a b c and+ w = d in
  (x, y, z, w)

(* ───── Choice ───── *)

let oneof gens st =
  match gens with
  | [] -> invalid_arg "Gen.oneof: empty list"
  | _ ->
      let i_tree =
        Tree.make_primitive
          (fun i -> Shrink.int_towards 0 i)
          (Random.State.int st (List.length gens))
      in
      Tree.bind i_tree (fun i -> (List.nth gens i) st)

let oneofl xs st =
  match xs with
  | [] -> invalid_arg "Gen.oneofl: empty list"
  | _ ->
      let i_tree =
        Tree.make_primitive
          (fun i -> Shrink.int_towards 0 i)
          (Random.State.int st (List.length xs))
      in
      Tree.map (List.nth xs) i_tree

let frequency weighted_gens st =
  match weighted_gens with
  | [] -> invalid_arg "Gen.frequency: empty list"
  | _ ->
      let total = List.fold_left (fun acc (w, _) -> acc + w) 0 weighted_gens in
      if total < 1 then invalid_arg "Gen.frequency: total weight < 1";
      let pick = Random.State.int st total in
      let rec choose acc = function
        | [] -> assert false
        | (w, g) :: rest ->
            if pick < acc + w then g st else choose (acc + w) rest
      in
      choose 0 weighted_gens

(* ───── Strings ───── *)

let string_size size_gen char_gen st =
  let size_tree = size_gen st in
  Tree.bind size_tree (fun size ->
      let st' = Random.State.copy st in
      let chars = List.init size (fun _ -> char_gen st') in
      Tree.map
        (fun char_list ->
          let a = Array.of_list char_list in
          String.init (Array.length a) (Array.get a))
        (Tree.sequence_list chars))

let string_of char_gen = string_size nat char_gen
let string = string_of char
let bytes st = map Bytes.of_string string st

(* ───── Size Control ───── *)

let sized f = bind nat f

(* ───── Shrink Control ───── *)

let no_shrink gen st =
  let (Tree.Tree (x, _)) = gen st in
  Tree.pure x

let add_shrink_invariant p gen st = Tree.add_shrink_invariant p (gen st)

(* ───── Recursive Generators ───── *)

let delay f st = f () st

let fix f =
  let rec gen st = f gen st in
  gen

(* ───── Search ───── *)

let find ?(count = 100) ~f gen st =
  let rec loop n =
    if n <= 0 then None
    else
      let (Tree.Tree (x, _)) = gen st in
      if f x then Some x else loop (n - 1)
  in
  loop count