Merge pull request #2 from eric-wieser/ulift

fix: allow sampling from higher universes

Plausible/Gen.lean

@@ -36,12 +36,14 @@ abbrev Gen (α : Type u) := ReaderT (ULift Nat) Rand α
 namespace Gen
 
 @[inline]
-def up (x : Gen.{u} α) : Gen.{max u v} (ULift α) := do
+def up (x : Gen.{u} α) : Gen (ULift.{v} α) := do
   let size ← read
-  let gen ← get
-  let ⟨val, gen⟩ := x.run ⟨size.down⟩ |>.run ⟨gen.down⟩
-  set <| ULift.up gen.down
-  return ⟨val⟩
+  Rand.up <| x.run ⟨size.down⟩
+
+@[inline]
+def down (x : Gen (ULift.{v} α)) : Gen α := do
+  let size ← read
+  Rand.down <| x.run ⟨size.down⟩
 
 /-- Lift `Random.random` to the `Gen` monad. -/
 def chooseAny (α : Type u) [Random Id α] : Gen α :=

Plausible/Random.lean

@@ -62,6 +62,27 @@ class BoundedRandom (m) (α : Type u) [LE α] where
   -/
   randomR {g : Type} (lo hi : α) (h : lo ≤ hi) [RandomGen g] : RandGT g m {a // lo ≤ a ∧ a ≤ hi}
 
+@[inline]
+protected def RandT.up {α : Type u} {m : Type u → Type w} {m' : Type (max u v) → Type w'}
+    {g : Type} [RandomGen g] [Monad m] [Monad m']
+    (m_up : ∀ {α}, m α → m' (ULift α)) (x : RandGT g m α) :
+    RandGT g m' (ULift.{v} α) := do
+  let ⟨val, gen⟩ ←  m_up <| x.run ⟨(← get).down⟩
+  set <| ULift.up gen.down
+  return ⟨val⟩
+
+@[inline]
+protected def RandT.down {α : Type u} {m : Type (max u v) → Type w} {m' : Type u → Type w'}
+    {g : Type} [RandomGen g] [Monad m] [Monad m']
+    (m_down : ∀ {α}, m (ULift α) → m' α) (x : RandGT g m (ULift.{v} α) ) :
+    RandGT g m' α := do
+  let gen := (← get).down
+  let ⟨val, gen⟩ ← m_down do
+    let ⟨⟨val⟩, ⟨gen⟩⟩ ← x.run ⟨gen⟩
+    pure <| .up (val, gen)
+  set <| ULift.up gen
+  return val
+
 namespace Rand
 /-- Generate one more `Nat` -/
 def next [RandomGen g] [Monad m] : RandGT g m Nat := do
@@ -81,6 +102,17 @@ def split {g : Type} [RandomGen g] [Monad m] : RandGT g m g := do
 def range {g : Type} [RandomGen g] [Monad m] : RandGT g m (Nat × Nat) := do
   let rng := (← get).down
   return RandomGen.range rng
+
+@[inline]
+protected def up {α : Type u} {g : Type} [RandomGen g] (x : RandG g α) :
+    RandG g (ULift.{v} α) := do
+  RandT.up (fun x => pure ⟨Id.run x⟩) x
+
+@[inline]
+protected def down {α : Type u} {g : Type} [RandomGen g] (x : RandG g (ULift.{v} α)) :
+    RandG g α :=
+  RandT.down (fun x => pure (Id.run x).down) x
+
 end Rand
 
 namespace Random

Plausible/Sampleable.lean

@@ -216,6 +216,9 @@ instance List.shrinkable [Shrinkable α] : Shrinkable (List α) where
     (L.mapIdx fun i _ => L.eraseIdx i) ++
     (L.mapIdx fun i a => (shrink a).map fun a' => L.modify (fun _ => a') i).join
 
+instance ULift.shrinkable [Shrinkable α] : Shrinkable (ULift α) where
+  shrink u := (shrink u.down).map ULift.up
+
 instance String.shrinkable : Shrinkable String where
   shrink s := (shrink s.toList).map String.mk
 
@@ -242,7 +245,7 @@ open SampleableExt
 instance Sum.SampleableExt [SampleableExt α] [SampleableExt β] : SampleableExt (Sum α β) where
   proxy := Sum (proxy α) (proxy β)
   sample := do
-    match ← chooseAny Bool with 
+    match ← chooseAny Bool with
     | true => return .inl (← sample)
     | false => return .inr (← sample)
   interp s :=
@@ -324,7 +327,7 @@ instance Char.sampleableDefault : SampleableExt Char :=
 instance Option.sampleableExt [SampleableExt α] : SampleableExt (Option α) where
   proxy := Option (proxy α)
   sample := do
-    match ← chooseAny Bool with 
+    match ← chooseAny Bool with
     | true => return none
     | false => return some (← sample)
   interp o := o.map interp
@@ -349,6 +352,11 @@ instance List.sampleableExt [SampleableExt α] : SampleableExt (List α) where
   sample := Gen.listOf sample
   interp := List.map interp
 
+instance ULift.sampleableExt [SampleableExt α] : SampleableExt (ULift α) where
+  proxy := proxy α
+  sample := sample
+  interp a := ⟨interp a⟩
+
 instance String.sampleableExt : SampleableExt String :=
   mkSelfContained do return String.mk (← Gen.listOf (Char.sampleableDefault.sample))
 
@@ -381,9 +389,18 @@ end NoShrink
 /--
 Print (at most) 10 samples of a given type to stdout for debugging.
 -/
-def printSamples {t : Type} [Repr t] (g : Gen t) : IO PUnit := do
-  for i in List.range 10 do
-    IO.println s!"{repr (← g.run i)}"
+def printSamples {t : Type u} [Repr t] (g : Gen t) : IO PUnit := do
+-- TODO: this should be a global instance
+  letI : MonadLift Id IO := ⟨fun f => pure <| Id.run f⟩
+  do
+    -- we can't convert directly from `Rand (List t)` to `RandT IO (List Std.Format)`
+    -- (and `RandT IO (List t)` isn't type-correct without
+    -- https://github.com/leanprover/lean4/issues/3011), so go via an intermediate
+    let xs : List Std.Format ← Plausible.runRand <| Rand.down <| do
+      let xs : List t ← (List.range 10).mapM (ReaderT.run g ∘ ULift.up)
+      pure <| ULift.up (xs.map repr)
+    for x in xs do
+      IO.println s!"{x}"
 
 open Lean Meta Elab
 
@@ -449,9 +466,8 @@ values of type `type` using an increasing size parameter.
 elab "#sample " e:term : command =>
   Command.runTermElabM fun _ => do
     let e ← Elab.Term.elabTermAndSynthesize e none
-    let g ← mkGenerator e
-    let ⟨0, α, repr, gen⟩ := g | throwError "Cannot sample from {g.1} due to its universe"
-    let printSamples := mkApp3 (mkConst ``printSamples) α repr gen
+    let ⟨u, α, repr, gen⟩ ← mkGenerator e
+    let printSamples := mkApp3 (mkConst ``printSamples [u]) α repr gen
     let code ← unsafe evalExpr (IO PUnit) (mkApp (mkConst ``IO) (mkConst ``PUnit [1])) printSamples
     _ ← code
 

Plausible/Testable.lean

@@ -312,6 +312,14 @@ instance forallTypesTestable {f : Type → Prop} [Testable (f Int)] :
     let r ← runProp (f Int) cfg min
     return addVarInfo var "Int" (· <| Int) r
 
+-- TODO: only in mathlib: @[pp_with_univ]
+instance (priority := 100) forallTypesULiftTestable.{u}
+    {f : Type u → Prop} [Testable (f (ULift.{u} Int))] :
+    Testable (NamedBinder var <| ∀ x, f x) where
+  run cfg min := do
+    let r ← runProp (f (ULift Int)) cfg min
+    pure <| addVarInfo var "ULift Int" (· <| ULift Int) r
+
 /--
 Format the counter-examples found in a test failure.
 -/

Test/Tactic.lean

@@ -149,3 +149,30 @@ theorem testBit_pred :
     testBit (pred x) i = (decide (0 < x) &&
       (Bool.xor ((List.range i).all fun j => ! testBit x j) (testBit x i))) := by
   plausible
+
+/--
+error:
+===================
+Found a counter-example!
+f := 1
+issue: ULift.up 1 = ULift.up 0 does not hold
+(0 shrinks)
+-------------------
+-/
+#guard_msgs in
+theorem ulift_nat (f : ULift.{1} Nat) : f = ⟨0⟩ := by
+  plausible
+
+/--
+error:
+===================
+Found a counter-example!
+α := "ULift Int"
+l := [0]
+issue: [ULift.up 0] = [ULift.up 0, ULift.up 0] does not hold
+(1 shrinks)
+-------------------
+-/
+#guard_msgs in
+theorem type_u (α : Type u) (l : List α) : l = l ++ l := by
+  plausible