@@ -2,9 +2,10 @@ module  -- shake: keep-all --deprecated_module: ignore
 
 public import RandomDo.ForMathlib.MeasureTheory.MeasurableSpace.Embedding
 public import RandomDo.Measurable
+public import RandomDo.MeasurableSpaceMonad.Defs
+public import RandomDo.MeasurableSpaceMonad.ForIn
 public import RandomDo.Monad.ForInInstances
 public import RandomDo.Monad.Instances
-public import RandomDo.Monad.MeasurableSpace
 public import RandomDo.Monad.Notation
 public import RandomDo.NumLean.Binomial
 public import RandomDo.NumLean.Distributions
 
@@ -5,8 +5,8 @@ Authors: David Ledvinka
 -/
 module
 
-public import Mathlib.Control.Random
-public import Mathlib.MeasureTheory.Measure.GiryMonad
+public import Mathlib.MeasureTheory.MeasurableSpace.Constructions
+public import Mathlib.Tactic.FunProp
 
 /-!
 # Measurable space monads
@@ -26,12 +26,10 @@ instance.
 `MeasurableSpaceFunctor` satisfies the laws of a functor on measurable spaces.
 * `LawfulMeasurableSpaceMonad`: A `Prop`-valued typeclass which states that a given
 `MeasurableSpaceMonad` satisfies the laws of a monad on measurable spaces.
-* `MeasurableSpaceForIn`: **TODO**
-* `MeasurableSpaceForIn'`: **TODO**
 
 ## Notation
 
-The following notation is available:
+The following notation is available in the `MeasureTheory` scope:
 
 * `f <$>ₘ x`: `MeasurableSpaceFunctor.mMap f x`.
 * `x >>=ₘ f`: `MeasurableSpaceBind.mBind x f`.
