From 9d89af68375f2f874ab081e1d4c3c24d2eb37592 Mon Sep 17 00:00:00 2001 From: David Ledvinka Date: Fri, 13 Nov 2026 19:03:36 +0000 Subject: [PATCH 1/2] Split rdo notation into Core and ForIn --- RandomDo.lean | 4 +- RandomDo/Monad/ForInInstances.lean | 2 +- RandomDo/RDo.lean | 4 + RandomDo/RDo/Core.lean | 82 +++++++++++++++++ .../{Monad/Notation.lean => RDo/ForIn.lean} | 57 +----------- RandomDo/Tactic/Computable/Example.lean | 2 +- Test.lean | 1 + Test/RDo/Core.lean | 89 +++++++++++++++++++ 8 files changed, 184 insertions(+), 57 deletions(-) create mode 100644 RandomDo/RDo.lean create mode 100644 RandomDo/RDo/Core.lean rename RandomDo/{Monad/Notation.lean => RDo/ForIn.lean} (80%) create mode 100644 Test/RDo/Core.lean diff --git a/RandomDo.lean b/RandomDo.lean index 4c24e0c..44782b1 100644 --- a/RandomDo.lean +++ b/RandomDo.lean @@ -6,13 +6,15 @@ public import RandomDo.MeasurableSpaceMonad.Defs public import RandomDo.MeasurableSpaceMonad.ForIn public import RandomDo.Monad.ForInInstances public import RandomDo.Monad.Instances -public import RandomDo.Monad.Notation public import RandomDo.NumLean.Binomial public import RandomDo.NumLean.Distributions public import RandomDo.NumLean.PCG64 public import RandomDo.NumLean.SeedSequence public import RandomDo.NumLean.Ziggurat public import RandomDo.NumLean.ZigguratSampler +public import RandomDo.RDo +public import RandomDo.RDo.Core +public import RandomDo.RDo.ForIn public import RandomDo.Tactic.Computable.Counterparts public import RandomDo.Tactic.Computable.Defs public import RandomDo.Tactic.Computable.Deriving diff --git a/RandomDo/Monad/ForInInstances.lean b/RandomDo/Monad/ForInInstances.lean index fde7c94..deb32fa 100644 --- a/RandomDo/Monad/ForInInstances.lean +++ b/RandomDo/Monad/ForInInstances.lean @@ -5,7 +5,7 @@ Authors: David Ledvinka -/ module -public import RandomDo.Monad.Notation +public import RandomDo.RDo.ForIn /-! # Instances for `MeasurableSpaceForIn` diff --git a/RandomDo/RDo.lean b/RandomDo/RDo.lean new file mode 100644 index 0000000..f92f846 --- /dev/null +++ b/RandomDo/RDo.lean @@ -0,0 +1,4 @@ +module -- shake: keep-all --deprecated_module: ignore + +public import RandomDo.RDo.Core +public import RandomDo.RDo.ForIn diff --git a/RandomDo/RDo/Core.lean b/RandomDo/RDo/Core.lean new file mode 100644 index 0000000..d49eb54 --- /dev/null +++ b/RandomDo/RDo/Core.lean @@ -0,0 +1,82 @@ +/- +Copyright (c) 2026 David Ledvinka. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: David Ledvinka +-/ +module + +public import RandomDo.MeasurableSpaceMonad.Defs +public import Mathlib.MeasureTheory.MeasurableSpace.Prod +public meta import Lean.Elab.BuiltinDo + +/-! +# `rdo` notation + +This file defines `rdo` notation for computations in a measurable space monad, where +`MeasurableSpacePure.mPure` and `MeasurableSpaceBind.mBind` take the roles of `Pure.pure` and +`Bind.bind` respectively. Use `open MeasureTheory` (or `open scoped MeasureTheory`) to enable the +notation. +-/ + +open Lean Lean.Parser Lean.Parser.Term Lean.Meta Lean.Elab Lean.Elab.Do Lean.Elab.Term Std Meta +open MeasureTheory MeasurableSpacePure MeasurableSpaceBind + +public meta section + +namespace MeasureTheory + +/-- `rdo` is the `do` notation for writing monadic programs in a measurable space monad. -/ +scoped syntax (name := randKind) "rdo" doSeq : term + +namespace RDo + +/-- The `DoOps` for `rdo` notation. Analogue of `do` notation's `Lean.Elab.Do.DoOps.default`. -/ +def randOps : DoOps where + mkPureApp α e := do + let info := (← read).monadInfo + if (← read).deadCode matches .deadSyntactically then + -- There is no dead syntax here. Just return a fresh metavariable so that we don't + -- do the `Term.ensureHasType` check below. + return ← mkFreshExprMVar (← mkMonadApp α) + let α ← Term.ensureHasType (mkSort (mkLevelSucc info.u)) α + let e ← Term.ensureHasType α e + let instPure ← Term.mkInstMVar (mkApp (mkConst ``MeasurableSpacePure [info.u, info.v]) info.m) + let instPure ← instantiateMVars instPure + let σ ← Term.mkInstMVar (mkApp (mkConst ``MeasurableSpace [info.u]) α) + let σ ← instantiateMVars σ + return mkApp5 (mkConst ``MeasurableSpacePure.mPure [info.u, info.v]) info.m instPure α σ e + mkBindApp α β e k := do + let info := (← read).monadInfo + let α ← Term.ensureHasType (mkSort (mkLevelSucc info.u)) α + let σα ← Term.mkInstMVar (mkApp (mkConst ``MeasurableSpace [info.u]) α) + let σβ ← Term.mkInstMVar (mkApp (mkConst ``MeasurableSpace [info.u]) β) + let mα := mkApp2 info.m α σα + let e ← Term.ensureHasType mα e + let k ← Term.ensureHasType (← mkArrow α (mkApp2 info.m β σβ)) k + let instBind ← Term.mkInstMVar (mkApp (mkConst ``MeasurableSpaceBind [info.u, info.v]) info.m) + return mkApp8 (mkConst ``MeasurableSpaceBind.mBind [info.u, info.v]) + info.m instBind α β σα σβ e k + isPureApp? e := + if e.isAppOfArity ``MeasurableSpacePure.mPure 5 then some (e.getArg! 4) else none + splitMonadApp? type := do + let .app m σ := type.consumeMData | return none + let .app m resultType := m.consumeMData | return none + unless ← isType resultType do return none + let σType ← whnf (← inferType σ) + unless σType.isAppOf ``MeasurableSpace do return none + let u ← getDecLevel resultType + let v ← getDecLevel type + return some ({ m, u := u.normalize, v := v.normalize }, resultType) + mkMonadApp α := do + let info := (← read).monadInfo + let σ ← Term.mkInstMVar (mkApp (mkConst ``MeasurableSpace [info.u]) α) + return mkApp2 info.m α σ + +/-- The `rdo` notation elaborator. Analogue of `do` notation's `Lean.Elab.Do.elabDo`. -/ +@[term_elab randKind] def elabRand : Term.TermElab := fun stx et? => do + let `(rdo $doSeq) := stx | throwUnsupportedSyntax + elabDoWith randOps doSeq et? + +end RDo + +end MeasureTheory diff --git a/RandomDo/Monad/Notation.lean b/RandomDo/RDo/ForIn.lean similarity index 80% rename from RandomDo/Monad/Notation.lean rename to RandomDo/RDo/ForIn.lean index 9b1d5f9..c1e3b4b 100644 --- a/RandomDo/Monad/Notation.lean +++ b/RandomDo/RDo/ForIn.lean @@ -5,74 +5,23 @@ Authors: David Ledvinka -/ module +public import RandomDo.RDo.Core public import RandomDo.MeasurableSpaceMonad.ForIn -public import Mathlib.MeasureTheory.MeasurableSpace.Prod -public meta import Lean.Elab.BuiltinDo public meta import Lean.Meta.ProdN /-! -# `rdo` Notation - -**TODO** +# `for` loops in `rdo` notation +This file defines the syntax, macro expansion, and elaboration for `for` loops in `rdo` notation. -/ open Lean Lean.Parser Lean.Parser.Term Lean.Meta Lean.Elab Lean.Elab.Do Lean.Elab.Term Std Meta - open MeasureTheory MeasurableSpacePure MeasurableSpaceBind public meta section namespace RDo -/-- Define `DoOp`s for `rdo` notation. -/ -def randOps : DoOps := { DoOps.default with - mkPureApp α e := do - let ⟨m, u, v, _, _⟩ := (← read).monadInfo - let α ← Term.ensureHasType (mkSort (mkLevelSucc u)) α - let e ← Term.ensureHasType α e - let instPure ← mkInstMVar (.app (mkConst ``MeasurableSpacePure [u,v]) m) - let instPure ← instantiateMVars instPure - let σ ← mkInstMVar (.app (mkConst ``MeasurableSpace [u]) α) - let σ ← instantiateMVars σ - return mkAppN (mkConst ``mPure [u, v]) #[m, instPure, α, σ, e] - mkBindApp α β e k := do - let ⟨m, u, v, _, _⟩ := (← read).monadInfo - let α ← Term.ensureHasType (mkSort (mkLevelSucc u)) α - let σα ← mkInstMVar (.app (mkConst ``MeasurableSpace [u]) α) - let σβ ← mkInstMVar (.app (mkConst ``MeasurableSpace [u]) β) - let mα := mkApp2 m α σα - let mβ := mkApp2 m β σβ - let e ← Term.ensureHasType mα e - let k ← Term.ensureHasType (← mkArrow α mβ) k - let σα ← instantiateMVars σα - let σβ ← instantiateMVars σβ - let instBind ← mkInstMVar (.app (mkConst ``MeasurableSpaceBind [u, v]) m) - let instBind ← instantiateMVars instBind - return mkAppN (mkConst ``mBind [u, v]) #[m, instBind, α, β, σα, σβ, e, k] - isPureApp? e := - if e.isAppOfArity ``mPure 5 then some (e.getArg! 4) else none - splitMonadApp? type := do - let .app m _ := type.consumeMData | return none - let .app m resultType := m.consumeMData | return none - unless ← isType resultType do return none - let u ← getDecLevel resultType - let v ← getDecLevel type - return some ({ m := m, u := u.normalize, v := v.normalize }, resultType) - mkMonadApp α := do - let ⟨m, u, _, _, _⟩ := (← read).monadInfo - let σ ← mkInstMVar (mkApp (mkConst ``MeasurableSpace [u]) α) - return mkApp2 m α σ - } - -/-- The `do` notation for writing monadic programs depending on a `MeasurableSpace` instance. -/ -syntax (name := randKind) "rdo" doSeq : term - -/-- Define `rdo` notation elaborator. -/ -@[term_elab randKind] def elabRand : Term.TermElab := fun stx et? => do - let `(rdo $doSeq) := stx | throwUnsupportedSyntax - elabDoWith randOps doSeq et? - section LoopElab /-- parses `_ : _ in` for `rdo` for loops -/ diff --git a/RandomDo/Tactic/Computable/Example.lean b/RandomDo/Tactic/Computable/Example.lean index 13f7e54..0c725c0 100644 --- a/RandomDo/Tactic/Computable/Example.lean +++ b/RandomDo/Tactic/Computable/Example.lean @@ -7,7 +7,7 @@ Authors: Gaëtan Serré module public import RandomDo.Monad.Instances -public import RandomDo.Monad.Notation +public import RandomDo.RDo public import RandomDo.NumLean.Distributions public import RandomDo.Tactic.Computable.Deriving public import Mathlib diff --git a/Test.lean b/Test.lean index bba557b..f178602 100644 --- a/Test.lean +++ b/Test.lean @@ -10,4 +10,5 @@ public import Test.IsMarkov public import Test.Loops public import Test.MonadLaws public import Test.Polymorphic +public import Test.RDo.Core public import Test.Sample diff --git a/Test/RDo/Core.lean b/Test/RDo/Core.lean new file mode 100644 index 0000000..bc4d89d --- /dev/null +++ b/Test/RDo/Core.lean @@ -0,0 +1,89 @@ +/- +Copyright (c) 2026 David Ledvinka. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: David Ledvinka +-/ +module + +import RandomDo.RDo.Core + +/-! +# Translation and dead-code regressions for `rdo` + +The equalities below use an abstract monad without laws. They check the generated operations, +so a dropped action or an accidental use of ordinary `pure`/`bind` cannot pass by reduction in +an `Id` instance. Independent universes also exercise construction of computation types. +-/ + +set_option linter.privateModule false + +assert_not_exists MeasureTheory.MeasurableSpaceForIn RDo.rdoFor + +open MeasureTheory MeasurableSpacePure MeasurableSpaceBind + +universe u v + +variable {m : (α : Type u) → [MeasurableSpace α] → Type v} [MeasurableSpaceMonad m] + {α β : Type u} [MeasurableSpace α] [MeasurableSpace β] + +example (a : α) : (rdo return a : m α) = mPure a := rfl + +example (x : m α) (f : α → m β) : (rdo + let a ← x + f a : m β) = x >>=ₘ f := rfl + +-- An action whose result is discarded must still be sequenced, exactly once and before `next`. +example (action : m PUnit) (next : m α) : (rdo + action + next : m α) = action >>=ₘ fun _ => next := rfl + +-- Branches must preserve both mutable values and the order of their monadic updates. +example (a : α) (b : β) (choose : Bool) (left : α → m α) (right : β → m β) : + (rdo + let mut x := a + let mut y := b + if choose then + x ← left x + y ← right y + return (x, y) : m (α × β)) = + (if choose then + left a >>=ₘ fun x => right b >>=ₘ fun y => mPure (x, y) + else mPure (a, b)) := by + cases choose <;> rfl + +/-- +warning: This `do` element and its control-flow region are dead code. Consider removing it. +-/ +#guard_msgs in +example (a : α) : (rdo + return a + return true : m α) = mPure a := rfl + +-- This directly exercises the `deadSyntactically` branch of `mkPureApp`: warning traversal +-- must neither typecheck the discarded Bool at α nor request a missing monad instance. +open Lean Meta Elab Term Do in +omit [MeasurableSpaceMonad m] in +example (x : m α) : x = x := by_elab + let x ← getFVarFromUserName `x + let expectedType ← inferType x + let ctx ← mkContext expectedType MeasureTheory.RDo.randOps + let e ← MeasureTheory.RDo.randOps.mkPureApp ctx.doBlockResultType (mkConst ``Bool.true) |>.run + { ctx with deadCode := .deadSyntactically } + unless e.isMVar do + throwError "Expected a fresh metavariable for syntactically dead code" + unless ← isDefEq (← inferType e) expectedType do + throwError "Incorrect monad application for syntactically dead code" + mkEqRefl x + +-- An alias's arguments need not be the result type and its measurable-space instance. +private abbrev TaggedResult (_tag : Type) + (n : (α : Type) → [MeasurableSpace α] → Type v) := n Nat + +example {n : (α : Type) → [MeasurableSpace α] → Type v} [MeasurableSpaceMonad n] : + (rdo return 1 : TaggedResult Unit n) = (mPure 1 : n Nat) := rfl + +-- Recognizing the instance argument must also allow aliases for its type. +private abbrev SpaceAlias (α : Type u) := MeasurableSpace α + +example {γ : Type u} [SpaceAlias γ] (a : γ) : + (rdo return a : m γ) = mPure a := rfl From d1e759dcf7d2c1c28dd19c2d59c69aadc7af0b2b Mon Sep 17 00:00:00 2001 From: David Ledvinka Date: Fri, 13 Nov 2026 19:10:21 +0000 Subject: [PATCH 2/2] Keep notation tests on the scaffold branch --- Test.lean | 1 - Test/RDo/Core.lean | 89 ---------------------------------------------- 2 files changed, 90 deletions(-) delete mode 100644 Test/RDo/Core.lean diff --git a/Test.lean b/Test.lean index f178602..bba557b 100644 --- a/Test.lean +++ b/Test.lean @@ -10,5 +10,4 @@ public import Test.IsMarkov public import Test.Loops public import Test.MonadLaws public import Test.Polymorphic -public import Test.RDo.Core public import Test.Sample diff --git a/Test/RDo/Core.lean b/Test/RDo/Core.lean deleted file mode 100644 index bc4d89d..0000000 --- a/Test/RDo/Core.lean +++ /dev/null @@ -1,89 +0,0 @@ -/- -Copyright (c) 2026 David Ledvinka. All rights reserved. -Released under Apache 2.0 license as described in the file LICENSE. -Authors: David Ledvinka --/ -module - -import RandomDo.RDo.Core - -/-! -# Translation and dead-code regressions for `rdo` - -The equalities below use an abstract monad without laws. They check the generated operations, -so a dropped action or an accidental use of ordinary `pure`/`bind` cannot pass by reduction in -an `Id` instance. Independent universes also exercise construction of computation types. --/ - -set_option linter.privateModule false - -assert_not_exists MeasureTheory.MeasurableSpaceForIn RDo.rdoFor - -open MeasureTheory MeasurableSpacePure MeasurableSpaceBind - -universe u v - -variable {m : (α : Type u) → [MeasurableSpace α] → Type v} [MeasurableSpaceMonad m] - {α β : Type u} [MeasurableSpace α] [MeasurableSpace β] - -example (a : α) : (rdo return a : m α) = mPure a := rfl - -example (x : m α) (f : α → m β) : (rdo - let a ← x - f a : m β) = x >>=ₘ f := rfl - --- An action whose result is discarded must still be sequenced, exactly once and before `next`. -example (action : m PUnit) (next : m α) : (rdo - action - next : m α) = action >>=ₘ fun _ => next := rfl - --- Branches must preserve both mutable values and the order of their monadic updates. -example (a : α) (b : β) (choose : Bool) (left : α → m α) (right : β → m β) : - (rdo - let mut x := a - let mut y := b - if choose then - x ← left x - y ← right y - return (x, y) : m (α × β)) = - (if choose then - left a >>=ₘ fun x => right b >>=ₘ fun y => mPure (x, y) - else mPure (a, b)) := by - cases choose <;> rfl - -/-- -warning: This `do` element and its control-flow region are dead code. Consider removing it. --/ -#guard_msgs in -example (a : α) : (rdo - return a - return true : m α) = mPure a := rfl - --- This directly exercises the `deadSyntactically` branch of `mkPureApp`: warning traversal --- must neither typecheck the discarded Bool at α nor request a missing monad instance. -open Lean Meta Elab Term Do in -omit [MeasurableSpaceMonad m] in -example (x : m α) : x = x := by_elab - let x ← getFVarFromUserName `x - let expectedType ← inferType x - let ctx ← mkContext expectedType MeasureTheory.RDo.randOps - let e ← MeasureTheory.RDo.randOps.mkPureApp ctx.doBlockResultType (mkConst ``Bool.true) |>.run - { ctx with deadCode := .deadSyntactically } - unless e.isMVar do - throwError "Expected a fresh metavariable for syntactically dead code" - unless ← isDefEq (← inferType e) expectedType do - throwError "Incorrect monad application for syntactically dead code" - mkEqRefl x - --- An alias's arguments need not be the result type and its measurable-space instance. -private abbrev TaggedResult (_tag : Type) - (n : (α : Type) → [MeasurableSpace α] → Type v) := n Nat - -example {n : (α : Type) → [MeasurableSpace α] → Type v} [MeasurableSpaceMonad n] : - (rdo return 1 : TaggedResult Unit n) = (mPure 1 : n Nat) := rfl - --- Recognizing the instance argument must also allow aliases for its type. -private abbrev SpaceAlias (α : Type u) := MeasurableSpace α - -example {γ : Type u} [SpaceAlias γ] (a : γ) : - (rdo return a : m γ) = mPure a := rfl