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