From 88a7081ce5ad4ca12752438bb30d1310f8ce26bd Mon Sep 17 00:00:00 2001 From: David Ledvinka Date: Wed, 4 Nov 2026 21:33:30 +0000 Subject: [PATCH] draft --- RandomDo.lean | 6 +- .../MeasurableSpace/Collections.lean | 53 ++ .../MeasureTheory/MeasurableSpace/Option.lean | 31 + RandomDo/Measurable.lean | 5 +- .../ForInInstances/Basic.lean | 153 ++++ .../ForInInstances/Iterators.lean | 347 ++++++++ .../ForInInstances/Ranges.lean | 822 ++++++++++++++++++ RandomDo/Monad/ForInInstances.lean | 129 --- RandomDo/Monad/Instances.lean | 6 + RandomDo/Monad/Notation.lean | 49 +- RandomDo/Tactic/IsMarkov/Lemmas.lean | 2 +- Test.lean | 2 + Test/ForInElab.lean | 137 +++ Test/ForInInstances.lean | 260 ++++++ Test/Gaps.lean | 24 +- 15 files changed, 1851 insertions(+), 175 deletions(-) create mode 100644 RandomDo/ForMathlib/MeasureTheory/MeasurableSpace/Collections.lean create mode 100644 RandomDo/ForMathlib/MeasureTheory/MeasurableSpace/Option.lean create mode 100644 RandomDo/MeasurableSpaceMonad/ForInInstances/Basic.lean create mode 100644 RandomDo/MeasurableSpaceMonad/ForInInstances/Iterators.lean create mode 100644 RandomDo/MeasurableSpaceMonad/ForInInstances/Ranges.lean delete mode 100644 RandomDo/Monad/ForInInstances.lean create mode 100644 Test/ForInElab.lean create mode 100644 Test/ForInInstances.lean diff --git a/RandomDo.lean b/RandomDo.lean index 4c24e0c..bf13ffb 100644 --- a/RandomDo.lean +++ b/RandomDo.lean @@ -1,10 +1,14 @@ module -- shake: keep-all --deprecated_module: ignore +public import RandomDo.ForMathlib.MeasureTheory.MeasurableSpace.Collections public import RandomDo.ForMathlib.MeasureTheory.MeasurableSpace.Embedding +public import RandomDo.ForMathlib.MeasureTheory.MeasurableSpace.Option public import RandomDo.Measurable public import RandomDo.MeasurableSpaceMonad.Defs public import RandomDo.MeasurableSpaceMonad.ForIn -public import RandomDo.Monad.ForInInstances +public import RandomDo.MeasurableSpaceMonad.ForInInstances.Basic +public import RandomDo.MeasurableSpaceMonad.ForInInstances.Iterators +public import RandomDo.MeasurableSpaceMonad.ForInInstances.Ranges public import RandomDo.Monad.Instances public import RandomDo.Monad.Notation public import RandomDo.NumLean.Binomial diff --git a/RandomDo/ForMathlib/MeasureTheory/MeasurableSpace/Collections.lean b/RandomDo/ForMathlib/MeasureTheory/MeasurableSpace/Collections.lean new file mode 100644 index 0000000..4663018 --- /dev/null +++ b/RandomDo/ForMathlib/MeasureTheory/MeasurableSpace/Collections.lean @@ -0,0 +1,53 @@ +/- +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 Mathlib.Data.List.OfFn +public import Mathlib.MeasureTheory.MeasurableSpace.Constructions + +/-! +# Measurable spaces on finite collections + +Lists carry the σ-algebra transported from `Σ n, Fin n → α` by `List.equivSigmaTuple`. +Arrays and subarrays inherit their measurable spaces through `toList`, and vectors through +`toArray`. In particular, the measurable space on a subarray depends only on its contents, +not on the unused portion of its backing array. + +The conversion functions `Array.toList`, `Subarray.toList`, and `Vector.toArray` are measurable. +These constructions are independent of the monad and iteration interfaces. +-/ + +@[expose] public section + +section MeasurableSpace + +variable {α : Type*} [MeasurableSpace α] + +instance instMeasurableSpaceList : MeasurableSpace (List α) := + MeasurableSpace.comap List.equivSigmaTuple inferInstance + +instance instMeasurableSpaceArray : MeasurableSpace (Array α) := + MeasurableSpace.comap Array.toList inferInstance + +instance instMeasurableSpaceVector {n : ℕ} : MeasurableSpace (Vector α n) := + MeasurableSpace.comap Vector.toArray inferInstance + +instance instMeasurableSpaceSubarray : MeasurableSpace (Subarray α) := + MeasurableSpace.comap (fun s : Subarray α ↦ s.toList) inferInstance + +@[fun_prop] +lemma measurable_toList : Measurable (Array.toList : Array α → List α) := + Measurable.of_comap_le fun _ a ↦ a + +@[fun_prop] +lemma measurable_subarray_toList : Measurable (fun s : Subarray α ↦ s.toList) := + Measurable.of_comap_le fun _ a ↦ a + +@[fun_prop] +lemma measurable_toArray {n : ℕ} : Measurable (Vector.toArray : Vector α n → Array α) := + Measurable.of_comap_le fun _ a ↦ a + +end MeasurableSpace diff --git a/RandomDo/ForMathlib/MeasureTheory/MeasurableSpace/Option.lean b/RandomDo/ForMathlib/MeasureTheory/MeasurableSpace/Option.lean new file mode 100644 index 0000000..564157b --- /dev/null +++ b/RandomDo/ForMathlib/MeasureTheory/MeasurableSpace/Option.lean @@ -0,0 +1,31 @@ +/- +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 Mathlib.MeasureTheory.MeasurableSpace.Basic + +/-! +# The measurable space on `Option` + +Equip `Option α` with the largest σ-algebra making `some` measurable. This is the disjoint-union +measurable structure, with `none` as the additional point. +-/ + +@[expose] public section + +namespace Option + +variable {α : Type*} [MeasurableSpace α] + +instance instMeasurableSpace : MeasurableSpace (Option α) := + MeasurableSpace.map some inferInstance + +@[simp] theorem measurableSet_iff {s : Set (Option α)} : + MeasurableSet s ↔ MeasurableSet (some ⁻¹' s) := Iff.rfl + +@[fun_prop] theorem measurable_some : Measurable (some : α → Option α) := fun _ hs => hs + +end Option diff --git a/RandomDo/Measurable.lean b/RandomDo/Measurable.lean index 2aab168..ac96dc2 100644 --- a/RandomDo/Measurable.lean +++ b/RandomDo/Measurable.lean @@ -5,7 +5,7 @@ Authors: Gaëtan Serré -/ module -public import RandomDo.Monad.ForInInstances +public import RandomDo.MeasurableSpaceMonad.ForInInstances.Basic public import RandomDo.ForMathlib.MeasureTheory.MeasurableSpace.Embedding public import Mathlib.MeasureTheory.MeasurableSpace.Prod public import Mathlib.Data.List.OfFn @@ -117,9 +117,6 @@ def Vector.measurableEquivTuple {n : ℕ} : Vector α n ≃ᵐ (Fin n → α) wh ext simp -instance instMeasurableSpaceOption : MeasurableSpace (Option α) := - MeasurableSpace.map some inferInstance - theorem measurableSet_option_iff {s : Set (Option α)} : MeasurableSet s ↔ MeasurableSet (some ⁻¹' s) := Iff.rfl diff --git a/RandomDo/MeasurableSpaceMonad/ForInInstances/Basic.lean b/RandomDo/MeasurableSpaceMonad/ForInInstances/Basic.lean new file mode 100644 index 0000000..6e09506 --- /dev/null +++ b/RandomDo/MeasurableSpaceMonad/ForInInstances/Basic.lean @@ -0,0 +1,153 @@ +/- +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.ForMathlib.MeasureTheory.MeasurableSpace.Collections +public import RandomDo.MeasurableSpaceMonad.ForIn +public import RandomDo.Monad.Notation +public import RandomDo.MeasurableSpaceMonad.ForInInstances.Ranges + +/-! +# Collection and iterator instances for measurable space monads + +This file provides `MeasurableSpaceForIn'` instances for `List`, `Array`, and `Vector` in any +`MeasurableSpaceMonad`, and `MeasurableSpaceForIn` instances for `Std.Iter` and `Std.IterM`, +including their `Total` and `Partial` wrappers. Together with `RandomDo.Monad.Notation`, these +instances support `for x in xs rdo ...`. The range instances additionally support direct range +syntax such as `for i in 0...n rdo ...`, including membership proofs. +The membership-aware instances additionally support +`for h : x in xs rdo ...`. + +Each consumer threads the loop state through the body. `ForInStep.yield` continues iteration +and `ForInStep.done` stops immediately. Only the loop state needs a `MeasurableSpace` instance; +the elements, collection, and internal iterator state do not. Lists, arrays, vectors, and pure +iterators allow the element universe to differ from the loop-state universe. + +## Main definitions + +* `List.measurableSpaceForIn'`, `Array.measurableSpaceForIn'`: iteration with a membership proof. + The vector instance delegates to the array implementation through `Vector.toArray`. +* `Iter.instMeasurableSpaceForIn'`, `IterM.instMeasurableSpaceForIn'`: opt-in membership-aware + consumers for iterators. As in core, the generic membership relation is an internal predicate, + so these definitions and their `Total`/`Partial` variants are not global instances. +* `Iter.forIn_eq_match_step`: unfolds a finite pure iterator loop by one step. + +## Implementation details + +The list, array, and vector implementations follow Lean core's `List.forIn'`, `Array.forIn'`, +and vector `ForIn'` instance in `Init/Data/List/Control.lean`, `Init/Data/Array/Basic.lean`, and +`Init/Data/Vector/Basic.lean`. They replace `Monad`, `do`, and `pure` with +`MeasurableSpaceMonad`, `rdo`, and `mPure`, and take a measurable space on the loop state. +As in core, lists use borrowing and a specialized recursive loop. Arrays have a structurally +recursive reference definition and a compiler implementation using `USize` indices, connected +by `implemented_by`. + +Iterator consumers are defined in `ForInInstances.Iterators` and `ForInInstances.Ranges`. +They adapt core's generic and specialized range loops to use measurable bind directly. The same +consumer serves `Measure` and adapted ordinary monads; producer effects and loop-body effects remain +interleaved. The `Total` wrappers retain core's `Finite` requirement. Ordinary consumers can also +traverse potentially infinite iterators. + +-/ + +@[expose] public section + +namespace MeasureTheory + +universe uρ uα u v w + +open MeasurableSpacePure MeasurableSpaceBind + +section Collections + +variable {α : Type uα} {m : (α : Type u) → [MeasurableSpace α] → Type v} + [MeasurableSpaceMonad m] {β : Type u} [MeasurableSpace β] + +section Array + +open _root_.Array + +/-- Compiler implementation of `Array.measurableSpaceForIn'` using `USize` indices. +Analogue of core's `Array.forIn'Unsafe`. As in core, this relies on the runtime bound on array +sizes to represent each valid index by a `USize`. -/ +@[inline] unsafe def Array.measurableSpaceForIn'Unsafe + (as : Array α) (b : β) (f : (a : α) → a ∈ as → β → m (ForInStep β)) : m β := + let sz := as.usize + let rec @[specialize, nolint docBlame] loop (i : USize) (b : β) : m β := rdo + if i < sz then + let a := as.uget i lcProof + match (← f a lcProof b) with + | ForInStep.done b => mPure b + | ForInStep.yield b => loop (i+1) b + else + mPure b + loop 0 b + +/-- Reference implementation of iteration over an array, with a membership proof. +Analogue of core's `Array.forIn'`. -/ +@[implemented_by Array.measurableSpaceForIn'Unsafe] +protected def Array.measurableSpaceForIn' + (as : Array α) (b : β) (f : (a : α) → a ∈ as → β → m (ForInStep β)) : m β := + let rec @[nolint docBlame] loop (i : Nat) (h : i ≤ as.size) (b : β) : m β := rdo + match i, h with + | 0, _ => mPure b + | i+1, h => + have h' : i < as.size := Nat.lt_of_lt_of_le (Nat.lt_succ_self i) h + have : as.size - 1 < as.size := Nat.sub_lt (Nat.zero_lt_of_lt h') (by decide) + have : as.size - 1 - i < as.size := Nat.lt_of_le_of_lt (Nat.sub_le (as.size - 1) i) this + match (← f as[as.size - 1 - i] (getElem_mem this) b) with + | ForInStep.done b => mPure b + | ForInStep.yield b => loop i (Nat.le_of_lt h') b + loop as.size (Nat.le_refl _) b + +instance : MeasurableSpaceForIn' m (Array α) α inferInstance where + forIn' := Array.measurableSpaceForIn' + +@[simp] theorem Array.measurableSpaceForIn'_eq_forIn' : + Array.measurableSpaceForIn' (m := m) (α := α) (β := β) = forIn' := rfl + +instance {n : ℕ} : MeasurableSpaceForIn' m (Vector α n) α inferInstance where + forIn' xs b f := Array.measurableSpaceForIn' xs.toArray b (fun a h b => f a (by simpa using h) b) + +end Array + +section List + +open _root_.List + +/-- Iteration over a list, with a membership proof. Analogue of core's `List.forIn'`. -/ +@[inline] +protected def List.measurableSpaceForIn' (as : @& List α) (init : β) + (f : (a : α) → a ∈ as → β → m (ForInStep β)) : m β := + let rec @[specialize, nolint docBlame] + loop : (as' : @& List α) → (b : β) → Exists (fun bs => bs ++ as' = as) → m β + | [], b, _ => mPure b + | a::as', b, h => rdo + have : a ∈ as := by + clear f + have ⟨bs, h⟩ := h + subst h + exact mem_append_right _ (Mem.head ..) + match (← f a this b) with + | ForInStep.done b => mPure b + | ForInStep.yield b => + have : Exists (fun bs => bs ++ as' = as) := + have ⟨bs, h⟩ := h + ⟨bs ++ [a], by rw [← h, append_cons (bs := as')]⟩ + loop as' b this + loop as init ⟨[], rfl⟩ + +instance : MeasurableSpaceForIn' m (List α) α inferInstance where + forIn' := List.measurableSpaceForIn' + +@[simp, grind =] theorem List.measurableSpaceForIn'_eq_forIn' : + List.measurableSpaceForIn' (m := m) (α := α) (β := β) = forIn' := rfl + +end List + +end Collections + +end MeasureTheory diff --git a/RandomDo/MeasurableSpaceMonad/ForInInstances/Iterators.lean b/RandomDo/MeasurableSpaceMonad/ForInInstances/Iterators.lean new file mode 100644 index 0000000..fcbf1e2 --- /dev/null +++ b/RandomDo/MeasurableSpaceMonad/ForInInstances/Iterators.lean @@ -0,0 +1,347 @@ +/- +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.ForIn +public import RandomDo.Monad.Notation +public import Init.Data.Iterators.Consumers.Loop +import Init.Data.Iterators.Lemmas.Consumers.Loop +import all Init.Data.Iterators.Basic + +/-! +# Iterator loops in measurable space monads + +This file adapts core's iterator consumers to `mPure` and `mBind`. A single implementation works +in every `MeasurableSpaceMonad`, including `Measure` and adapted ordinary monads. The loop state +needs a measurable space; the elements and iterator internals do not. + +## Main definitions + +* `MeasurableSpaceIteratorLoop`: an extension point for producer-specific loop implementations. +* `LawfulMeasurableSpaceIteratorLoop`: agreement with the default consumer on well-founded loops. +* `MeasurableSpaceLiftBind`: lifting of ordinary monadic producer effects into a measurable monad. +* `Iter.instMeasurableSpaceForIn'`, `IterM.instMeasurableSpaceForIn'`: opt-in membership-aware + consumers. Ordinary `for` loops and `Total`/`Partial` wrappers have global instances. +* `Iter.forIn_eq_match_step`: the one-step equation for finite pure iterators. + +## Implementation details + +The definitions follow `Init/Data/Iterators/Consumers/Monadic/Loop.lean` and +`Init/Data/Iterators/Consumers/Loop.lean`. The default consumer uses the same +`WellFounded.extrinsicFix₃` recursion and specialization as core. Loop-body effects run directly +in the measurable monad. There is no continuation-monad adapter or final continuation. + +The low-priority default supports any iterator. Producer-specific optimizations require their +own `MeasurableSpaceIteratorLoop` instance; an ordinary `Std.IteratorLoop` instance does not +provide one automatically. A function abstract over the iterator type should take the loop +instance as a parameter if it needs to preserve a caller's specialized consumer, as in core. +`ForInInstances.Ranges` supplies the range specializations. + +`IterM` producers use an ordinary monad. `MeasurableSpaceLiftBind` runs a producer action and then +passes its internal result to the consumer, without assigning a measurable space to that result. +For adapted ordinary monads such as `PseudoRandomM`, this uses the existing `MonadLiftT` and bind. +Pure iterators apply the function directly and require no lifting instance from the user. +-/ + +@[expose] public section + +namespace MeasureTheory + +open Std Std.Iterators MeasurableSpacePure MeasurableSpaceBind + +universe u v w w' + +/-- Lift an ordinary computation and pass its result to a measurable computation. The source +result type needs no measurable space. This is the `liftBind` argument of core's iterator consumers, +with a measurable target type. -/ +class MeasurableSpaceLiftBind (n : Type w → Type w') + (m : (α : Type u) → [MeasurableSpace α] → Type v) where + /-- Run the producer action, then apply the function to its result. -/ + liftBind : {α : Type w} → {β : Type u} → [MeasurableSpace β] → + (α → m β) → n α → m β + +instance instMeasurableSpaceLiftBindId {m : (α : Type u) → [MeasurableSpace α] → Type v} : + MeasurableSpaceLiftBind Id m where + liftBind f x := f x.run + +/-- Loop consumers for an iterator and a measurable space monad. Analogue of +`Std.IteratorLoop`, specialized to bodies returning `m (ForInStep β)`. + +Core's extra predicate on loop-body steps is fixed to `True`, as in its `ForIn` instances. +Keeping the body result as `ForInStep β` avoids inserting a monadic bind and pure solely to attach +a proof of `True`. Such an insertion would need measurability hypotheses to simplify in general. -/ +class MeasurableSpaceIteratorLoop (ι : Type w) (n : Type w → Type w') {α : Type w} + [Iterator ι n α] (m : (β : Type u) → [MeasurableSpace β] → Type v) where + /-- Traverse an iterator, lifting producer effects and threading a measurable loop state. -/ + forIn : (lift : ∀ (γ : Type w) (δ : Type u) [MeasurableSpace δ], (γ → m δ) → n γ → m δ) → + (β : Type u) → [MeasurableSpace β] → (it : Std.IterM (α := ι) n α) → β → + ((a : α) → it.IsPlausibleIndirectOutput a → β → m (ForInStep β)) → m β + +variable {n : Type w → Type w'} {ι α : Type w} [Iterator ι n α] + {m : (β : Type u) → [MeasurableSpace β] → Type v} [MeasurableSpaceMonad m] + +/-- Direct iteration through `IterM.step`. Analogue of `Std.IterM.DefaultConsumers.forIn'`. +The predicate `P` carries membership information without restricting the iterator's state type. -/ +@[always_inline, inline] +def IterM.DefaultConsumers.measurableSpaceForIn' + (lift : ∀ (γ : Type w) (δ : Type u) [MeasurableSpace δ], (γ → m δ) → n γ → m δ) + (β : Type u) [MeasurableSpace β] (it : Std.IterM (α := ι) n α) (init : β) + (P : α → Prop) (hP : ∀ a, it.IsPlausibleIndirectOutput a → P a) + (f : (a : α) → P a → β → m (ForInStep β)) : m β := + haveI : Nonempty (m β) := ⟨mPure init⟩ + WellFounded.extrinsicFix₃ (C₃ := fun _ _ _ => m β) + (InvImage (IteratorLoop.rel ι n (fun _ _ _ => True)) (fun x => (x.1, x.2.1))) + (fun it acc (hP : ∀ a, it.IsPlausibleIndirectOutput a → P a) recur => + (lift _ _ · it.step) fun s => rdo + match s.inflate with + | .yield it' out h => + match ← f out (hP out <| .direct ⟨_, h⟩) acc with + | .yield b => + recur it' b (fun _ h' => hP _ <| .indirect ⟨_, rfl, h⟩ h') (Or.inl ⟨out, h, trivial⟩) + | .done b => return b + | .skip it' h => + recur it' acc (fun _ h' => hP _ <| .indirect ⟨_, rfl, h⟩ h') (Or.inr ⟨h, rfl⟩) + | .done _ => return acc) it init hP + +section DefaultLemmas + +variable {β : Type u} [MeasurableSpace β] + {lift : ∀ (γ : Type w) (δ : Type u) [MeasurableSpace δ], (γ → m δ) → n γ → m δ} + {it : Std.IterM (α := ι) n α} {init : β} + {P : α → Prop} {hP : ∀ a, it.IsPlausibleIndirectOutput a → P a} + {f : (a : α) → P a → β → m (ForInStep β)} + +/-- Unfold the default consumer when its iteration relation is well-founded. -/ +theorem IterM.DefaultConsumers.measurableSpaceForIn'_eq_match_step + (wf : IteratorLoop.WellFounded ι n (γ := β) (fun _ _ _ => True)) : + measurableSpaceForIn' lift β it init P hP f = + (lift _ _ · it.step) (fun s => + match s.inflate with + | .yield it' out h => + f out (hP _ <| .direct ⟨_, h⟩) init >>=ₘ fun + | .yield b => measurableSpaceForIn' lift β it' b P + (fun _ h' => hP _ <| .indirect ⟨_, rfl, h⟩ h') f + | .done b => mPure b + | .skip it' h => measurableSpaceForIn' lift β it' init P + (fun _ h' => hP _ <| .indirect ⟨_, rfl, h⟩ h') f + | .done _ => mPure init) := by + have : Nonempty (m β) := ⟨mPure init⟩ + rw [measurableSpaceForIn', WellFounded.extrinsicFix₃_eq_apply (InvImage.wf _ wf)] + congr + +/-- The default consumer is independent of the predicate used to carry membership information. -/ +theorem IterM.DefaultConsumers.measurableSpaceForIn'_congr + {it : Std.IterM (α := ι) n α} {init : β} + {P : α → Prop} {hP : ∀ a, it.IsPlausibleIndirectOutput a → P a} + {f : (a : α) → P a → β → m (ForInStep β)} + {Q : α → Prop} {hQ : ∀ a, it.IsPlausibleIndirectOutput a → Q a} + {g : (a : α) → Q a → β → m (ForInStep β)} + (wf : IteratorLoop.WellFounded ι n (γ := β) (fun _ _ _ => True)) + (hfg : ∀ a b (ha : P a) (ha' : Q a), f a ha b = g a ha' b) : + measurableSpaceForIn' lift β it init P hP f = + measurableSpaceForIn' lift β it init Q hQ g := by + rw [measurableSpaceForIn'_eq_match_step wf, measurableSpaceForIn'_eq_match_step wf] + congr; ext step + split + · congr 1 + · apply hfg + · funext result + cases result + · rfl + · apply measurableSpaceForIn'_congr wf hfg + · apply measurableSpaceForIn'_congr wf hfg + · rfl +termination_by IteratorLoop.WithWF.mk it init (hwf := wf) +decreasing_by + · exact Or.inl ⟨_, ‹_›, trivial⟩ + · exact Or.inr ⟨‹_›, rfl⟩ + +end DefaultLemmas + +/-- The generic step-based consumer. Analogue of `Std.IteratorLoop.defaultImplementation`. +Producers can supply a higher-priority instance with a specialized loop. -/ +@[always_inline, inline, instance_reducible] +def MeasurableSpaceIteratorLoop.defaultImplementation : MeasurableSpaceIteratorLoop ι n m where + forIn lift β _ it init := + IterM.DefaultConsumers.measurableSpaceForIn' lift β it init _ (fun _ => id) + +instance (priority := 100) : MeasurableSpaceIteratorLoop ι n m := .defaultImplementation + +/-- An optimized loop agrees with the default consumer on well-founded iterations, provided +lifting a pure producer value applies the supplied function to that value. -/ +class LawfulMeasurableSpaceIteratorLoop (ι : Type w) (n : Type w → Type w') + {α : Type w} [Monad n] [Iterator ι n α] + (m : (β : Type u) → [MeasurableSpace β] → Type v) [MeasurableSpaceMonad m] + [i : MeasurableSpaceIteratorLoop ι n m] : Prop where + /-- Equality with the default implementation. Analogue of `Std.LawfulIteratorLoop.lawful`. -/ + lawful (lift : ∀ (γ : Type w) (δ : Type u) [MeasurableSpace δ], (γ → m δ) → n γ → m δ) + (hlift : ∀ γ δ [MeasurableSpace δ] (f : γ → m δ) a, lift γ δ f (pure a) = f a) + (β : Type u) [MeasurableSpace β] (it : Std.IterM (α := ι) n α) (init : β) + (wf : IteratorLoop.WellFounded ι n (γ := β) (fun _ _ _ => True)) + (f : (a : α) → it.IsPlausibleIndirectOutput a → β → m (ForInStep β)) : + i.forIn lift β it init f = + MeasurableSpaceIteratorLoop.defaultImplementation.forIn lift β it init f + +instance [Monad n] : + letI : MeasurableSpaceIteratorLoop ι n m := .defaultImplementation + LawfulMeasurableSpaceIteratorLoop ι n m where + lawful _ _ _ _ _ _ _ _ := rfl + +/-- Unfold a lawful loop consumer by one producer step. -/ +theorem MeasurableSpaceIteratorLoop.forIn_eq_match_step [Monad n] + [MeasurableSpaceIteratorLoop ι n m] [LawfulMeasurableSpaceIteratorLoop ι n m] + (lift : ∀ (δ : Type w) (ε : Type u) [MeasurableSpace ε], (δ → m ε) → n δ → m ε) + (hlift : ∀ δ ε [MeasurableSpace ε] (f : δ → m ε) a, lift δ ε f (pure a) = f a) + {β : Type u} [MeasurableSpace β] (it : Std.IterM (α := ι) n α) (init : β) + (wf : IteratorLoop.WellFounded ι n (γ := β) (fun _ _ _ => True)) + (f : (a : α) → it.IsPlausibleIndirectOutput a → β → m (ForInStep β)) : + MeasurableSpaceIteratorLoop.forIn lift β it init f = + (lift _ _ · it.step) (fun s => + match s.inflate with + | .yield it' out h => f out (.direct ⟨_, h⟩) init >>=ₘ fun + | .yield b => MeasurableSpaceIteratorLoop.forIn lift β it' b + (fun a ha b => f a (.indirect ⟨_, rfl, h⟩ ha) b) + | .done b => mPure b + | .skip it' h => MeasurableSpaceIteratorLoop.forIn lift β it' init + (fun a ha b => f a (.indirect ⟨_, rfl, h⟩ ha) b) + | .done _ => mPure init) := by + rw [LawfulMeasurableSpaceIteratorLoop.lawful lift hlift β it init wf] + simp only [MeasurableSpaceIteratorLoop.forIn] + rw [IterM.DefaultConsumers.measurableSpaceForIn'_eq_match_step wf] + congr 1; funext s + cases s.inflate using Std.PlausibleIterStep.casesOn + · apply MeasurableSpaceBind.bind_congr + intro step + cases step + · rfl + · simp only + rw [LawfulMeasurableSpaceIteratorLoop.lawful lift hlift β _ _ wf] + exact IterM.DefaultConsumers.measurableSpaceForIn'_congr wf (by intros; rfl) + · simp only + rw [LawfulMeasurableSpaceIteratorLoop.lawful lift hlift β _ _ wf] + exact IterM.DefaultConsumers.measurableSpaceForIn'_congr wf (by intros; rfl) + · rfl + +/-- Membership-aware iteration. Analogue of `Std.IterM.instForIn'`. +As in core, this is opt-in because the generic membership relation is an internal predicate. -/ +@[always_inline, inline, instance_reducible] +def IterM.instMeasurableSpaceForIn' [MeasurableSpaceLiftBind n m] + [MeasurableSpaceIteratorLoop ι n m] : + MeasurableSpaceForIn' m (Std.IterM (α := ι) n α) α + ⟨fun it out => it.IsPlausibleIndirectOutput out⟩ where + forIn' {β} _ it init f := + MeasurableSpaceIteratorLoop.forIn (m := m) + (fun _ _ _ => MeasurableSpaceLiftBind.liftBind) β it init f + +instance IterM.instMeasurableSpaceForIn [MeasurableSpaceLiftBind n m] + [MeasurableSpaceIteratorLoop ι n m] : + MeasurableSpaceForIn m (Std.IterM (α := ι) n α) α := + haveI := IterM.instMeasurableSpaceForIn' (m := m) (n := n) (ι := ι) + instMeasurableSpaceForInOfForIn' + +/-- Membership-aware iteration. Analogue of `Std.Iter.instForIn'`. +As in core, this is opt-in because the generic membership relation is an internal predicate. -/ +@[always_inline, inline, instance_reducible] +def Iter.instMeasurableSpaceForIn' {ι α : Type w} [Iterator ι Id α] + [MeasurableSpaceIteratorLoop ι Id m] : + MeasurableSpaceForIn' m (Std.Iter (α := ι) α) α + ⟨fun it out => it.IsPlausibleIndirectOutput out⟩ where + forIn' {β} _ it init f := + MeasurableSpaceIteratorLoop.forIn (fun _ _ _ f x => f x.run) β + it.toIterM init fun out h acc => + f out (Std.Iter.isPlausibleIndirectOutput_iff_isPlausibleIndirectOutput_toIterM.mpr h) acc + +instance Iter.instMeasurableSpaceForIn {ι α : Type w} [Iterator ι Id α] + [MeasurableSpaceIteratorLoop ι Id m] : + MeasurableSpaceForIn m (Std.Iter (α := ι) α) α := + haveI := Iter.instMeasurableSpaceForIn' (m := m) (ι := ι) + instMeasurableSpaceForInOfForIn' + +/-- Membership-aware iteration. Analogue of `Std.Iter.Partial.instForIn'`. +As in core, this is opt-in because the generic membership relation is an internal predicate. -/ +@[always_inline, inline, instance_reducible] +def Iter.Partial.instMeasurableSpaceForIn' + {ι α : Type w} [Iterator ι Id α] [MeasurableSpaceIteratorLoop ι Id m] : + MeasurableSpaceForIn' m (Std.Iter.Partial (α := ι) α) α + ⟨fun it out => it.it.IsPlausibleIndirectOutput out⟩ where + forIn' it init f := Iter.instMeasurableSpaceForIn'.forIn' it.it init f + +instance Iter.Partial.instMeasurableSpaceForIn + {ι α : Type w} [Iterator ι Id α] [MeasurableSpaceIteratorLoop ι Id m] : + MeasurableSpaceForIn m (Std.Iter.Partial (α := ι) α) α := + haveI := Iter.Partial.instMeasurableSpaceForIn' (m := m) (ι := ι) + instMeasurableSpaceForInOfForIn' + +/-- Membership-aware iteration. Analogue of `Std.Iter.Total.instForIn'`. +As in core, this is opt-in because the generic membership relation is an internal predicate. +The `Finite` assumption restricts the total wrapper to terminating iterators; its implementation +reuses the ordinary consumer. -/ +@[always_inline, inline, instance_reducible, nolint unusedArguments] +def Iter.Total.instMeasurableSpaceForIn' + {ι α : Type w} [Iterator ι Id α] [MeasurableSpaceIteratorLoop ι Id m] [Finite ι Id] : + MeasurableSpaceForIn' m (Std.Iter.Total (α := ι) α) α + ⟨fun it out => it.it.IsPlausibleIndirectOutput out⟩ where + forIn' it init f := Iter.instMeasurableSpaceForIn'.forIn' it.it init f + +instance Iter.Total.instMeasurableSpaceForIn + {ι α : Type w} [Iterator ι Id α] [MeasurableSpaceIteratorLoop ι Id m] [Finite ι Id] : + MeasurableSpaceForIn m (Std.Iter.Total (α := ι) α) α := + haveI := Iter.Total.instMeasurableSpaceForIn' (m := m) (ι := ι) + instMeasurableSpaceForInOfForIn' + +/-- Membership-aware iteration. Analogue of `Std.IterM.Partial.instForIn'`. +As in core, this is opt-in because the generic membership relation is an internal predicate. -/ +@[always_inline, inline, instance_reducible] +def IterM.Partial.instMeasurableSpaceForIn' + [MeasurableSpaceLiftBind n m] [MeasurableSpaceIteratorLoop ι n m] : + MeasurableSpaceForIn' m (Std.IterM.Partial (α := ι) n α) α + ⟨fun it out => it.it.IsPlausibleIndirectOutput out⟩ where + forIn' it init f := IterM.instMeasurableSpaceForIn'.forIn' it.it init f + +instance IterM.Partial.instMeasurableSpaceForIn + [MeasurableSpaceLiftBind n m] [MeasurableSpaceIteratorLoop ι n m] : + MeasurableSpaceForIn m (Std.IterM.Partial (α := ι) n α) α := + haveI := IterM.Partial.instMeasurableSpaceForIn' (m := m) (n := n) (ι := ι) + instMeasurableSpaceForInOfForIn' + +/-- Membership-aware iteration. Analogue of `Std.IterM.Total.instForIn'`. +As in core, this is opt-in because the generic membership relation is an internal predicate. +The `Finite` assumption restricts the total wrapper to terminating iterators; its implementation +reuses the ordinary consumer. -/ +@[always_inline, inline, instance_reducible, nolint unusedArguments] +def IterM.Total.instMeasurableSpaceForIn' + [MeasurableSpaceLiftBind n m] [MeasurableSpaceIteratorLoop ι n m] [Finite ι n] : + MeasurableSpaceForIn' m (Std.IterM.Total (α := ι) n α) α + ⟨fun it out => it.it.IsPlausibleIndirectOutput out⟩ where + forIn' it init f := IterM.instMeasurableSpaceForIn'.forIn' it.it init f + +instance IterM.Total.instMeasurableSpaceForIn + [MeasurableSpaceLiftBind n m] [MeasurableSpaceIteratorLoop ι n m] [Finite ι n] : + MeasurableSpaceForIn m (Std.IterM.Total (α := ι) n α) α := + haveI := IterM.Total.instMeasurableSpaceForIn' (m := m) (n := n) (ι := ι) + instMeasurableSpaceForInOfForIn' + +/-- Unfold a finite pure iterator loop by one step. Analogue of `Std.Iter.forIn_eq_match_step`. +This requires no monad laws or measurable space on the iterator's internal state or elements. -/ +theorem Iter.forIn_eq_match_step {ι α : Type w} [Iterator ι Id α] + [MeasurableSpaceIteratorLoop ι Id m] [LawfulMeasurableSpaceIteratorLoop ι Id m] + [Finite ι Id] {β : Type u} [MeasurableSpace β] + (it : Std.Iter (α := ι) α) (init : β) (f : α → β → m (ForInStep β)) : + MeasureTheory.forIn it init f = + match it.step with + | .yield it' out _ => f out init >>=ₘ fun + | .yield b => MeasureTheory.forIn it' b f + | .done b => mPure b + | .skip it' _ => MeasureTheory.forIn it' init f + | .done _ => mPure init := by + simp only [MeasureTheory.forIn, MeasureTheory.forIn'] + rw [MeasurableSpaceIteratorLoop.forIn_eq_match_step _ (by intros; rfl) + _ _ IteratorLoop.wellFounded_of_finite] + simp only [Std.Iter.step] + cases it.toIterM.step.run.inflate using Std.PlausibleIterStep.casesOn <;> rfl + + +end MeasureTheory diff --git a/RandomDo/MeasurableSpaceMonad/ForInInstances/Ranges.lean b/RandomDo/MeasurableSpaceMonad/ForInInstances/Ranges.lean new file mode 100644 index 0000000..c582d69 --- /dev/null +++ b/RandomDo/MeasurableSpaceMonad/ForInInstances/Ranges.lean @@ -0,0 +1,822 @@ +/- +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.ForInInstances.Iterators +public import Init.Data.Range.Polymorphic.Iterators +import all Init.Data.Range.Polymorphic.RangeIterator + +/-! +# Range instances for measurable space monads + +These three consumers follow `Std.Rxc.Iterator.instIteratorLoop`, +`Std.Rxo.Iterator.instIteratorLoop`, and `Std.Rxi.Iterator.instIteratorLoop` in +`Init/Data/Range/Polymorphic/RangeIterator.lean`. They keep the next value directly in the loop +arguments, avoiding the `Option` in the generic range iterator state. + +Each implementation is shared by all measurable space monads. The accompanying lawfulness +proofs follow core's well-founded reference loops and show agreement with the default consumer. +They use congruence of bind, without assuming monad laws or measurability of the body. + +The nine direct range instances follow `Init/Data/Range/Polymorphic/Iterators.lean`, delegating +to these consumers and converting iterator membership proofs to range membership proofs. +They support `for i in a...b rdo ...` and the other open, closed, and unbounded range forms, +with the same finiteness requirements as core. The `forIn_eq_forIn_iter` lemmas expose this +connection for proofs about range loops. +-/ + +@[expose] public section + +namespace MeasureTheory + +open Std Std.Iterators Std.PRange MeasurableSpacePure MeasurableSpaceBind + +open MeasureTheory.IterM.DefaultConsumers + +universe u v w + +namespace Rxc + +open Std.Rxc Std.Rxc.Iterator + +variable {α : Type w} + +/-- Direct range iteration without allocating the `Option` that records exhaustion on each step. +Analogue of `Std.Rxc.Iterator.instIteratorLoop`. -/ +@[always_inline, inline] +instance Iterator.instMeasurableSpaceIteratorLoop [UpwardEnumerable α] [LE α] [DecidableLE α] + [LawfulUpwardEnumerable α] [LawfulUpwardEnumerableLE α] + {m : (β : Type u) → [MeasurableSpace β] → Type v} [MeasurableSpaceMonad m] : + MeasurableSpaceIteratorLoop (Rxc.Iterator α) Id m where + forIn _ γ _ it init f := + match it with + | ⟨⟨some next, upperBound⟩⟩ => + loop γ (next ≤ ·) (fun a b hab hna => ?hle) + upperBound init next ?hle'' (fun a ha₁ ha₂ c => f a ?hf c) + | ⟨⟨none, _⟩⟩ => mPure init + where + @[always_inline, inline] + loop (γ : Type u) [MeasurableSpace γ] + (LargeEnough : α → Prop) + (hl : ∀ a b : α, a ≤ b → LargeEnough a → LargeEnough b) + (upperBound : α) (acc : γ) (next : α) (h : LargeEnough next) + (f : (out : α) → LargeEnough out → out ≤ upperBound → (c : γ) → m (ForInStep γ)) : m γ := + haveI : Nonempty (m γ) := ⟨mPure acc⟩ + WellFounded.extrinsicFix₃ (C₃ := fun _ _ _ => m γ) + (InvImage (IteratorLoop.rel _ Id (fun _ _ _ => True)) + (fun x => (⟨Rxc.Iterator.mk (some x.1) upperBound⟩, x.2.1))) + (fun next acc (h : LargeEnough next) G => rdo + if hu : next ≤ upperBound then + match ← f next h hu acc with + | .yield acc' => + match hs : UpwardEnumerable.succ? next with + | some next' => G next' acc' (hl _ _ ?hle' h) ?decreasing + | none => mPure acc' + | .done acc' => mPure acc' + else + mPure acc) next acc h + finally + case hf => + rw [Monadic.isPlausibleIndirectOutput_iff] + simp only [UpwardEnumerable.le_iff] at ha₁ + obtain ⟨n, hn⟩ := ha₁ + exact ⟨n, hn, ha₂⟩ + case hle => + simp only [UpwardEnumerable.le_iff] at hna hab ⊢ + exact UpwardEnumerable.le_trans hna hab + case hle' => + simp only [UpwardEnumerable.le_iff] + refine ⟨1, ?_⟩ + simpa [succMany?_one] using hs + case hle'' => + exact UpwardEnumerable.le_iff.mpr (UpwardEnumerable.le_refl _) + case decreasing => + simp_wf + simp [IteratorLoop.rel, Monadic.isPlausibleStep_iff, Monadic.step, *] + +private noncomputable def Iterator.instMeasurableSpaceIteratorLoop.loop.wf + [UpwardEnumerable α] [LE α] [DecidableLE α] + [LawfulUpwardEnumerable α] [LawfulUpwardEnumerableLE α] + {m : (β : Type u) → [MeasurableSpace β] → Type v} [MeasurableSpaceMonad m] + (γ : Type u) [MeasurableSpace γ] + (wf : IteratorLoop.WellFounded (Rxc.Iterator α) Id (γ := γ) (fun _ _ _ => True)) + (LargeEnough : α → Prop) (hl : ∀ a b : α, a ≤ b → LargeEnough a → LargeEnough b) + (upperBound : α) (acc : γ) (next : α) (h : LargeEnough next) + (f : (out : α) → LargeEnough out → out ≤ upperBound → (c : γ) → m (ForInStep γ)) : + m γ := rdo + if hu : next ≤ upperBound then + match ← f next h hu acc with + | .yield acc' => + match hs : UpwardEnumerable.succ? next with + | some next' => + loop.wf γ wf LargeEnough hl upperBound acc' next' (hl _ _ ?hle h) f + | none => return acc' + | .done acc' => return acc' + else + return acc +termination_by IteratorLoop.WithWF.mk ⟨⟨some next, upperBound⟩⟩ acc (hwf := wf) +decreasing_by + simp [IteratorLoop.rel, Monadic.isPlausibleStep_iff, Monadic.step, *] +where finally + case hle => + simp only [UpwardEnumerable.le_iff] + refine ⟨1, ?_⟩ + simpa [succMany?_one] using hs + +private theorem Iterator.instMeasurableSpaceIteratorLoop.loop_eq_wf + [UpwardEnumerable α] [LE α] [DecidableLE α] + [LawfulUpwardEnumerable α] [LawfulUpwardEnumerableLE α] + {m : (β : Type u) → [MeasurableSpace β] → Type v} [MeasurableSpaceMonad m] + {γ : Type u} [MeasurableSpace γ] {LargeEnough hl upperBound} {next hn} {acc} (wf f) : + loop γ LargeEnough hl upperBound acc next hn f = + loop.wf (α := α) (m := m) γ wf LargeEnough hl upperBound acc next hn f := by + have : Nonempty (m γ) := ⟨mPure acc⟩ + rw [loop, WellFounded.extrinsicFix₃_eq_fix]; rotate_left + · exact InvImage.wf _ wf + · fun_induction loop.wf γ wf LargeEnough hl upperBound acc next hn f + · rw [WellFounded.fix_eq] + simp only [↓reduceDIte, *] + apply MeasurableSpaceBind.bind_congr; intro forInStep + split + · split + · simp_all + · simp + · simp + · rw [WellFounded.fix_eq] + simp_all + +private theorem Iterator.instMeasurableSpaceIteratorLoop.loopWf_eq + [UpwardEnumerable α] [LE α] [DecidableLE α] + [LawfulUpwardEnumerable α] [LawfulUpwardEnumerableLE α] + {m : (β : Type u) → [MeasurableSpace β] → Type v} [MeasurableSpaceMonad m] + (γ : Type u) [MeasurableSpace γ] + {lift : ∀ (δ : Type w) (ε : Type u) [MeasurableSpace ε], (δ → m ε) → Id δ → m ε} + (hlift : ∀ δ ε [MeasurableSpace ε] (f : δ → m ε) a, lift δ ε f (pure a) = f a) + (wf : IteratorLoop.WellFounded (Rxc.Iterator α) Id (γ := γ) (fun _ _ _ => True)) + (LargeEnough : α → Prop) (hl : ∀ a b : α, a ≤ b → LargeEnough a → LargeEnough b) + (upperBound : α) (acc : γ) (next : α) (h : LargeEnough next) + (f : (out : α) → LargeEnough out → out ≤ upperBound → (c : γ) → m (ForInStep γ)) : + loop.wf γ wf LargeEnough hl upperBound acc next h f = (rdo + if hu : next ≤ upperBound then + match ← f next h hu acc with + | .yield acc' => + letI it' : IterM (α := Rxc.Iterator α) Id α := ⟨⟨succ? next, upperBound⟩⟩ + measurableSpaceForIn' (n := Id) (m := m) lift γ it' acc' + it'.IsPlausibleIndirectOutput (fun _ => id) + fun next' h acc' => f next' + (by + refine hl next next' ?_ ‹_› + simp only [it', Monadic.isPlausibleIndirectOutput_iff, + ← succMany?_add_one_eq_succ?_bind_succMany?] at h + exact UpwardEnumerable.le_iff.mpr ⟨h.choose + 1, h.choose_spec.1⟩) + (by + simp only [it', Monadic.isPlausibleIndirectOutput_iff] at h + exact h.choose_spec.2) + acc' + | .done acc' => return acc' + else return acc) := by + have : Nonempty (m γ) := ⟨mPure acc⟩ + rw [loop.wf] + congr 1; ext hu + apply MeasurableSpaceBind.bind_congr; intro forInStep + split + · split + · rw [loopWf_eq (lift := lift) _ hlift wf] + rw [measurableSpaceForIn'_eq_match_step (lift := lift) wf] + · simp only [IterM.step_mk, Monadic.step_eq_step, Monadic.step, + Shrink.inflate_deflate, *] + split + · apply MeasurableSpaceBind.bind_congr; intro forInStep + split + · apply measurableSpaceForIn'_congr wf ; intros; rfl + · simp + · simp + · rw [measurableSpaceForIn'_eq_match_step wf] + simp only [IterM.step_eq, Shrink.inflate_deflate, *] + -- Eliminate the lift and `Shrink` before unfolding the step and its proof indices. + simp [Monadic.step] + · simp +termination_by IteratorLoop.WithWF.mk ⟨⟨some next, upperBound⟩⟩ acc (hwf := wf) +decreasing_by + simp [IteratorLoop.rel, Monadic.isPlausibleStep_iff, Monadic.step, *] + +instance Iterator.instLawfulMeasurableSpaceIteratorLoop [UpwardEnumerable α] [LE α] [DecidableLE α] + [LawfulUpwardEnumerable α] [LawfulUpwardEnumerableLE α] + {m : (β : Type u) → [MeasurableSpace β] → Type v} [MeasurableSpaceMonad m] : + LawfulMeasurableSpaceIteratorLoop (Rxc.Iterator α) Id m where + lawful := by + intro lift hlift γ _ it init wf f + simp only [MeasurableSpaceIteratorLoop.forIn] + rw [measurableSpaceForIn'_eq_match_step wf] + split; rotate_left + · simp only [Std.IterM.step_eq, hlift, Shrink.inflate_deflate] + simp [Monadic.step] + rename_i next _ + rw [instMeasurableSpaceIteratorLoop.loop_eq_wf wf, + instMeasurableSpaceIteratorLoop.loopWf_eq (lift := lift) _ hlift wf] + simp only [Std.IterM.step_eq, hlift, Shrink.inflate_deflate] + simp only [Monadic.step] + split + · apply MeasurableSpaceBind.bind_congr; intro forInStep + split + · apply measurableSpaceForIn'_congr wf + intros; rfl + · rfl + · rfl + +end Rxc + +namespace Rxo + +open Std.Rxo Std.Rxo.Iterator + +variable {α : Type w} + +/-- Direct range iteration without allocating the `Option` that records exhaustion on each step. +Analogue of `Std.Rxo.Iterator.instIteratorLoop`. -/ +@[always_inline, inline] +instance Iterator.instMeasurableSpaceIteratorLoop [UpwardEnumerable α] [LT α] [DecidableLT α] + [LawfulUpwardEnumerable α] [LawfulUpwardEnumerableLT α] + {m : (β : Type u) → [MeasurableSpace β] → Type v} [MeasurableSpaceMonad m] : + MeasurableSpaceIteratorLoop (Rxo.Iterator α) Id m where + forIn _ γ _ it init f := + match it with + | ⟨⟨some next, upperBound⟩⟩ => + loop γ (UpwardEnumerable.LE next ·) (fun a b hab hna => ?hle) + upperBound init next ?hle'' (fun a ha₁ ha₂ c => f a ?hf c) + | ⟨⟨none, _⟩⟩ => mPure init + where + @[always_inline, inline] + loop (γ : Type u) [MeasurableSpace γ] (LargeEnough : α → Prop) + (hl : ∀ a b : α, UpwardEnumerable.LE a b → LargeEnough a → LargeEnough b) + (upperBound : α) (acc : γ) (next : α) (h : LargeEnough next) + (f : (out : α) → LargeEnough out → out < upperBound → (c : γ) → m (ForInStep γ)) : m γ := + haveI : Nonempty (m γ) := ⟨mPure acc⟩ + WellFounded.extrinsicFix₃ (C₃ := fun _ _ _ => m γ) + (InvImage (IteratorLoop.rel _ Id (fun _ _ _ => True)) + (fun x => (⟨Rxo.Iterator.mk (some x.1) upperBound⟩, x.2.1))) + (fun next acc (h : LargeEnough next) G => rdo + if hu : next < upperBound then + match ← f next h hu acc with + | .yield acc' => + match hs : UpwardEnumerable.succ? next with + | some next' => G next' acc' (hl _ _ ?hle' h) ?decreasing + | none => mPure acc' + | .done acc' => mPure acc' + else + mPure acc) next acc h + finally + case hf => + rw [Monadic.isPlausibleIndirectOutput_iff] + obtain ⟨n, hn⟩ := ha₁ + exact ⟨n, hn, ha₂⟩ + case hle => + exact UpwardEnumerable.le_trans hna hab + case hle' => + refine ⟨1, ?_⟩ + simpa [succMany?_one] using hs + case hle'' => + exact UpwardEnumerable.le_refl _ + case decreasing => + simp_wf; simp [IteratorLoop.rel, Monadic.isPlausibleStep_iff, Monadic.step, *] + +private noncomputable def Iterator.instMeasurableSpaceIteratorLoop.loop.wf + [UpwardEnumerable α] [LT α] [DecidableLT α] + [LawfulUpwardEnumerable α] [LawfulUpwardEnumerableLT α] + {m : (β : Type u) → [MeasurableSpace β] → Type v} [MeasurableSpaceMonad m] + (γ : Type u) [MeasurableSpace γ] + (wf : IteratorLoop.WellFounded (Rxo.Iterator α) Id (γ := γ) (fun _ _ _ => True)) + (LargeEnough : α → Prop) + (hl : ∀ a b : α, UpwardEnumerable.LE a b → LargeEnough a → LargeEnough b) + (upperBound : α) (acc : γ) (next : α) (h : LargeEnough next) + (f : (out : α) → LargeEnough out → out < upperBound → (c : γ) → m (ForInStep γ)) : + m γ := rdo + if hu : next < upperBound then + match ← f next h hu acc with + | .yield acc' => + match hs : UpwardEnumerable.succ? next with + | some next' => + loop.wf γ wf LargeEnough hl upperBound acc' next' (hl _ _ ?hle h) f + | none => return acc' + | .done acc' => return acc' + else + return acc +termination_by IteratorLoop.WithWF.mk ⟨⟨some next, upperBound⟩⟩ acc (hwf := wf) +decreasing_by + simp [IteratorLoop.rel, Monadic.isPlausibleStep_iff, Monadic.step, *] +where finally + case hle => + refine ⟨1, ?_⟩ + simpa [succMany?_one] using hs + +private theorem Iterator.instMeasurableSpaceIteratorLoop.loop_eq_wf + [UpwardEnumerable α] [LT α] [DecidableLT α] + [LawfulUpwardEnumerable α] [LawfulUpwardEnumerableLT α] + {m : (β : Type u) → [MeasurableSpace β] → Type v} [MeasurableSpaceMonad m] + {γ : Type u} [MeasurableSpace γ] {LargeEnough hl upperBound} {next hn} {acc} (wf f) : + loop γ LargeEnough hl upperBound acc next hn f = + loop.wf (α := α) (m := m) γ wf LargeEnough hl upperBound acc next hn f := by + have : Nonempty (m γ) := ⟨mPure acc⟩ + rw [loop, WellFounded.extrinsicFix₃_eq_fix]; rotate_left + · exact InvImage.wf _ wf + · fun_induction loop.wf γ wf LargeEnough hl upperBound acc next hn f + · rw [WellFounded.fix_eq] + simp only [↓reduceDIte, *] + apply MeasurableSpaceBind.bind_congr; intro forInStep + split + · split + · simp_all + · simp + · simp + · rw [WellFounded.fix_eq] + simp_all + +private theorem Iterator.instMeasurableSpaceIteratorLoop.loopWf_eq + [UpwardEnumerable α] [LT α] [DecidableLT α] + [LawfulUpwardEnumerable α] [LawfulUpwardEnumerableLT α] + {m : (β : Type u) → [MeasurableSpace β] → Type v} [MeasurableSpaceMonad m] + (γ : Type u) [MeasurableSpace γ] + {lift : ∀ (δ : Type w) (ε : Type u) [MeasurableSpace ε], (δ → m ε) → Id δ → m ε} + (hlift : ∀ δ ε [MeasurableSpace ε] (f : δ → m ε) a, lift δ ε f (pure a) = f a) + (wf : IteratorLoop.WellFounded (Rxo.Iterator α) Id (γ := γ) (fun _ _ _ => True)) + (LargeEnough : α → Prop) + (hl : ∀ a b : α, UpwardEnumerable.LE a b → LargeEnough a → LargeEnough b) + (upperBound : α) (acc : γ) (next : α) (h : LargeEnough next) + (f : (out : α) → LargeEnough out → out < upperBound → (c : γ) → m (ForInStep γ)) : + loop.wf γ wf LargeEnough hl upperBound acc next h f = (rdo + if hu : next < upperBound then + match ← f next h hu acc with + | .yield acc' => + letI it' : IterM (α := Rxo.Iterator α) Id α := ⟨⟨succ? next, upperBound⟩⟩ + measurableSpaceForIn' (n := Id) (m := m) lift γ it' acc' + it'.IsPlausibleIndirectOutput (fun _ => id) + fun next' h acc' => f next' + (by + refine hl next next' ?_ ‹_› + simp only [it', Monadic.isPlausibleIndirectOutput_iff, + ← succMany?_add_one_eq_succ?_bind_succMany?] at h + exact ⟨h.choose + 1, h.choose_spec.1⟩) + (by + simp only [it', Monadic.isPlausibleIndirectOutput_iff] at h + exact h.choose_spec.2) + acc' + | .done acc' => return acc' + else return acc) := by + have : Nonempty (m γ) := ⟨mPure acc⟩ + rw [loop.wf] + congr 1; ext hu + apply MeasurableSpaceBind.bind_congr; intro forInStep + split + · split + · rw [loopWf_eq (lift := lift) _ hlift wf] + rw [measurableSpaceForIn'_eq_match_step (lift := lift) wf] + · simp only [Std.IterM.step_eq, Monadic.step, + Shrink.inflate_deflate, *] + split + · apply MeasurableSpaceBind.bind_congr; intro forInStep + split + · apply measurableSpaceForIn'_congr wf ; intros; rfl + · simp + · simp + · rw [measurableSpaceForIn'_eq_match_step wf] + simp [Std.IterM.step_eq, Monadic.step, *] + · simp +termination_by IteratorLoop.WithWF.mk ⟨⟨some next, upperBound⟩⟩ acc (hwf := wf) +decreasing_by + simp [IteratorLoop.rel, Monadic.isPlausibleStep_iff, Monadic.step, *] + +instance Iterator.instLawfulMeasurableSpaceIteratorLoop [UpwardEnumerable α] [LT α] [DecidableLT α] + [LawfulUpwardEnumerable α] [LawfulUpwardEnumerableLT α] + {m : (β : Type u) → [MeasurableSpace β] → Type v} [MeasurableSpaceMonad m] : + LawfulMeasurableSpaceIteratorLoop (Rxo.Iterator α) Id m where + lawful := by + intro lift hlift γ _ it init wf f + simp only [MeasurableSpaceIteratorLoop.forIn] + rw [measurableSpaceForIn'_eq_match_step wf] + split; rotate_left + · simp only [Std.IterM.step_eq, hlift, Shrink.inflate_deflate] + simp [Monadic.step] + rename_i next _ + rw [instMeasurableSpaceIteratorLoop.loop_eq_wf wf, + instMeasurableSpaceIteratorLoop.loopWf_eq (lift := lift) _ hlift wf] + simp only [Std.IterM.step_eq, hlift, Shrink.inflate_deflate] + simp only [Monadic.step] + split + · apply MeasurableSpaceBind.bind_congr; intro forInStep + split + · apply measurableSpaceForIn'_congr wf + intros; rfl + · rfl + · rfl + +end Rxo + +namespace Rxi + +open Std.Rxi Std.Rxi.Iterator + +variable {α : Type w} + +/-- Direct range iteration without allocating the `Option` that records exhaustion on each step. +Analogue of `Std.Rxi.Iterator.instIteratorLoop`. -/ +@[always_inline, inline] +instance Iterator.instMeasurableSpaceIteratorLoop [UpwardEnumerable α] [LawfulUpwardEnumerable α] + {m : (β : Type u) → [MeasurableSpace β] → Type v} [MeasurableSpaceMonad m] : + MeasurableSpaceIteratorLoop (Rxi.Iterator α) Id m where + forIn _ γ _ it init f := + match it with + | ⟨⟨some next⟩⟩ => + loop γ (UpwardEnumerable.LE next ·) (fun a b hab hna => ?hle) + init next ?hle'' (fun a ha c => f a ?hf c) + | ⟨⟨none⟩⟩ => mPure init + where + @[always_inline, inline] + loop (γ : Type u) [MeasurableSpace γ] + (LargeEnough : α → Prop) + (hl : ∀ a b : α, UpwardEnumerable.LE a b → LargeEnough a → LargeEnough b) + (acc : γ) (next : α) (h : LargeEnough next) + (f : (out : α) → LargeEnough out → (c : γ) → m (ForInStep γ)) : m γ := + haveI : Nonempty (m γ) := ⟨mPure acc⟩ + WellFounded.extrinsicFix₃ (C₃ := fun _ _ _ => m γ) + (InvImage (IteratorLoop.rel _ Id (fun _ _ _ => True)) + (fun x => (⟨Rxi.Iterator.mk (some x.1)⟩, x.2.1))) + (fun next acc (h : LargeEnough next) G => rdo + match ← f next h acc with + | .yield acc' => + match hs : UpwardEnumerable.succ? next with + | some next' => G next' acc' (hl _ _ ?hle' h) ?decreasing + | none => mPure acc' + | .done acc' => mPure acc') next acc h + finally + case hf => + rw [Monadic.isPlausibleIndirectOutput_iff] + exact ha + case hle => + exact UpwardEnumerable.le_trans hna hab + case hle' => + refine ⟨1, ?_⟩ + simpa [succMany?_one] using hs + case hle'' => + exact UpwardEnumerable.le_refl _ + case decreasing => + simp_wf; simp [IteratorLoop.rel, Monadic.isPlausibleStep_iff, Monadic.step, *] + +private noncomputable def Iterator.instMeasurableSpaceIteratorLoop.loop.wf [UpwardEnumerable α] + [LawfulUpwardEnumerable α] + {m : (β : Type u) → [MeasurableSpace β] → Type v} [MeasurableSpaceMonad m] + (γ : Type u) [MeasurableSpace γ] + (wf : IteratorLoop.WellFounded (Rxi.Iterator α) Id (γ := γ) (fun _ _ _ => True)) + (LargeEnough : α → Prop) + (hl : ∀ a b : α, UpwardEnumerable.LE a b → LargeEnough a → LargeEnough b) + (acc : γ) (next : α) (h : LargeEnough next) + (f : (out : α) → LargeEnough out → (c : γ) → m (ForInStep γ)) : + m γ := rdo + match ← f next h acc with + | .yield acc' => + match hs : UpwardEnumerable.succ? next with + | some next' => + loop.wf γ wf LargeEnough hl acc' next' (hl _ _ ?hle h) f + | none => return acc' + | .done acc' => return acc' +termination_by IteratorLoop.WithWF.mk ⟨⟨some next⟩⟩ acc (hwf := wf) +decreasing_by + simp [IteratorLoop.rel, Monadic.isPlausibleStep_iff, Monadic.step, *] +where finally + case hle => + refine ⟨1, ?_⟩ + simpa [succMany?_one] using hs + +private theorem Iterator.instMeasurableSpaceIteratorLoop.loop_eq_wf [UpwardEnumerable α] + [LawfulUpwardEnumerable α] + {m : (β : Type u) → [MeasurableSpace β] → Type v} [MeasurableSpaceMonad m] + {γ : Type u} [MeasurableSpace γ] {LargeEnough hl} {next hn} {acc} (wf f) : + loop γ LargeEnough hl acc next hn f = + loop.wf (α := α) (m := m) γ wf LargeEnough hl acc next hn f := by + have : Nonempty (m γ) := ⟨mPure acc⟩ + rw [loop, WellFounded.extrinsicFix₃_eq_fix]; rotate_left + · exact InvImage.wf _ wf + · fun_induction loop.wf γ wf LargeEnough hl acc next hn f + · rw [WellFounded.fix_eq] + dsimp only + apply MeasurableSpaceBind.bind_congr; intro forInStep + split + · split + · simp_all + · simp + · simp + +private theorem Iterator.instMeasurableSpaceIteratorLoop.loopWf_eq [UpwardEnumerable α] + [LawfulUpwardEnumerable α] + {m : (β : Type u) → [MeasurableSpace β] → Type v} [MeasurableSpaceMonad m] + (γ : Type u) [MeasurableSpace γ] + {lift : ∀ (δ : Type w) (ε : Type u) [MeasurableSpace ε], (δ → m ε) → Id δ → m ε} + (hlift : ∀ δ ε [MeasurableSpace ε] (f : δ → m ε) a, lift δ ε f (pure a) = f a) + (wf : IteratorLoop.WellFounded (Rxi.Iterator α) Id (γ := γ) (fun _ _ _ => True)) + (LargeEnough : α → Prop) + (hl : ∀ a b : α, UpwardEnumerable.LE a b → LargeEnough a → LargeEnough b) + (acc : γ) (next : α) (h : LargeEnough next) + (f : (out : α) → LargeEnough out → (c : γ) → m (ForInStep γ)) : + loop.wf γ wf LargeEnough hl acc next h f = (rdo + match ← f next h acc with + | .yield acc' => + letI it' : IterM (α := Rxi.Iterator α) Id α := ⟨⟨succ? next⟩⟩ + measurableSpaceForIn' (n := Id) (m := m) lift γ it' acc' + it'.IsPlausibleIndirectOutput (fun _ => id) + fun next' h acc' => f next' + (by + refine hl next next' ?_ ‹_› + simp only [it', Monadic.isPlausibleIndirectOutput_iff, + ← succMany?_add_one_eq_succ?_bind_succMany?] at h + exact ⟨h.choose + 1, h.choose_spec⟩) + acc' + | .done acc' => return acc') := by + have : Nonempty (m γ) := ⟨mPure acc⟩ + rw [loop.wf] + apply MeasurableSpaceBind.bind_congr; intro forInStep + split + · split + · rw [loopWf_eq (lift := lift) _ hlift wf] + rw [measurableSpaceForIn'_eq_match_step (lift := lift) wf] + · simp only [Monadic.step_eq_step, Monadic.step, + Shrink.inflate_deflate, *] + apply MeasurableSpaceBind.bind_congr; intro forInStep + split + · apply measurableSpaceForIn'_congr wf ; intros; rfl + · simp + · rw [measurableSpaceForIn'_eq_match_step wf] + simp [Monadic.step_eq_step, Monadic.step, *] + · simp +termination_by IteratorLoop.WithWF.mk ⟨⟨some next⟩⟩ acc (hwf := wf) +decreasing_by + simp [IteratorLoop.rel, Monadic.isPlausibleStep_iff, Monadic.step, *] + +instance Iterator.instLawfulMeasurableSpaceIteratorLoop [UpwardEnumerable α] + [LawfulUpwardEnumerable α] + {m : (β : Type u) → [MeasurableSpace β] → Type v} [MeasurableSpaceMonad m] : + LawfulMeasurableSpaceIteratorLoop (Rxi.Iterator α) Id m where + lawful := by + intro lift hlift γ _ it init wf f + simp only [MeasurableSpaceIteratorLoop.forIn] + rw [measurableSpaceForIn'_eq_match_step wf] + split; rotate_left + · simp only [Monadic.step_eq_step, hlift, Shrink.inflate_deflate] + simp [Monadic.step] + rename_i next _ + rw [instMeasurableSpaceIteratorLoop.loop_eq_wf wf, + instMeasurableSpaceIteratorLoop.loopWf_eq (lift := lift) _ hlift wf] + simp only [Monadic.step_eq_step, hlift, Shrink.inflate_deflate] + simp only [Monadic.step] + apply MeasurableSpaceBind.bind_congr; intro forInStep + split + · apply measurableSpaceForIn'_congr wf + intros; rfl + · rfl + +end Rxi + +/-! ## Direct range instances + +As in core, these instances are defaults so that the element type can determine the range's +endpoint type. The `Finite` assumptions retain core's termination requirements; the implementations +reuse the consumers above. +-/ + +variable {α : Type w} {m : (β : Type u) → [MeasurableSpace β] → Type v} + {β : Type u} [MeasurableSpace β] + +namespace Rcc + +open Std.Rcc + +variable [UpwardEnumerable α] + [LE α] [DecidableLE α] [LawfulUpwardEnumerable α] [LawfulUpwardEnumerableLE α] + [MeasurableSpaceMonad m] + +@[no_expose, default_instance, nolint unusedArguments] +instance instMeasurableSpaceForIn' [Finite (Rxc.Iterator α) Id] : + MeasurableSpaceForIn' m (Rcc α) α inferInstance where + forIn' r init f := by + haveI := Iter.instMeasurableSpaceForIn' (ι := Rxc.Iterator α) (m := m) + refine MeasurableSpaceForIn'.forIn' (α := α) (Internal.iter r) init (fun a ha acc => f a ?_ acc) + simp only [Membership.mem] at ha + rwa [Internal.isPlausibleIndirectOutput_iter_iff] at ha + +/-- A range loop is the corresponding iterator loop. -/ +theorem forIn_eq_forIn_iter [Finite (Rxc.Iterator α) Id] + (r : Rcc α) (init : β) (f : α → β → m (ForInStep β)) : + MeasureTheory.forIn r init f = MeasureTheory.forIn (Internal.iter r) init f := by + rfl + +end Rcc + +namespace Rco + +open Std.Rco + +variable [UpwardEnumerable α] [LE α] [LT α] [DecidableLT α] + [LawfulUpwardEnumerable α] [LawfulUpwardEnumerableLE α] [LawfulUpwardEnumerableLT α] + [MeasurableSpaceMonad m] + +@[no_expose, default_instance, nolint unusedArguments] +instance instMeasurableSpaceForIn' [Finite (Rxo.Iterator α) Id] : + MeasurableSpaceForIn' m (Rco α) α inferInstance where + forIn' r init f := by + haveI := Iter.instMeasurableSpaceForIn' (ι := Rxo.Iterator α) (m := m) + refine MeasurableSpaceForIn'.forIn' (α := α) (Internal.iter r) init (fun a ha acc => f a ?_ acc) + simp only [Membership.mem] at ha + rwa [Internal.isPlausibleIndirectOutput_iter_iff] at ha + +/-- A range loop is the corresponding iterator loop. -/ +theorem forIn_eq_forIn_iter [Finite (Rxo.Iterator α) Id] + (r : Rco α) (init : β) (f : α → β → m (ForInStep β)) : + MeasureTheory.forIn r init f = MeasureTheory.forIn (Internal.iter r) init f := by + rfl + +end Rco + +namespace Rci + +open Std.Rci + +variable [UpwardEnumerable α] + [LE α] + [LawfulUpwardEnumerable α] [LawfulUpwardEnumerableLE α] + [MeasurableSpaceMonad m] + +@[no_expose, default_instance, nolint unusedArguments] +instance instMeasurableSpaceForIn' [Finite (Rxi.Iterator α) Id] : + MeasurableSpaceForIn' m (Rci α) α inferInstance where + forIn' r init f := by + haveI := Iter.instMeasurableSpaceForIn' (ι := Rxi.Iterator α) (m := m) + refine MeasurableSpaceForIn'.forIn' (α := α) (Internal.iter r) init (fun a ha acc => f a ?_ acc) + simp only [Membership.mem] at ha + rwa [Internal.isPlausibleIndirectOutput_iter_iff] at ha + +/-- A range loop is the corresponding iterator loop. -/ +theorem forIn_eq_forIn_iter [Finite (Rxi.Iterator α) Id] + (r : Rci α) (init : β) (f : α → β → m (ForInStep β)) : + MeasureTheory.forIn r init f = MeasureTheory.forIn (Internal.iter r) init f := by + rfl + +end Rci + +namespace Roc + +open Std.Roc + +variable [UpwardEnumerable α] + [LE α] [DecidableLE α] [LawfulUpwardEnumerable α] [LawfulUpwardEnumerableLE α] + [LT α] [LawfulUpwardEnumerableLT α] + [MeasurableSpaceMonad m] + +@[no_expose, default_instance, nolint unusedArguments] +instance instMeasurableSpaceForIn' [Finite (Rxc.Iterator α) Id] : + MeasurableSpaceForIn' m (Roc α) α inferInstance where + forIn' r init f := by + haveI := Iter.instMeasurableSpaceForIn' (ι := Rxc.Iterator α) (m := m) + refine MeasurableSpaceForIn'.forIn' (α := α) (Internal.iter r) init (fun a ha acc => f a ?_ acc) + simp only [Membership.mem] at ha + rwa [Internal.isPlausibleIndirectOutput_iter_iff] at ha + +/-- A range loop is the corresponding iterator loop. -/ +theorem forIn_eq_forIn_iter [Finite (Rxc.Iterator α) Id] + (r : Roc α) (init : β) (f : α → β → m (ForInStep β)) : + MeasureTheory.forIn r init f = MeasureTheory.forIn (Internal.iter r) init f := by + rfl + +end Roc + +namespace Roo + +open Std.Roo + +variable [UpwardEnumerable α] + [LT α] [DecidableLT α] + [LawfulUpwardEnumerable α] [LawfulUpwardEnumerableLT α] + [MeasurableSpaceMonad m] + +@[no_expose, default_instance, nolint unusedArguments] +instance instMeasurableSpaceForIn' [Finite (Rxo.Iterator α) Id] : + MeasurableSpaceForIn' m (Roo α) α inferInstance where + forIn' r init f := by + haveI := Iter.instMeasurableSpaceForIn' (ι := Rxo.Iterator α) (m := m) + refine MeasurableSpaceForIn'.forIn' (α := α) (Internal.iter r) init (fun a ha acc => f a ?_ acc) + simp only [Membership.mem] at ha + rwa [Internal.isPlausibleIndirectOutput_iter_iff] at ha + +/-- A range loop is the corresponding iterator loop. -/ +theorem forIn_eq_forIn_iter [Finite (Rxo.Iterator α) Id] + (r : Roo α) (init : β) (f : α → β → m (ForInStep β)) : + MeasureTheory.forIn r init f = MeasureTheory.forIn (Internal.iter r) init f := by + rfl + +end Roo + +namespace Roi + +open Std.Roi + +variable [UpwardEnumerable α] + [LT α] [LawfulUpwardEnumerable α] [LawfulUpwardEnumerableLT α] + [MeasurableSpaceMonad m] + +@[no_expose, default_instance, nolint unusedArguments] +instance instMeasurableSpaceForIn' [Finite (Rxi.Iterator α) Id] : + MeasurableSpaceForIn' m (Roi α) α inferInstance where + forIn' r init f := by + haveI := Iter.instMeasurableSpaceForIn' (ι := Rxi.Iterator α) (m := m) + refine MeasurableSpaceForIn'.forIn' (α := α) (Internal.iter r) init (fun a ha acc => f a ?_ acc) + simp only [Membership.mem] at ha + rwa [Internal.isPlausibleIndirectOutput_iter_iff] at ha + +/-- A range loop is the corresponding iterator loop. -/ +theorem forIn_eq_forIn_iter [Finite (Rxi.Iterator α) Id] + (r : Roi α) (init : β) (f : α → β → m (ForInStep β)) : + MeasureTheory.forIn r init f = MeasureTheory.forIn (Internal.iter r) init f := by + rfl + +end Roi + +namespace Ric + +open Std.Ric + +variable [UpwardEnumerable α] + [LE α] [DecidableLE α] [Least? α] + [LawfulUpwardEnumerable α] [LawfulUpwardEnumerableLE α] + [LawfulUpwardEnumerableLeast? α] + [MeasurableSpaceMonad m] + +@[no_expose, default_instance, nolint unusedArguments] +instance instMeasurableSpaceForIn' [Finite (Rxc.Iterator α) Id] : + MeasurableSpaceForIn' m (Ric α) α inferInstance where + forIn' r init f := by + haveI := Iter.instMeasurableSpaceForIn' (ι := Rxc.Iterator α) (m := m) + refine MeasurableSpaceForIn'.forIn' (α := α) (Internal.iter r) init (fun a ha acc => f a ?_ acc) + simp only [Membership.mem] at ha + rwa [Internal.isPlausibleIndirectOutput_iter_iff] at ha + +/-- A range loop is the corresponding iterator loop. -/ +theorem forIn_eq_forIn_iter [Finite (Rxc.Iterator α) Id] + (r : Ric α) (init : β) (f : α → β → m (ForInStep β)) : + MeasureTheory.forIn r init f = MeasureTheory.forIn (Internal.iter r) init f := by + rfl + +end Ric + +namespace Rio + +open Std.Rio + +variable [UpwardEnumerable α] + [LT α] [DecidableLT α] [Least? α] + [LawfulUpwardEnumerable α] [LawfulUpwardEnumerableLT α] + [LawfulUpwardEnumerableLeast? α] + [MeasurableSpaceMonad m] + +@[no_expose, default_instance, nolint unusedArguments] +instance instMeasurableSpaceForIn' [Finite (Rxo.Iterator α) Id] : + MeasurableSpaceForIn' m (Rio α) α inferInstance where + forIn' r init f := by + haveI := Iter.instMeasurableSpaceForIn' (ι := Rxo.Iterator α) (m := m) + refine MeasurableSpaceForIn'.forIn' (α := α) (Internal.iter r) init (fun a ha acc => f a ?_ acc) + simp only [Membership.mem] at ha + rwa [Internal.isPlausibleIndirectOutput_iter_iff] at ha + +/-- A range loop is the corresponding iterator loop. -/ +theorem forIn_eq_forIn_iter [Finite (Rxo.Iterator α) Id] + (r : Rio α) (init : β) (f : α → β → m (ForInStep β)) : + MeasureTheory.forIn r init f = MeasureTheory.forIn (Internal.iter r) init f := by + rfl + +end Rio + +namespace Rii + +open Std.Rii + +variable [UpwardEnumerable α] [Least? α] + [LawfulUpwardEnumerable α] [LawfulUpwardEnumerableLeast? α] + [MeasurableSpaceMonad m] + +@[no_expose, default_instance, nolint unusedArguments] +instance instMeasurableSpaceForIn' [Finite (Rxi.Iterator α) Id] : + MeasurableSpaceForIn' m (Rii α) α inferInstance where + forIn' r init f := by + haveI := Iter.instMeasurableSpaceForIn' (ι := Rxi.Iterator α) (m := m) + refine MeasurableSpaceForIn'.forIn' (α := α) (Internal.iter r) init (fun a ha acc => f a ?_ acc) + simp only [Membership.mem] at ha + rwa [Internal.isPlausibleIndirectOutput_iter_iff] at ha + +/-- A range loop is the corresponding iterator loop. -/ +theorem forIn_eq_forIn_iter [Finite (Rxi.Iterator α) Id] + (r : Rii α) (init : β) (f : α → β → m (ForInStep β)) : + MeasureTheory.forIn r init f = MeasureTheory.forIn (Internal.iter r) init f := by + rfl + +end Rii + +end MeasureTheory diff --git a/RandomDo/Monad/ForInInstances.lean b/RandomDo/Monad/ForInInstances.lean deleted file mode 100644 index fde7c94..0000000 --- a/RandomDo/Monad/ForInInstances.lean +++ /dev/null @@ -1,129 +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 - -public import RandomDo.Monad.Notation - -/-! -# Instances for `MeasurableSpaceForIn` - -**TODO** - --/ - -@[expose] public section - -section MeasurableSpace - -variable {α : Type*} [MeasurableSpace α] - -instance instMeasurableSpaceList : MeasurableSpace (List α) := - MeasurableSpace.comap List.equivSigmaTuple inferInstance - -instance instMeasurableSpaceArray : MeasurableSpace (Array α) := - MeasurableSpace.comap Array.toList inferInstance - -instance instMeasurableSpaceVector {n : ℕ} : MeasurableSpace (Vector α n) := - MeasurableSpace.comap Vector.toArray inferInstance - -instance instMeasurableSpaceSubarray : MeasurableSpace (Subarray α) := - MeasurableSpace.comap (fun s : Subarray α ↦ s.toList) inferInstance - -@[fun_prop] -lemma measurable_toList : Measurable (Array.toList : Array α → List α) := - Measurable.of_comap_le fun _ a ↦ a - -@[fun_prop] -lemma measurable_subarray_toList : Measurable (fun s : Subarray α ↦ s.toList) := - Measurable.of_comap_le fun _ a ↦ a - -@[fun_prop] -lemma measurable_toArray {n : ℕ} : Measurable (Vector.toArray : Vector α n → Array α) := - Measurable.of_comap_le fun _ a ↦ a - -end MeasurableSpace - -universe u v - -open MeasureTheory MeasurableSpacePure - -variable {α : Type u} [mα : MeasurableSpace α] {m : (α : Type u) → [MeasurableSpace α] → Type v} - -section Array - -/-- Compiler implementation for `forIn` -/ -@[inline] unsafe def Array.measurableSpaceForIn'Unsafe [MeasurableSpaceMonad m] - {β : Type u} [mβ : MeasurableSpace β] - (as : Array α) (b : β) (f : (a : α) → a ∈ as → β → m (ForInStep β)) : m β := - let sz := as.usize - let rec @[specialize, nolint docBlame] loop (i : USize) (b : β) : m β := rdo - if i < sz then - let a := as.uget i lcProof - match (← f a lcProof b) with - | ForInStep.done b => mPure b - | ForInStep.yield b => loop (i+1) b - else - mPure b - loop 0 b - -/-- Reference implementation for `forIn'` -/ -@[implemented_by Array.measurableSpaceForIn'Unsafe] -protected def Array.measurableSpaceForIn' [MeasurableSpaceMonad m] - {β : Type u} [mβ : MeasurableSpace β] - (as : Array α) (b : β) (f : (a : α) → a ∈ as → β → m (ForInStep β)) : m β := - let rec @[nolint docBlame] loop (i : Nat) (h : i ≤ as.size) (b : β) : m β := rdo - match i, h with - | 0, _ => mPure b - | i+1, h => - have h' : i < as.size := Nat.lt_of_lt_of_le (Nat.lt_succ_self i) h - have : as.size - 1 < as.size := Nat.sub_lt (Nat.zero_lt_of_lt h') (by decide) - have : as.size - 1 - i < as.size := Nat.lt_of_le_of_lt (Nat.sub_le (as.size - 1) i) this - match (← f as[as.size - 1 - i] (getElem_mem this) b) with - | ForInStep.done b => mPure b - | ForInStep.yield b => loop i (Nat.le_of_lt h') b - loop as.size (Nat.le_refl _) b - -instance [MeasurableSpaceMonad m] : MeasurableSpaceForIn' m (Array α) α inferInstance where - forIn' := Array.measurableSpaceForIn' - -instance {n : ℕ} [MeasurableSpaceMonad m] : - MeasurableSpaceForIn' m (Vector α n) α inferInstance where - forIn' xs b f := Array.measurableSpaceForIn' xs.toArray b (fun a h b => f a (by simpa using h) b) - -end Array - -section List - -variable {α β : Type*} [MeasurableSpace α] [MeasurableSpace β] [Ring α] - -/-- implementation for `forIn'` -/ -@[inline] -protected def List.measurableSpaceForIn' [MeasurableSpaceMonad m] - {β : Type u} [mβ : MeasurableSpace β] (as : @& List α) (init : β) - (f : (a : α) → a ∈ as → β → m (ForInStep β)) : m β := - let rec @[specialize, nolint docBlame] - loop : (as' : @& List α) → (b : β) → Exists (fun bs => bs ++ as' = as) → m β - | [], b, _ => mPure b - | a::as', b, h => rdo - have : a ∈ as := by - clear f - have ⟨bs, h⟩ := h - subst h - exact mem_append_right _ (Mem.head ..) - match (← f a this b) with - | ForInStep.done b => mPure b - | ForInStep.yield b => - have : Exists (fun bs => bs ++ as' = as) := - have ⟨bs, h⟩ := h - ⟨bs ++ [a], by rw [← h, append_cons (bs := as')]⟩ - loop as' b this - loop as init ⟨[], rfl⟩ - -@[simps] -instance [MeasurableSpaceMonad m] : MeasurableSpaceForIn' m (List α) α inferInstance where - forIn' := List.measurableSpaceForIn' - -end List diff --git a/RandomDo/Monad/Instances.lean b/RandomDo/Monad/Instances.lean index e716032..b1a27b7 100644 --- a/RandomDo/Monad/Instances.lean +++ b/RandomDo/Monad/Instances.lean @@ -6,6 +6,7 @@ Authors: David Ledvinka module public import RandomDo.MeasurableSpaceMonad.Defs +public import RandomDo.MeasurableSpaceMonad.ForInInstances.Iterators public import Mathlib.Control.Random public import Mathlib.MeasureTheory.Measure.GiryMonad public import Mathlib.Probability.ProductMeasure @@ -33,6 +34,11 @@ instance {m : Type u → Type v} [Monad m] : mPure := pure mBind := bind +instance (priority := 900) Monad.toMeasurableSpaceMonad.instMeasurableSpaceLiftBind + {n : Type u → Type w} {p : Type u → Type v} [Monad p] [MonadLiftT n p] : + MeasureTheory.MeasurableSpaceLiftBind n (Monad.toMeasurableSpaceMonad p) where + liftBind f x := monadLift x >>= f + /-- A measurable space monad for pseudo random number generation. -/ abbrev PseudoRandomM := Monad.toMeasurableSpaceMonad Rand diff --git a/RandomDo/Monad/Notation.lean b/RandomDo/Monad/Notation.lean index 9b1d5f9..cdc396d 100644 --- a/RandomDo/Monad/Notation.lean +++ b/RandomDo/Monad/Notation.lean @@ -6,6 +6,7 @@ Authors: David Ledvinka module public import RandomDo.MeasurableSpaceMonad.ForIn +public import RandomDo.ForMathlib.MeasureTheory.MeasurableSpace.Option public import Mathlib.MeasureTheory.MeasurableSpace.Prod public meta import Lean.Elab.BuiltinDo public meta import Lean.Meta.ProdN @@ -75,17 +76,16 @@ syntax (name := randKind) "rdo" doSeq : term section LoopElab -/-- parses `_ : _ in` for `rdo` for loops -/ +/-- A declaration in an `rdo` loop. Analogue of `do` notation's `Lean.Parser.Term.doForDecl`. -/ def rdoForDecl := leading_parser Lean.Parser.optional (atomic (Term.ident >> " : ")) >> termParser >> " in " >> withForbidden "rdo" termParser -/-- parser for `rdo` for loops -/ +/-- A `for` loop in `rdo`. Analogue of `do` notation's `Lean.Parser.Term.doFor`. -/ @[doElem_parser] def rdoFor := leading_parser "for " >> sepBy1 rdoForDecl ", " >> "rdo " >> doSeq -/-- Define expander for loops in `rdo` notation. Note this code mirrors core's implementation for -`do` notation as much as possible. -/ +/-- Expand loop binders. Analogue of `do` notation's `Lean.Elab.Do.expandDoFor`. -/ @[macro rdoFor] def expandRDoFor : Macro := fun stx => do match stx with | `(rdoFor| for $[$_ : ]? $_:ident in $_ rdo $_) => @@ -142,12 +142,26 @@ def rdoForDecl := leading_parser | some ($y, s') => $s:ident := s' do $body) - doElems := doElems.push (← `(doSeqItem| for%$tk $[$h? : ]? $x:ident in $xs rdo $body)) + doElems := doElems.push + (← `(doSeqItem| for%$tk $[$h? : ]? $x:ident in $xs rdo $body)) `(doElem| do $doElems*) | _ => Macro.throwUnsupported -/-- Define loop elaborator for `rdo` notation. Note this code mirrors core's implementation for -`do` notation as much as possible. -/ +/-- +Infer a loop's control effects. Analogue of `do` notation's `Lean.Elab.Do.InferControlInfo.ofElem` +case for `for` loops. +-/ +@[doElem_control_info rdoFor] def controlInfoRDoFor : ControlInfoHandler := fun stx => do + let `(rdoFor| for $[$_ : ]? $_:ident in $_ rdo $body) := stx | throwUnsupportedSyntax + let info ← inferControlInfoSeq body + return { info with -- keep only reassigns and returnsEarly + numRegularExits := 1, + continues := false, + breaks := false, + noFallthrough := false, + } + +/-- Elaborate a loop. Analogue of `do` notation's `Lean.Elab.Do.elabDoFor`. -/ @[doElem_elab rdoFor] def elabRDoFor : DoElab := fun stx dec => do let `(rdoFor| for%$tk $[$h? : ]? $x:ident in $xs rdo $body) := stx | throwUnsupportedSyntax let dec ← dec.ensureUnitAt tk @@ -179,17 +193,17 @@ def rdoForDecl := leading_parser | some e => mkSome oldReturnCont.resultType e defs := defs.push returnVar for x in loopMutVars do - let defn ← getLocalDeclFromUserName x.getId - Term.addTermInfo' x.ident defn.toExpr + Term.addTermInfo' x.ident (← getFVarFromUserName x.getId) + let v ← x.stateValue -- ForIn forces the mut tuple into the universe mi.u: that of the do block result type. -- If we don't do this, then we are stuck on solving constraints such as -- `max ?u.46 ?u.47 =?= max (max ?u.22 ?u.46) ?u.47` -- It's important we do this as a separate isLevelDefEq check on the decremented level because -- otherwise (`ensureHasType (mkSort mi.u.succ)`) we are stuck on constraints like -- `max (?u+1) (?v+1) =?= ?u+1` - let u ← getDecLevel defn.type + let u ← getDecLevel (← inferType v) discard <| isLevelDefEq u mi.u - defs := defs.push defn.toExpr + defs := defs.push v if info.returnsEarly && loopMutVars.isEmpty then defs := defs.push (mkConst ``Unit.unit) return defs @@ -198,16 +212,16 @@ def rdoForDecl := leading_parser let (app, p?) ← match h? with | none => let instForIn ← Term.mkInstMVar <| - mkApp3 (mkConst ``MeasurableSpaceForIn [uρ, uα, mi.u, mi.v]) mi.m ρ α + mkApp3 (mkConst ``MeasurableSpaceForIn [uρ, uα, mi.u, mi.v]) mi.m ρ α let app := Lean.mkConst ``MeasurableSpaceForIn.forIn [uρ, uα, mi.u, mi.v] - let app := mkApp8 app mi.m ρ α instForIn σ mσ xs preS + let app := mkApp8 app mi.m ρ α instForIn σ mσ xs preS -- 1 arg remaining: loop body pure (app, none) | some _ => let d ← mkFreshExprMVar (mkApp2 (mkConst ``Membership [uα, uρ]) α ρ) (userName := `d) let instForIn ← Term.mkInstMVar <| - mkApp4 (mkConst ``MeasurableSpaceForIn' [uρ, uα, mi.u, mi.v]) mi.m ρ α d + mkApp4 (mkConst ``MeasurableSpaceForIn' [uρ, uα, mi.u, mi.v]) mi.m ρ α d let app := Lean.mkConst ``MeasurableSpaceForIn'.forIn' [uρ, uα, mi.u, mi.v] - let app := mkApp9 app mi.m ρ α d instForIn σ mσ xs preS + let app := mkApp9 app mi.m ρ α d instForIn σ mσ xs preS -- 1 arg remaining: loop body pure (app, some d) let s ← mkFreshUserName `__s let xh : Array (Name × (Array Expr → DoElabM Expr)) := match h?, p? with @@ -250,9 +264,8 @@ def rdoForDecl := leading_parser if info.returnsEarly then let ret ← getFVarFromUserName returnVarName let ret ← if loopMutVars.isEmpty then mkAppM ``Prod.fst #[ret] else pure ret - /- Both branches produce an `m γ`, so the motive is constant and the matcher generated - for `Break.runK` is not needed: `Break.runK` itself is that case analysis. Note that it - takes the `none` branch first, the other way round from the matcher. -/ + -- Both branches produce `m γ`, so the motive is constant. Use the public + -- `Break.runK` operation, whose `none` branch precedes its `some` branch. let none := mkSimpleThunk (← dec.continueWithUnit) let some ← withLocalDeclD (← mkFreshUserName `r) oldReturnCont.resultType fun r => do mkLambdaFVars #[r] (← oldReturnCont.k r) diff --git a/RandomDo/Tactic/IsMarkov/Lemmas.lean b/RandomDo/Tactic/IsMarkov/Lemmas.lean index 38cc0d0..3d47e67 100644 --- a/RandomDo/Tactic/IsMarkov/Lemmas.lean +++ b/RandomDo/Tactic/IsMarkov/Lemmas.lean @@ -6,7 +6,7 @@ Authors: Gaëtan Serré module public import RandomDo.Monad.Instances -public import RandomDo.Monad.ForInInstances +public import RandomDo.MeasurableSpaceMonad.ForInInstances.Basic public import RandomDo.Measurable public import RandomDo.Tactic.IsMarkov.ForInStep public import Mathlib.MeasureTheory.Measure.ProbabilityMeasure diff --git a/Test.lean b/Test.lean index bba557b..5d7e83d 100644 --- a/Test.lean +++ b/Test.lean @@ -4,6 +4,8 @@ public import Test.Bind public import Test.Common public import Test.Computable public import Test.Control +public import Test.ForInElab +public import Test.ForInInstances public import Test.Gaps public import Test.Instances public import Test.IsMarkov diff --git a/Test/ForInElab.lean b/Test/ForInElab.lean new file mode 100644 index 0000000..2a081dd --- /dev/null +++ b/Test/ForInElab.lean @@ -0,0 +1,137 @@ +/- +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.Monad.Notation +-- Executable checks also need interpreter access to the imported declarations. +meta import RandomDo.Monad.Notation + +/-! +# State and control flow in `rdo` loops + +Abstract translation checks cover the iteration interfaces without assuming measurable +collections or elements. Executable checks use a test-only state monad and core's list iterator: +the log makes dropped, duplicated, or reordered actions observable. These are tests of the +elaborator, not implementations of the production collection instances. +-/ + +-- This module intentionally contains private fixtures and executable test assertions. +set_option linter.privateModule false +set_option linter.hashCommand false + +open MeasureTheory MeasurableSpacePure MeasurableSpaceBind + +universe u v uρ uι + +section Abstract + +variable {m : (α : Type u) → [MeasurableSpace α] → Type v} [MeasurableSpaceMonad m] + {ρ : Type uρ} {ι : Type uι} {α : Type u} [MeasurableSpace α] + +-- Only the loop state needs a measurable-space instance; collection and element universes +-- need not coincide with the state or computation universe. +example [MeasurableSpaceForIn m ρ ι] (xs : ρ) (init : α) (step : ι → α → m α) : + (rdo + let mut state := init + for a in xs rdo + state ← step a state + return state : m α) = + (MeasureTheory.forIn xs init fun a state => + step a state >>=ₘ fun state => mPure (ForInStep.yield state)) >>=ₘ mPure := rfl + +-- A proof binder must pass the actual membership proof to the body. +example [Membership ι ρ] [MeasurableSpaceForIn' m ρ ι inferInstance] + (xs : ρ) (init : α) (step : (a : ι) → a ∈ xs → α → m α) : + (rdo + let mut state := init + for h : a in xs rdo + state ← step a h state + return state : m α) = + (MeasureTheory.forIn' xs init fun a h state => + step a h state >>=ₘ fun state => mPure (ForInStep.yield state)) >>=ₘ mPure := rfl + +end Abstract + +namespace Test.ForInElab + +private def TraceM (α : Type) [_mα : MeasurableSpace α] := StateM (List Nat) α + +local instance : MeasurableSpaceMonad TraceM where + mPure := pure + mBind := bind + +local instance {α : Type} : MeasurableSpaceForIn' TraceM (List α) α inferInstance where + forIn' := _root_.ForIn'.forIn' (m := StateM (List Nat)) + +private def run {α : Type} [MeasurableSpace α] (action : TraceM α) : α × List Nat := + Id.run (action []) + +private def record (n : Nat) : TraceM Unit := fun log => ((), log ++ [n]) + +-- `continue` and `break` must preserve updates to both mutable variables. The log records +-- which iterations ran, including the skipped value and the final value before `break`. +private def control (xs : List Nat) : TraceM (Nat × Nat) := rdo + let mut visits := 0 + let mut total := 0 + for n in xs rdo + record n + visits := visits + 1 + if n = 0 then continue + total := total + n + if n = 3 then break + return (visits, total) + +#guard run (control [1, 0, 3, 7]) == ((3, 4), [1, 0, 3]) +#guard run (control [1, 2]) == ((2, 3), [1, 2]) + +-- Early return must preserve the updated state and skip both remaining iterations and the +-- continuation after the loop. Normal completion must run that continuation. +private def returnWithState (xs : List Nat) : TraceM Nat := rdo + let mut total := 0 + for n in xs rdo + record n + total := total + n + if total ≥ 4 then return total + record 99 + return total + +#guard run (returnWithState [1, 3, 5]) == (4, [1, 3]) +#guard run (returnWithState [1, 2]) == (3, [1, 2, 99]) + +-- The elaborator uses a different state tuple when an early return carries no mutable variables. +private def returnWithoutState (xs : List Nat) : TraceM Nat := rdo + for n in xs rdo + record n + if n = 3 then return n + return 0 + +#guard run (returnWithoutState [1, 3, 5]) == (3, [1, 3]) +#guard run (returnWithoutState [1, 2]) == (0, [1, 2]) + +-- Inner loop jumps must not jump out of the outer loop. This also exercises the custom +-- `ControlInfo` handler, whose absence previously prevented nested loops from elaborating. +private def nested : TraceM Unit := rdo + for a in [1, 2] rdo + for b in [0, 1, 2, 3] rdo + if b = 0 then continue + if b = 2 then break + record (10 * a + b) + record (10 * a) + +#guard run nested == ((), [11, 10, 21, 20]) + +-- Parallel iteration stores the second stream in the loop state. The particular σ-algebra is +-- irrelevant here: this test is about macro expansion and stopping at the shorter collection. +local instance : MeasurableSpace (List Nat) := ⊤ + +private def parallel (xs : List (Nat × Nat)) (ys : List Nat) : TraceM Unit := rdo + for (a, b) in xs, c in ys rdo + record (a + b + c) + +#guard run (parallel [(1, 2), (3, 4), (5, 6)] [10, 20]) == ((), [13, 27]) +#guard run (parallel [(1, 2)] [10, 20]) == ((), [13]) + +end Test.ForInElab diff --git a/Test/ForInInstances.lean b/Test/ForInInstances.lean new file mode 100644 index 0000000..f5a1f65 --- /dev/null +++ b/Test/ForInInstances.lean @@ -0,0 +1,260 @@ +/- +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.MeasurableSpaceMonad.ForInInstances.Basic +import RandomDo.Monad.Instances +import RandomDo.Monad.Notation +import Std.Data.Iterators +import all Init.Data.Range.Polymorphic.RangeIterator +import all Std.Data.Iterators.Producers.Range +meta import RandomDo.MeasurableSpaceMonad.ForInInstances.Basic +meta import RandomDo.Monad.Instances +meta import RandomDo.Monad.Notation +meta import Std.Data.Iterators + +/-! +# Collection and iterator loops in the supported measurable space monads + +The abstract checks keep the element and state universes independent and require no measurable +space on the elements. Executable checks use an adapted state monad to observe iteration order +and early termination, including the array's compiler implementation. A kernel check covers its +reference implementation, and measure programs check that loop-body effects are retained. Iterator +checks cover lazy pipelines, producer effects, membership proofs, and termination wrappers. Direct +range checks cover endpoint conventions, bounded-type exhaustion, and inference from typed binders. +-/ + +-- This module contains private fixtures and executable regression checks. +set_option linter.privateModule false +set_option linter.hashCommand false + +open MeasureTheory MeasurableSpacePure MeasurableSpaceBind + +universe uα u v + +section Inference + +variable {α : Type uα} {m : (α : Type u) → [MeasurableSpace α] → Type v} + [MeasurableSpaceMonad m] + +example : MeasurableSpaceForIn' m (List α) α inferInstance := inferInstance +example : MeasurableSpaceForIn' m (Array α) α inferInstance := inferInstance +example {n : Nat} : MeasurableSpaceForIn' m (Vector α n) α inferInstance := inferInstance + +end Inference + +namespace Test.ForInInstances + +-- These fixtures use main's ordinary-monad adapter; no sampler refactor is required. +private abbrev TraceT (m : Type → Type) := Monad.toMeasurableSpaceMonad (RandGT Nat m) + +private abbrev Trace := TraceT Id + +-- Record visited elements as decimal digits in the sampler state. This makes both their order +-- and any actions performed after the requested stop observable. +private def record (n : Nat) : Trace Unit := fun s => ((), ⟨10 * s.down + n⟩) + +private def visit {ρ : Type} [MeasurableSpaceForIn (Trace) ρ Nat] + (xs : ρ) : Trace Nat := rdo + let mut total := 0 + for n in xs rdo + record n + total := total + n + if n = 3 then break + return total + +private def run (action : Trace Nat) (s : Nat := 0) : Nat × Nat := + let (a, s') := action ⟨s⟩ + (a, s'.down) + +-- Kernel reduction checks the safe array reference implementation. +example : run (visit #[1, 3, 7]) = (4, 13) := rfl + +-- Executable checks use the array implementation selected by `implemented_by`. +#guard run (visit [1, 3, 7]) == (4, 13) +#guard run (visit #[1, 3, 7]) == (4, 13) +#guard run (visit (#[1, 3, 7].toVector)) == (4, 13) +#guard run (visit (#[] : Array Nat)) 42 == (0, 42) +#guard run (visit #[1, 2]) == (3, 12) +#guard run (visit [1, 2]) == (3, 12) + +-- A zero-measure action inside a nonempty loop must annihilate the result measure. +example : + (rdo + let mut total := 0 + for n in [1, 2] rdo + let a ← (0 : Measure Nat) + total := total + n + a + return total : Measure Nat) = 0 := by + simp [MeasureTheory.forIn, MeasureTheory.forIn', + MeasureTheory.List.measurableSpaceForIn', MeasureTheory.List.measurableSpaceForIn'.loop, + mPure, mBind] + +-- Pipelines must stay lazy: filtering skips elements, take stops the producer, and break can +-- terminate an unbounded iterator. The sampler state records precisely the visited elements. +#guard run (visit ([0, 1, 2, 3, 4, 5].iter.map (· + 1) |>.filter (· % 2 == 0) |>.take 2)) == + (6, 24) +#guard run (visit (([1, 2, 9].iter.zip #[0, 0].iter).map fun p => p.1 + p.2)) == (3, 12) +#guard run (visit ((0...* : Std.Rci Nat).iter)) == (6, 123) +#guard run (visit [1, 3, 7].iter.ensureTermination) == (4, 13) + +-- Iterator internals and elements need no measurable-space instance. In particular, the +-- elements can inhabit a larger universe than the sampler's return values. +example {m : (α : Type) → [MeasurableSpace α] → Type} [MeasurableSpaceMonad m] + (xs : List (Type)) : m Nat := rdo + let mut count := 0 + for _ in xs.iter rdo + count := count + 1 + return count + +-- Producer effects and body effects must interleave. Breaking at 3 must prevent the producer +-- from processing 7. The expected log is producer(1), body(1), producer(3), body(3). +private def produce (n : Nat) : RandGT Nat Id Nat := do + record (2 * n) + return n + +#guard run (visit ([1, 3, 7].iter.mapM produce)) == (4, 2163) +#guard run (visit ([1, 3, 7].iter.mapM produce).ensureTermination) == (4, 2163) + +-- Producer effects can also lift from the base monad of a transformer. An absent producer +-- value aborts the computation, even though the loop body itself never fails. +private def fromBase (xs : List (Option Nat)) : TraceT Option Nat := rdo + let mut total := 0 + for n in xs.iter.mapM id rdo + total := total + n + return total + +#guard (fromBase [some 1, some 2]).run ⟨42⟩ == some (3, ⟨42⟩) +#guard (fromBase [some 1, none, some 2]).run ⟨42⟩ == none + +-- Check the compatibility wrappers without using the deprecated allowNontermination function. +#guard run (visit (Std.Iter.Partial.mk [1, 3, 7].iter)) == (4, 13) +#guard run (visit (Std.IterM.Partial.mk ([1, 3, 7].iter.mapM produce))) == (4, 2163) + +-- As in core, membership-aware iterator loops are enabled explicitly. +private def withIteratorProof : Trace Nat := + let it := [1, 2].iter + letI := MeasureTheory.Iter.instMeasurableSpaceForIn' (m := Trace) + (ι := Std.Iterators.Types.ListIterator Nat) + rdo + let mut total := 0 + for h : n in it rdo + have : it.IsPlausibleIndirectOutput n := h + record n + total := total + n + return total + +#guard run withIteratorProof == (3, 12) + +-- The unfolding theorem supports kernel proofs about Measure, not just executable samplers. +example : MeasureTheory.forIn ([] : List Nat).iter 17 + (fun _ _ => (0 : Measure (ForInStep Nat))) = Measure.dirac 17 := by + rw [MeasureTheory.Iter.forIn_eq_match_step] + simp [mPure] + +-- One polymorphic program must support both the executable and measure interpretations. +-- The sampler argument affects the sum, and breaking must prevent subsequent sampling. +private def sampledSum {m : (α : Type) → [MeasurableSpace α] → Type} + [MeasurableSpaceMonad m] (sample : Nat → m Nat) (xs : List Nat) : m Nat := rdo + let mut total := 0 + for n in xs.iter rdo + let a ← sample n + total := total + a + if n = 3 then break + return total + +private def sample (n : Nat) : Trace Nat := rdo + record n + return n + +#guard run (sampledSum sample [1, 3, 7]) == (4, 13) +#guard run (sampledSum sample []) 42 == (0, 42) + +example : sampledSum (m := Measure) (fun _ => 0) [1, 3, 7] = 0 := by + unfold sampledSum + dsimp only + rw [MeasureTheory.Iter.forIn_eq_match_step] + simp [mBind, mPure] + +-- The closed-range specialization must exhaust a bounded type without wrapping around. +private def boundaryRange {m : (α : Type) → [MeasurableSpace α] → Type} + [MeasurableSpaceMonad m] (sample : UInt8 → m Nat) : m Nat := rdo + let mut total := 0 + for n in (253 : UInt8)...=255 rdo + let a ← sample n + total := total + a + return total + +#guard run (boundaryRange (fun n => rdo record 1; return n.toNat)) == (762, 111) + +-- This proof uses the specialized range's lawfulness instance, not just the default consumer. +example : boundaryRange (m := Measure) (fun _ => 0) = 0 := by + unfold boundaryRange + dsimp only + rw [Rcc.forIn_eq_forIn_iter, MeasureTheory.Iter.forIn_eq_match_step] + simp [Std.Rxc.Iterator.step_eq_step, Std.Rxc.Iterator.step, Std.Rcc.Internal.iter, mBind, mPure] + +-- Each direct range form must preserve its endpoint convention and stop at type exhaustion. +private def rangeElements {ρ : Type} [MeasurableSpaceForIn (Trace) ρ UInt8] + (r : ρ) : Array Nat := + ((rdo + let mut elements := #[] + for n in r rdo + elements := elements.push n.toNat + return elements : Trace (Array Nat)) ⟨0⟩).1 + +#guard rangeElements ((2 : UInt8)...=4) == #[2, 3, 4] +#guard rangeElements ((2 : UInt8)...4) == #[2, 3] +#guard rangeElements ((253 : UInt8)...*) == #[253, 254, 255] +#guard rangeElements ((2 : UInt8)<...=4) == #[3, 4] +#guard rangeElements ((2 : UInt8)<...4) == #[3] +#guard rangeElements ((253 : UInt8)<...*) == #[254, 255] +#guard rangeElements (*...=(2 : UInt8)) == #[0, 1, 2] +#guard rangeElements (*...(2 : UInt8)) == #[0, 1] +#guard rangeElements (*...* : Std.Rii UInt8) == Array.ofFn (fun i : Fin 256 => i.val) +#guard rangeElements ((255 : UInt8)<...*) == #[] +#guard rangeElements ((4 : UInt8)...2) == #[] + +-- The element type inferred from the body must determine a range with literal or absent bounds. +-- Known notation limitation: `for (i : Int) in 1...3 rdo ...` still defaults the bounds to +-- Nat during pattern expansion, unlike core. Annotating a bound or using the body as below works. +-- Keep that elaborator issue separate from the range instances. +#guard run (rdo + let mut total := 0 + for i in 1...3 rdo + total := total + (i : Int).toNat + for i in *...* rdo + total := total + (i : Fin 3).val + return total) == (6, 0) + +-- Range membership must supply a usable bounds proof, so indexing needs no run-time check. +private def indexedSum {m : (α : Type) → [MeasurableSpace α] → Type} + [MeasurableSpaceMonad m] (xs : Array Nat) : m Nat := rdo + let mut total := 0 + for h : i in 0...xs.size rdo + total := total + xs[i]'h.2 + return total + +#guard run (indexedSum #[2, 4, 8]) == (14, 0) + +-- Direct ranges also preserve early exit and skip effects after continue or return. +#guard run (visit (1...8 : Std.Rco Nat)) == (6, 123) +#guard run (rdo + let mut total := 0 + for i in 0...5 rdo + if i = 0 then continue + record i + if i = 2 then return total + total := total + i + return total) == (1, 12) + +-- As in core, a potentially infinite Nat range needs explicit iterator syntax instead. +example : True := by + fail_if_success + have : MeasurableSpaceForIn (Trace) (Std.Rci Nat) Nat := inferInstance + trivial + +end Test.ForInInstances diff --git a/Test/Gaps.lean b/Test/Gaps.lean index 39aeb1c..70ab24e 100644 --- a/Test/Gaps.lean +++ b/Test/Gaps.lean @@ -18,26 +18,6 @@ open MeasureTheory ProbabilityTheory namespace Test.Gaps -/-! ## Nested loops - -TODO: register a `ControlInfo` inference handler for `RDo.rdoFor`, mirroring the rule core states -inline for `doFor` in `Lean/Elab/Do/InferControlInfo.lean`. --/ - -/-- -error: No `ControlInfo` inference handler found for `RDo.rdoFor` in syntax - for y in ys rdo - s := s + x * y -Register a handler with `@[doElem_control_info RDo.rdoFor]`. --/ -#guard_msgs (whitespace := lax) in -def nestedLoops (xs ys : List ℕ) : IdM ℕ := rdo - let mut s := 0 - for x in xs rdo - for y in ys rdo - s := s + x * y - return s - /-! ## Unbounded and conditional iteration TODO: `while`, `repeat` and `repeat … until` all expand to `for _ in Loop.mk do …`, which reaches @@ -79,8 +59,8 @@ noncomputable def tryCatch : Measure Bool := rdo /-! ## Collections with no `MeasurableSpaceForIn` instance -TODO: instances exist for `List`, `Array` and `Vector` only. A range is the one most missed, since -`for i in [0:n]` has to be written `for i in List.range n` today. +Polymorphic ranges (`0...n`) and iterators are supported. The legacy range syntax (`[0:n]`) +and direct `Finset` iteration still need separate instances. -/ /--