@@ -55,9 +53,9 @@ see R. J. Aumann, [*Borel structures for function spaces*](https://doi.org/10.12
 
 @[expose] public section
 
-open Function
+namespace MeasureTheory
 
-section
+open Function
 
 universe u v
 
@@ -97,7 +95,7 @@ class MeasurableSpaceFunctor (f : (α : Type u) → [MeasurableSpace α] → Typ
   mMapConst {α β : Type u} [MeasurableSpace α] [MeasurableSpace β] : α → f β → f α :=
     mMap ∘ (const _)
 
-@[inherit_doc] infixr:100 " <$>ₘ " => MeasurableSpaceFunctor.mMap
+@[inherit_doc] scoped infixr:100 " <$>ₘ " => MeasurableSpaceFunctor.mMap
 
 /--
 The `mPure` function is overloaded via `MeasurableSpacePure` instances.
@@ -123,7 +121,7 @@ class MeasurableSpaceBind (m : (α : Type u) → [MeasurableSpace α] → Type v
   -/
   mBind {α β : Type u} [MeasurableSpace α] [MeasurableSpace β] : m α → (α → m β) → m β
 
-@[inherit_doc] infixl:55  " >>=ₘ " => MeasurableSpaceBind.mBind
+@[inherit_doc] scoped infixl:55  " >>=ₘ " => MeasurableSpaceBind.mBind
 
 /--
 A monad on types with a `MeasurableSpace` instance.
@@ -232,8 +230,8 @@ class LawfulMeasurableSpaceMonad
 open LawfulMeasurableSpaceMonad
 
 attribute [fun_prop] measurable_mPure measurable_mBind
-attribute [simp] pure_bind bind_assoc bind_pure_comp
-attribute [grind <=] pure_bind
+attribute [simp] mPure_mBind mBind_assoc mBind_mPure_comp
+attribute [grind <=] mPure_mBind
 
 variable {m : (α : Type u) → [MeasurableSpace α] → Type v} [MeasurableSpaceMonad m]
   [∀ α, [MeasurableSpace α] → MeasurableSpace (m α)] [LawfulMeasurableSpaceMonad m]
@@ -263,50 +261,8 @@ theorem mBind_mPure_unit {x : m PUnit} : (x >>=ₘ fun _ => mPure ⟨⟩) = x :=
     {g : β → m γ} (hg : Measurable g) :
     ((f <$>ₘ x) >>=ₘ fun b => g b) = (x >>=ₘ fun a => g (f a)) := by
   rw [← mBind_mPure_comp hf]
-  simp (disch := fun_prop) [mBind_assoc, mPure_mBind]
+  simp (disch := fun_prop) only [mBind_assoc, mPure_mBind]
 
 end Lawful
 
-end
-
-section MeasurableSpaceFor
-
-universe uρ uα u v
-
-variable {α : Type u} [mα : MeasurableSpace α] (m : (α : Type u) → [MeasurableSpace α] → Type v)
-  {m' : (α : Type u) → Type v}
-
-instance instMeasurableSpace {β : Type u} [mβ : MeasurableSpace β] :
-    MeasurableSpace (ForInStep β) := mβ.map ForInStep.yield ⊓ mβ.map ForInStep.done
-
-/--
-Monadic iteration in `rdo`-blocks, using the `for x in xs` notation.
--/
-class MeasurableSpaceForIn (ρ : Type uρ) (α : outParam (Type uα)) where
-  /--
-  Monadically iterates over the contents of a collection `xs`, with a local state `b` and the
-  possibility of early termination.
-  -/
-  forIn {β : Type u} [MeasurableSpace β] (xs : ρ) (b : β)
-    (f : α → β → m (ForInStep β)) : m β
-
-/--
-Monadic iteration in `rdo`-blocks with a membership proof, using the `for h : x in xs` notation.
--/
-class MeasurableSpaceForIn' (ρ : Type uρ) (α : outParam (Type uα))
-    (d : outParam (Membership α ρ)) where
-  /--
-  Monadically iterates over the contents of a collection `xs`, with a local state `b` and the
-  possibility of early termination. At each iteration, the body of the loop is provided with a proof
-  that the current element is in the collection.
-  -/
-  forIn' {β : Type u} [MeasurableSpace β] (xs : ρ) (b : β)
-    (f : (a : α) → a ∈ xs → β → m (ForInStep β)) : m β
-
-@[simps]
-instance (priority := 500) instMeasurableSpaceForInOfForIn'
-    {ρ : Type uρ} {α : Type uα} {d : Membership α ρ} [MeasurableSpaceForIn' m ρ α d] :
-    MeasurableSpaceForIn m ρ α where
-  forIn x b f := MeasurableSpaceForIn'.forIn' x b fun a _ s => f a s
-
-end MeasurableSpaceFor
+end MeasureTheory
@@ -0,0 +1,102 @@
+/-
+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
+
+/-!
+# Iteration in measurable space monads
+
+This file defines the typeclasses `MeasurableSpaceForIn` and `MeasurableSpaceForIn'`, the
+`MeasurableSpaceMonad` analogues of Lean core's `ForIn` and `ForIn'`, respectively.
+Instances of `MeasurableSpaceForIn` specify how to iterate over a collection of
+elements, running a monadic loop body at each step while carrying a local state and allowing
+early termination. These instances enable the use of `for x in xs rdo ...` in `rdo` notation.
+`MeasurableSpaceForIn'` provides the same operation, but also passes the loop body a proof that the
+current element belongs to the collection. It enables the use of `for h : x in xs rdo ...` in
+`rdo` notation (as well as `for x in xs rdo ...` since a `MeasurableSpaceForIn'` instance also
+provides a `MeasurableSpaceForIn` instance).
+
+-/
+
+@[expose] public section
+
+namespace MeasureTheory
+
+universe uρ uα u v
+
+instance instMeasurableSpaceForInStep {β : Type u} [mβ : MeasurableSpace β] :
+    MeasurableSpace (ForInStep β) :=
+  mβ.map ForInStep.yield ⊓ mβ.map ForInStep.done
+
+/--
+Monadic iteration in `rdo` blocks, using the `for x in xs rdo ...` notation.
+
+The parameter `m` is the monad of the `rdo` block in which iteration is performed, `ρ` is the type
+of the collection being iterated over, and `α` is the type of elements.
+
+This is the measurable space analogue of `ForIn`.
+-/
+class MeasurableSpaceForIn (m : (α : Type u) → [MeasurableSpace α] → Type v)
+    (ρ : Type uρ) (α : outParam (Type uα)) where
+  /--
+  Monadically iterates over the contents of a collection `xs`, with a local state `b` and the
+  possibility of early termination. This is the measurable space analogue of `ForIn.forIn`.
+  -/
+  forIn {β : Type u} [MeasurableSpace β] (xs : ρ) (b : β) (f : α → β → m (ForInStep β)) : m β
+
+export MeasurableSpaceForIn (forIn)
+
+/--
+Monadic iteration in `rdo` blocks with a membership proof, using the `for h : x in xs rdo ...`
+notation.
+
+The parameter `m` is the monad of the `rdo` block in which iteration is performed, `ρ` is the type
+of the collection being iterated over, `α` is the type of elements, and `d` is the specific
+membership predicate to provide.
+
+This is the measurable space analogue of `ForIn'`.
+-/
+class MeasurableSpaceForIn' (m : (α : Type u) → [MeasurableSpace α] → Type v)
+    (ρ : Type uρ) (α : outParam (Type uα)) (d : outParam (Membership α ρ)) where
+  /--
+  Monadically iterates over the contents of a collection `xs`, with a local state `b` and the
+  possibility of early termination. At each iteration, the body of the loop is provided with a
+  proof that the current element is in the collection. This is the measurable space analogue of
+  `ForIn'.forIn'`.
+  -/
+  forIn' {β : Type u} [MeasurableSpace β] (x : ρ) (b : β)
+    (f : (a : α) → a ∈ x → β → m (ForInStep β)) : m β
+
+export MeasurableSpaceForIn' (forIn')
+
+variable {m : (α : Type u) → [MeasurableSpace α] → Type v} {ρ : Type uρ} {α : Type uα}
+
+/--
+A `MeasurableSpaceForIn'` instance also provides a `MeasurableSpaceForIn` instance by ignoring the
+membership proof. We set the priority to 500 so it is below the default.
+-/
+@[default_instance]
+instance (priority := 500) instMeasurableSpaceForInOfForIn' {d : Membership α ρ}
+    [MeasurableSpaceForIn' m ρ α d] : MeasurableSpaceForIn m ρ α where
+  forIn x b f := forIn' x b fun a _ => f a
+
+@[simp] theorem forIn'_eq_forIn [d : Membership α ρ] [MeasurableSpaceForIn' m ρ α d]
+    {β : Type u} [MeasurableSpace β] (x : ρ) (b : β)
+    (f : (a : α) → a ∈ x → β → m (ForInStep β)) (g : (a : α) → β → m (ForInStep β))
+    (h : ∀ a h b, f a h b = g a b) :
+    forIn' x b f = forIn x b g := by
+  simp only [forIn]
+  congr
+  funext a ha b
+  exact h a ha b
+
+@[wf_preprocess] theorem forIn_eq_forIn' [d : Membership α ρ] [MeasurableSpaceForIn' m ρ α d]
+    {β : Type u} [MeasurableSpace β] (x : ρ) (b : β) (f : α → β → m (ForInStep β)) :
+    forIn x b f = forIn' x b (fun x h => binderNameHint x f <| binderNameHint h () <| f x) := by
+  rfl
+
+end MeasureTheory
@@ -48,7 +48,7 @@ end MeasurableSpace
 
 universe u v
 
-open MeasurableSpacePure
+open MeasureTheory MeasurableSpacePure
 
 variable {α : Type u} [mα : MeasurableSpace α] {m : (α : Type u) → [MeasurableSpace α] → Type v}
 
 
@@ -5,7 +5,9 @@ Authors: David Ledvinka
 -/
 module
 
-public import RandomDo.Monad.MeasurableSpace
+public import RandomDo.MeasurableSpaceMonad.Defs
+public import Mathlib.Control.Random
+public import Mathlib.MeasureTheory.Measure.GiryMonad
 public import Mathlib.Probability.ProductMeasure
 
 /-!
@@ -17,6 +19,8 @@ public import Mathlib.Probability.ProductMeasure
 
 @[expose] public section
 
+open MeasureTheory
+
 universe u v w
 
 /-- A (core) monad automatically defines a (not necessarily lawful) measurable space monad by
@@ -32,8 +36,6 @@ instance {m : Type u → Type v} [Monad m] :
 /-- A measurable space monad for pseudo random number generation. -/
 abbrev PseudoRandomM := Monad.toMeasurableSpaceMonad Rand
 
-open MeasureTheory
-
 open MeasurableSpacePure MeasurableSpaceBind MeasurableSpaceFunctor MeasurableSpaceMonad
 
 @[simps]
 
@@ -5,7 +5,9 @@ Authors: David Ledvinka
 -/
 module
 
-public import RandomDo.Monad.MeasurableSpace
+public import RandomDo.MeasurableSpaceMonad.ForIn
+public import Mathlib.MeasureTheory.MeasurableSpace.Prod
+public meta import Lean.Elab.BuiltinDo
 public meta import Lean.Meta.ProdN
 
 /-!
@@ -17,7 +19,7 @@ public meta import Lean.Meta.ProdN
 
 open Lean Lean.Parser Lean.Parser.Term Lean.Meta Lean.Elab Lean.Elab.Do Lean.Elab.Term Std Meta
 
-open MeasureTheory Measure MeasurableSpacePure MeasurableSpaceBind
+open MeasureTheory MeasurableSpacePure MeasurableSpaceBind
 
 public meta section
 
 
@@ -6,7 +6,7 @@ Authors: Gaëtan Serré
 module
 
 public meta import RandomDo.Tactic.Computable.Counterparts
-public import RandomDo.Monad.MeasurableSpace
+public import RandomDo.MeasurableSpaceMonad.ForIn
 public meta import Lean.Elab.Tactic.Basic
 
 /-!
@@ -44,7 +44,7 @@ module that writes it: the interpreter then asks for the code of what it calls,
 
 public meta section
 
-open Lean Meta NumLean
+open Lean Meta NumLean MeasureTheory
 
 namespace RDo.Tactic
 
 
@@ -6,7 +6,7 @@ Authors: Gaëtan Serré
 module
 
 public import RandomDo.Tactic.IsMarkov.Defs
-public import RandomDo.Monad.MeasurableSpace
+public import RandomDo.MeasurableSpaceMonad.ForIn
 
 /-!
 # Measurability of `ForInStep`