feat: lemmas about iterator collectors

src/Std/Data/Iterators.lean

@@ -8,6 +8,7 @@ import Std.Data.Iterators.Basic
 import Std.Data.Iterators.Producers
 import Std.Data.Iterators.Consumers
 import Std.Data.Iterators.Internal
+import Std.Data.Iterators.Lemmas
 
 /-!
 # Iterators
@@ -92,7 +93,7 @@ All of the following module names are prefixed with `Std.Data.Iterators`.
 
 ### Verification API
 
-`Lemmas` will provide the means to verify programs that use iterators.
+`Lemmas` provides the means to verify programs that use iterators.
 
 ### Implementation details
 

src/Std/Data/Iterators/Basic.lean

@@ -19,16 +19,6 @@ namespace Std
 
 namespace Iterators
 
-/--
-`BaseIter` is the common data structure underlying `Iter` and `IterM`. API users should never
-use `BaseIter` directly, only `Iter` and `IterM`.
--/
-structure BaseIter {α : Type w} (m : Type w → Type w') (β : Type w) : Type w where
-  /--
-  Internal implementation detail of the iterator.
-  -/
-  internalState : α
-
 /--
 An iterator that sequentially emits values of type `β` in the monad `m`. It may be finite
 or infinite.
@@ -68,7 +58,9 @@ def x := [1, 2, 3].iterM IO
 def x := ([1, 2, 3].iterM IO : IterM IO Nat)
 ```
 -/
-def IterM {α : Type w} (m : Type w → Type w') (β : Type w) := BaseIter (α := α) m β
+structure IterM {α : Type w} (m : Type w → Type w') (β : Type w) where
+  /-- Internal implementation detail of the iterator. -/
+  internalState : α
 
 /--
 An iterator that sequentially emits values of type `β`. It may be finite
@@ -109,29 +101,41 @@ def x := [1, 2, 3].iter
 def x := ([1, 2, 3].iter : Iter Nat)
 ```
 -/
-def Iter {α : Type w} (β : Type w) := BaseIter (α := α) Id β
+structure Iter {α : Type w} (β : Type w) where
+  /-- Internal implementation detail of the iterator. -/
+  internalState : α
 
 /--
 Converts a pure iterator (`Iter β`) into a monadic iterator (`IterM Id β`) in the
 identity monad `Id`.
 -/
 def Iter.toIterM {α : Type w} {β : Type w} (it : Iter (α := α) β) : IterM (α := α) Id β :=
-  it
+  ⟨it.internalState⟩
 
 /--
 Converts a monadic iterator (`IterM Id β`) over `Id` into a pure iterator (`Iter β`).
 -/
-def IterM.toPureIter {α : Type w} {β : Type w} (it : IterM (α := α) Id β) : Iter (α := α) β :=
-  it
+def IterM.toIter {α : Type w} {β : Type w} (it : IterM (α := α) Id β) : Iter (α := α) β :=
+  ⟨it.internalState⟩
+
+@[simp]
+theorem Iter.toIter_toIterM {α : Type w} {β : Type w} (it : Iter (α := α) β) :
+    it.toIterM.toIter = it :=
+  rfl
 
 @[simp]
-theorem Iter.toPureIter_toIterM {α : Type w} {β : Type w} (it : Iter (α := α) β) :
-    it.toIterM.toPureIter = it :=
+theorem Iter.toIter_comp_toIterM {α : Type w} {β : Type w} :
+    IterM.toIter ∘ Iter.toIterM (α := α) (β := β) = id :=
   rfl
 
 @[simp]
-theorem Iter.toIterM_toPureIter {α : Type w} {β : Type w} (it : IterM (α := α) Id β) :
-    it.toPureIter.toIterM = it :=
+theorem Iter.toIterM_toIter {α : Type w} {β : Type w} (it : IterM (α := α) Id β) :
+    it.toIter.toIterM = it :=
+  rfl
+
+@[simp]
+theorem Iter.toIterM_comp_toIter {α : Type w} {β : Type w} :
+    Iter.toIterM ∘ IterM.toIter (α := α) (β := β) = id :=
   rfl
 
 section IterStep
@@ -169,6 +173,27 @@ def IterStep.successor : IterStep α β → Option α
   | .skip it => some it
   | .done => none
 
+/--
+If present, applies `f` to the iterator of an `IterStep` and replaces the iterator
+with the result of the application of `f`.
+-/
+@[always_inline, inline]
+def IterStep.mapIterator {α' : Type u'} (f : α → α') : IterStep α β → IterStep α' β
+  | .yield it out => .yield (f it) out
+  | .skip it => .skip (f it)
+  | .done => .done
+
+@[simp]
+theorem IterStep.mapIterator_mapIterator {α' : Type u'} {α'' : Type u''}
+    {f : α → α'} {g : α' → α''} {step : IterStep α β} :
+    (step.mapIterator f).mapIterator g = step.mapIterator (g ∘ f) := by
+  cases step <;> rfl
+
+@[simp]
+theorem IterStep.mapIterator_id {step : IterStep α β} :
+    step.mapIterator id = step := by
+  cases step <;> rfl
+
 /--
 A variant of `IterStep` that bundles the step together with a proof that it is "plausible".
 The plausibility predicate will later be chosen to assert that a state is a plausible successor
@@ -203,6 +228,20 @@ def PlausibleIterStep.done {IsPlausibleStep : IterStep α β → Prop}
     (h : IsPlausibleStep .done) : PlausibleIterStep IsPlausibleStep :=
   ⟨.done, h⟩
 
+/--
+A more convenient `cases` eliminator for `PlausibleIterStep`.
+-/
+@[elab_as_elim, cases_eliminator]
+abbrev PlausibleIterStep.casesOn {IsPlausibleStep : IterStep α β → Prop}
+    {motive : PlausibleIterStep IsPlausibleStep → Sort x} (s : PlausibleIterStep IsPlausibleStep)
+    (yield : ∀ it' out h, motive ⟨.yield it' out, h⟩)
+    (skip : ∀ it' h, motive ⟨.skip it', h⟩)
+    (done : ∀ h, motive ⟨.done, h⟩) : motive s :=
+  match s with
+  | .yield it' out h => yield it' out h
+  | .skip it' h => skip it' h
+  | .done h => done h
+
 end IterStep
 
 /--
@@ -296,7 +335,7 @@ is up to the `Iterator` instance but it should be strong enough to allow termina
 -/
 def Iter.IsPlausibleStep {α : Type w} {β : Type w} [Iterator α Id β]
     (it : Iter (α := α) β) (step : IterStep (Iter (α := α) β) β) : Prop :=
-  it.toIterM.IsPlausibleStep step
+  it.toIterM.IsPlausibleStep (step.mapIterator Iter.toIterM)
 
 /--
 The type of the step object returned by `Iter.step`, containing an `IterStep`
@@ -305,6 +344,37 @@ and a proof that this is a plausible step for the given iterator.
 def Iter.Step {α : Type w} {β : Type w} [Iterator α Id β] (it : Iter (α := α) β) :=
   PlausibleIterStep (Iter.IsPlausibleStep it)
 
+/--
+Converts an `Iter.Step` into an `IterM.Step`.
+-/
+@[always_inline, inline]
+def Iter.Step.toMonadic {α : Type w} {β : Type w} [Iterator α Id β] {it : Iter (α := α) β}
+    (step : it.Step) : it.toIterM.Step :=
+  ⟨step.val.mapIterator Iter.toIterM, step.property⟩
+
+/--
+Converts an `IterM.Step` into an `Iter.Step`.
+-/
+@[always_inline, inline]
+def IterM.Step.toPure {α : Type w} {β : Type w} [Iterator α Id β] {it : IterM (α := α) Id β}
+    (step : it.Step) : it.toIter.Step :=
+  ⟨step.val.mapIterator IterM.toIter, (by simp [Iter.IsPlausibleStep, step.property])⟩
+
+@[simp]
+theorem IterM.Step.toPure_yield {α β : Type w} [Iterator α Id β] {it : IterM (α := α) Id β}
+    {it' out h} : IterM.Step.toPure (⟨.yield it' out, h⟩ : it.Step) = .yield it'.toIter out h :=
+  rfl
+
+@[simp]
+theorem IterM.Step.toPure_skip {α β : Type w} [Iterator α Id β] {it : IterM (α := α) Id β}
+    {it' h} : IterM.Step.toPure (⟨.skip it', h⟩ : it.Step) = .skip it'.toIter h :=
+  rfl
+
+@[simp]
+theorem IterM.Step.toPure_done {α β : Type w} [Iterator α Id β] {it : IterM (α := α) Id β}
+    {h} : IterM.Step.toPure (⟨.done, h⟩ : it.Step) = .done h :=
+  rfl
+
 /--
 Asserts that a certain output value could plausibly be emitted by the given iterator in its next
 step.
@@ -319,15 +389,15 @@ given iterator `it`.
 -/
 def Iter.IsPlausibleSuccessorOf {α : Type w} {β : Type w} [Iterator α Id β]
     (it' it : Iter (α := α) β) : Prop :=
-  it'.toIterM.IsPlausibleSuccessorOf it
+  it'.toIterM.IsPlausibleSuccessorOf it.toIterM
 
 /--
 Asserts that a certain iterator `it'` could plausibly be the directly succeeding iterator of another
 given iterator `it` while no value is emitted (see `IterStep.skip`).
 -/
 def Iter.IsPlausibleSkipSuccessorOf {α : Type w} {β : Type w} [Iterator α Id β]
     (it' it : Iter (α := α) β) : Prop :=
-  it'.toIterM.IsPlausibleSkipSuccessorOf it
+  it'.toIterM.IsPlausibleSkipSuccessorOf it.toIterM
 
 /--
 Makes a single step with the given iterator `it`, potentially emitting a value and providing a
@@ -336,7 +406,7 @@ the termination measures `it.finitelyManySteps` and `it.finitelyManySkips`.
 -/
 @[always_inline, inline]
 def Iter.step {α β : Type w} [Iterator α Id β] (it : Iter (α := α) β) : it.Step :=
-  it.toIterM.step
+  it.toIterM.step.run.toPure
 
 end Pure
 
@@ -415,14 +485,14 @@ with `IterM.finitelyManySteps`.
 theorem Iter.TerminationMeasures.Finite.rel_of_yield
     {α : Type w} {β : Type w} [Iterator α Id β]
     {it it' : Iter (α := α) β} {out : β} (h : it.IsPlausibleStep (.yield it' out)) :
-    IterM.TerminationMeasures.Finite.Rel ⟨it'⟩ ⟨it⟩ :=
+    IterM.TerminationMeasures.Finite.Rel ⟨it'.toIterM⟩ ⟨it.toIterM⟩ :=
   IterM.TerminationMeasures.Finite.rel_of_yield h
 
 @[inherit_doc Iter.TerminationMeasures.Finite.rel_of_yield]
 theorem Iter.TerminationMeasures.Finite.rel_of_skip
     {α : Type w} {β : Type w} [Iterator α Id β]
     {it it' : Iter (α := α) β} (h : it.IsPlausibleStep (.skip it')) :
-    IterM.TerminationMeasures.Finite.Rel ⟨it'⟩ ⟨it⟩ :=
+    IterM.TerminationMeasures.Finite.Rel ⟨it'.toIterM⟩ ⟨it.toIterM⟩ :=
   IterM.TerminationMeasures.Finite.rel_of_skip h
 
 macro_rules | `(tactic| decreasing_trivial) => `(tactic|

src/Std/Data/Iterators/Consumers.lean

@@ -5,6 +5,7 @@ Authors: Paul Reichert
 -/
 prelude
 import Std.Data.Iterators.Consumers.Monadic
+import Std.Data.Iterators.Consumers.Access
 import Std.Data.Iterators.Consumers.Collect
 import Std.Data.Iterators.Consumers.Loop
 import Std.Data.Iterators.Consumers.Partial

src/Std/Data/Iterators/Consumers/Access.lean

@@ -0,0 +1,52 @@
+/-
+Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Paul Reichert
+-/
+prelude
+import Std.Data.Iterators.Consumers.Partial
+
+namespace Std.Iterators
+
+/--
+If possible, takes `n` steps with the iterator `it` and
+returns the `n`-th emitted value, or `none` if `it` finished
+before emitting `n` values.
+
+This function requires a `Productive` instance proving that the iterator will always emit a value
+after a finite number of skips. If the iterator is not productive or such an instance is not
+available, consider using `it.allowNontermination.atIdxSlow?` instead of `it.atIdxSlow?`. However,
+it is not possible to formally verify the behavior of the partial variant.
+-/
+@[specialize]
+def Iter.atIdxSlow? {α β} [Iterator α Id β] [Productive α Id]
+    (n : Nat) (it : Iter (α := α) β) : Option β :=
+  match it.step with
+  | .yield it' out _ =>
+    match n with
+    | 0 => some out
+    | k + 1 => it'.atIdxSlow? k
+  | .skip it' _ => it'.atIdxSlow? n
+  | .done _ => none
+termination_by (n, it.finitelyManySkips)
+
+/--
+If possible, takes `n` steps with the iterator `it` and
+returns the `n`-th emitted value, or `none` if `it` finished
+before emitting `n` values.
+
+This is a partial, potentially nonterminating, function. It is not possible to formally verify
+its behavior. If the iterator has a `Productive` instance, consider using `Iter.atIdxSlow?` instead.
+-/
+@[specialize]
+partial def Iter.Partial.atIdxSlow? {α β} [Iterator α Id β] [Monad Id]
+    (n : Nat) (it : Iter.Partial (α := α) β) : Option β := do
+  match it.it.step with
+  | .yield it' out _ =>
+    match n with
+    | 0 => some out
+    | k + 1 => (⟨it'⟩ : Iter.Partial (α := α) β).atIdxSlow? k
+  | .skip it' _ => (⟨it'⟩ : Iter.Partial (α := α) β).atIdxSlow? n
+  | .done _ => none
+
+end Std.Iterators

src/Std/Data/Iterators/Consumers/Monadic/Collect.lean

@@ -91,6 +91,12 @@ class LawfulIteratorCollect (α : Type w) (m : Type w → Type w') [Monad m] [It
     [i : IteratorCollect α m] where
   lawful : i = .defaultImplementation
 
+theorem LawfulIteratorCollect.toArrayMapped_eq {α β γ : Type w} {m : Type w → Type w'} [Monad m]
+    [Iterator α m β] [Finite α m] [IteratorCollect α m] [hl : LawfulIteratorCollect α m]
+    {f : β → m γ} {it : IterM (α := α) m β} :
+    IteratorCollect.toArrayMapped f it = IterM.DefaultConsumers.toArrayMapped f it := by
+  cases hl.lawful; rfl
+
 instance (α : Type w) (m : Type w → Type w') [Monad m] [Iterator α m β]
     [Monad m] [Iterator α m β] [Finite α m] :
     haveI : IteratorCollect α m := .defaultImplementation

src/Std/Data/Iterators/Lemmas.lean

@@ -0,0 +1,9 @@
+/-
+Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Paul Reichert
+-/
+prelude
+import Std.Data.Iterators.Lemmas.Basic
+import Std.Data.Iterators.Lemmas.Monadic
+import Std.Data.Iterators.Lemmas.Consumers

src/Std/Data/Iterators/Lemmas/Basic.lean

@@ -0,0 +1,45 @@
+/-
+Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Paul Reichert
+-/
+prelude
+import Std.Data.Iterators.Basic
+
+namespace Std.Iterators
+
+/--
+Induction principle for finite iterators: One can define a function `f` that maps every
+iterator `it` to an element of `motive it` by defining `f it` in terms of the values of `f` on
+the plausible successors of `it'.
+-/
+@[specialize]
+def Iter.inductSteps {α β} [Iterator α Id β] [Finite α Id]
+  (motive : Iter (α := α) β → Sort x)
+  (step : (it : Iter (α := α) β) →
+    (ih_yield : ∀ {it' : Iter (α := α) β} {out : β},
+      it.IsPlausibleStep (.yield it' out) → motive it') →
+    (ih_skip : ∀ {it' : Iter (α := α) β}, it.IsPlausibleStep (.skip it') → motive it') →
+    motive it)
+  (it : Iter (α := α) β) : motive it :=
+  step it
+    (fun {it' _} _ => inductSteps motive step it')
+    (fun {it'} _ => inductSteps motive step it')
+termination_by it.finitelyManySteps
+
+/--
+Induction principle for productive iterators: One can define a function `f` that maps every
+iterator `it` to an element of `motive it` by defining `f it` in terms of the values of `f` on
+the plausible skip successors of `it'.
+-/
+@[specialize]
+def Iter.inductSkips {α β} [Iterator α Id β] [Productive α Id]
+  (motive : Iter (α := α) β → Sort x)
+  (step : (it : Iter (α := α) β) →
+    (ih_skip : ∀ {it' : Iter (α := α) β}, it.IsPlausibleStep (.skip it') → motive it') →
+    motive it)
+  (it : Iter (α := α) β) : motive it :=
+  step it (fun {it'} _ => inductSkips motive step it')
+termination_by it.finitelyManySkips
+
+end Std.Iterators

src/Std/Data/Iterators/Lemmas/Consumers.lean

@@ -0,0 +1,8 @@
+/-
+Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Paul Reichert
+-/
+prelude
+import Std.Data.Iterators.Lemmas.Consumers.Monadic
+import Std.Data.Iterators.Lemmas.Consumers.Collect