Loops

RandomDo.lean

@@ -1,11 +1,14 @@
 module  -- shake: keep-all --deprecated_module: ignore
 
 public import RandomDo.ForMathlib.MeasureTheory.MeasurableSpace.Embedding
+public import RandomDo.ForMathlib.MeasureTheory.Measure.GiryMonad
+public import RandomDo.ForMathlib.Probability.Distributions.Bernoulli
 public import RandomDo.Measurable
 public import RandomDo.Monad.ForInInstances
 public import RandomDo.Monad.Instances
 public import RandomDo.Monad.MeasurableSpace
 public import RandomDo.Monad.Notation
+public import RandomDo.Monad.While
 public import RandomDo.NumLean.Binomial
 public import RandomDo.NumLean.Distributions
 public import RandomDo.NumLean.PCG64
@@ -23,3 +26,6 @@ public import RandomDo.Tactic.IsMarkov.Deriving
 public import RandomDo.Tactic.IsMarkov.Elab
 public import RandomDo.Tactic.IsMarkov.ForInStep
 public import RandomDo.Tactic.IsMarkov.Lemmas
+public import RandomDo.Tactic.IsMarkov.While.LoopInvariant
+public import RandomDo.Tactic.IsMarkov.While.Tactic
+public import RandomDo.Tactic.IsMarkov.While.Termination

RandomDo/ForMathlib/MeasureTheory/Measure/GiryMonad.lean

@@ -0,0 +1,39 @@
+/-
+Copyright (c) 2026 Gaëtan Serré. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Gaëtan Serré
+-/
+module
+
+public import Mathlib.MeasureTheory.Measure.GiryMonad
+
+/-!
+# The bind of a sum of two measures, and of a Dirac mass
+
+-/
+
+@[expose] public section
+
+open MeasureTheory
+
+namespace MeasureTheory.Measure
+
+variable {α β : Type*} [MeasurableSpace α] [MeasurableSpace β]
+
+theorem bind_add {μ ν : Measure α} {f : α → Measure β} (hf : AEMeasurable f (μ + ν)) :
+    (μ + ν).bind f = μ.bind f + ν.bind f := by
+  obtain ⟨hμ, hν⟩ := aemeasurable_add_measure_iff.1 hf
+  ext s hs
+  rw [add_apply, bind_apply hs hf, bind_apply hs hμ, bind_apply hs hν, lintegral_add_measure]
+
+/-- Binding a Dirac mass at a point of a space whose points are measurable: the continuation needs
+no measurability. -/
+@[simp]
+theorem dirac_bind' [MeasurableSingletonClass α] (a : α) (f : α → Measure β) :
+    (dirac a).bind f = f a := by
+  rw [Measure.bind, map_congr (ae_eq_dirac f)]
+  change (map (fun _ ↦ f a) (dirac a)).join = f a
+  rw [map_const]
+  simp
+
+end MeasureTheory.Measure

RandomDo/ForMathlib/Probability/Distributions/Bernoulli.lean

@@ -0,0 +1,39 @@
+/-
+Copyright (c) 2026 Gaëtan Serré. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Gaëtan Serré
+-/
+module
+
+public import Mathlib.Probability.Distributions.Bernoulli
+public import RandomDo.ForMathlib.MeasureTheory.Measure.GiryMonad
+
+/-!
+# Binding a Bernoulli distribution
+
+-/
+
+@[expose] public section
+
+open MeasureTheory Measure unitInterval
+open scoped ENNReal
+
+namespace ProbabilityTheory
+
+variable {X Y : Type*} [MeasurableSpace X] [MeasurableSpace Y]
+
+/-- Binding a Bernoulli distribution on a space whose points are measurable: the continuation needs
+no measurability, and the two weights are real numbers, so that `simp` can use it and `norm_num`
+can compute with the result. -/
+@[simp]
+lemma bernoulliMeasure_bind [MeasurableSingletonClass X] (x y : X) (p : I) (g : X → Measure Y) :
+    Ber(x, y, p).bind g = ENNReal.ofReal p • g x + ENNReal.ofReal (1 - p) • g y := by
+  have h (q : I) : ((toNNReal q : NNReal) : ℝ≥0∞) = ENNReal.ofReal q := by
+    rw [ENNReal.ofReal, Real.toNNReal_of_nonneg q.2.1]
+    rfl
+  rw [bernoulliMeasure_def, bind_add ((aemeasurable_dirac.smul_measure _).add_measure
+    (aemeasurable_dirac.smul_measure _)), bind_smul, bind_smul, dirac_bind', dirac_bind']
+  change (toNNReal p : ℝ≥0∞) • g x + (toNNReal (σ p) : ℝ≥0∞) • g y = _
+  rw [h, h, coe_symm_eq]
+
+end ProbabilityTheory

RandomDo/Monad/Notation.lean

@@ -260,6 +260,22 @@ def rdoForDecl := leading_parser
           dec.continueWithUnit
   mkBindApp σ γ forIn rest
 
+/-- parser for `rdo` while loops -/
+@[doElem_parser] def rdoWhile := leading_parser
+  "while " >> withForbidden "rdo" doIfCond >> " rdo " >> doSeq
+
+/-- Define expander for `while` loops in `rdo` notation. As in core, `while c rdo body` is a loop
+over `Loop.mk` that runs `body` while `c` holds and breaks otherwise. -/
+@[macro rdoWhile] def expandRDoWhile : Macro
+  | `(rdoWhile| while%$tk $cond:doIfCond rdo $body) =>
+    `(doElem| for%$tk _ in Lean.Loop.mk rdo if $cond:doIfCond then $body else break)
+  | _ => Macro.throwUnsupported
+
+/-- Infer the `ControlInfo` of an `rdo` loop as that of the core `for` loop with the same body. -/
+@[doElem_control_info rdoFor] def controlInfoRDoFor : ControlInfoHandler := fun stx => do
+  let `(rdoFor| for $_:rdoForDecl,* rdo $body) := stx | throwUnsupportedSyntax
+  inferControlInfoElem (← `(doElem| for _ in #[()] do $body))
+
 end LoopElab
 
 end RDo