diff --git a/Plausible/ArbitraryFueled.lean b/Plausible/ArbitraryFueled.lean index 0ea3a6c6..23a5732f 100644 --- a/Plausible/ArbitraryFueled.lean +++ b/Plausible/ArbitraryFueled.lean @@ -29,8 +29,4 @@ class ArbitraryFueled (α : Type) where meta instance [ArbitraryFueled α] : Arbitrary α where arbitrary := Gen.sized ArbitraryFueled.arbitraryFueled -/-- Raised when a fueled generator fails due to insufficient fuel. -/ -meta def Gen.outOfFuel : GenError := - .genError "out of fuel" - end Plausible diff --git a/Plausible/Gen.lean b/Plausible/Gen.lean index 02f1613a..6db97916 100644 --- a/Plausible/Gen.lean +++ b/Plausible/Gen.lean @@ -109,9 +109,13 @@ The code for these combinators closely mirrors those used in Rocq/Coq QuickChick -/ +/-- Raised when a fueled generator fails due to insufficient fuel. -/ +meta def Gen.outOfFuel : GenError := + .genError "out of fuel" + /-- `pick default xs n` chooses a weight & a generator `(k, gen)` from the list `xs` such that `n < k`. If `xs` is empty, the `default` generator with weight 0 is returned. -/ -private def pick (default : Gen α) (xs : List (Nat × Gen α)) (n : Nat) : Nat × Gen α := +def pick (default : Gen α) (xs : List (Nat × Gen α)) (n : Nat) : Nat × Gen α := match xs with | [] => (0, default) | (k, x) :: xs => @@ -120,6 +124,40 @@ private def pick (default : Gen α) (xs : List (Nat × Gen α)) (n : Nat) : Nat else pick default xs (n - k) +instance inhabitedGen : Inhabited (Gen α) where + default := throw <| .genError "inhabitedWitness" + + +private def pickDrop (xs : List (Nat × Gen α)) (n : Nat) : Nat × Gen α × List (Nat × Gen α) := + let fail : GenError := .genError "Plausible.Chamelean.Gen.pickDrop: out of options." + match xs with + | [] => (0, throw fail, []) + | (k, x) :: xs => + if n < k then + (k, x, xs) + else + let (k', x', xs') := pickDrop xs (n - k) + (k', x', (k, x)::xs') + +private def sumFst (gs : List (Nat × α)) : Nat := List.sum <| List.map Prod.fst gs + +/- Helper function for `backtrack` which picks one out of `total` generators with some initial amount of `fuel` -/ +private def backtrackFuel {α : Type u} (fuel : Nat) (total : Nat) (gs : List (Nat × Gen α)) : Gen α := + match fuel with + | .zero => throw Gen.outOfFuel + | .succ fuel' => do + let n ← (Gen.choose Nat 0 (total - 1) (by omega)).up + let (k, g, gs') := pickDrop gs n.down + -- Try to generate a value using `g`, if it fails, backtrack with `fuel'` + -- and pick one out of the `total - k` remaining generators + tryCatch g (fun _ => backtrackFuel fuel' (total - k) gs') + +/-- Tries all generators until one returns a `Some` value or all the generators failed once with `None`. + The generators are picked at random according to their weights (like `frequency` in Haskell QuickCheck), + and each generator is run at most once. -/ +def backtrack (gs : List (Nat × Gen α)) : Gen α := + backtrackFuel (gs.length) (sumFst gs) gs + /-- Picks one of the generators in `gs` at random, returning the `default` generator if `gs` is empty. @@ -220,7 +258,7 @@ partial def Gen.runUntil {α : Type} (attempts : Option Nat := .none) (x : Gen where repeatGen (attempts : Option Nat) (x : Gen α) : Gen α := match attempts with - | .some 0 => throw <| GenError.genError "Gen.runUtil: Out of attempts" + | .some 0 => throw <| GenError.genError "Gen.runUntil: Out of attempts" | _ => try x catch | GenError.genError _ => do @@ -230,19 +268,23 @@ partial def Gen.runUntil {α : Type} (attempts : Option Nat := .none) (x : Gen | .some n => .some (n-1) | .none => .none -private def test : Gen Nat := +private def test (n : Nat) : Gen Nat := do let x : Nat ← Gen.choose Nat 0 (← Gen.getSize) (Nat.zero_le _) - if x % 10 == 0 + if x % n == 0 then return x else throw <| .genError "uh oh" --- This fails 9/10 times --- #eval Gen.run test 9 +-- -- This fails `9/10` times +-- #eval Gen.run (test 10) 9 + +-- -- This succeeds almost always. +-- #eval Gen.runUntil (attempts := .some 1000) (test 10) 9 + --- This succeeds almost always. --- #eval Gen.runUntil (attempts := .some 1000) test 9 +-- -- This test usually returns an even number, but sometimes `1` +-- #eval Gen.run (Gen.backtrack [(10, test 2), (1, return 1)]) 4 end Plausible