feat: lemmas about iterator collectors (#8380)

This PR provides simple lemmas about `toArray`, `toList` and `toListRev`
for the iterator library.

It also changes the definition of `Iter` and `IterM` so that they aren't
equal anymore and in particular not definitionally equal. While it was
very convenient to have them be definitionally equal when working with
dependent code, it was also confusing and annoying that one would
sometimes end up with something like `it.toList = IterM.toList it`,
where `it : Iter β`.

Author: datokrat

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