feat: ForIn, fold(M), drain lemmas for iterators

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

@@ -47,10 +47,10 @@ number of steps. If the iterator is not finite or such an instance is not availa
 verify the behavior of the partial variant.
 -/
 @[always_inline, inline]
-def Iter.foldM {n : Type w → Type w} [Monad n]
+def Iter.foldM {m : Type w → Type w'} [Monad m]
     {α : Type w} {β : Type w} {γ : Type w} [Iterator α Id β] [Finite α Id]
-    [IteratorLoop α Id n] (f : γ → β → n γ)
-    (init : γ) (it : Iter (α := α) β) : n γ :=
+    [IteratorLoop α Id m] (f : γ → β → m γ)
+    (init : γ) (it : Iter (α := α) β) : m γ :=
   ForIn.forIn it init (fun x acc => ForInStep.yield <$> f acc x)
 
 /--
@@ -63,10 +63,10 @@ This is a partial, potentially nonterminating, function. It is not possible to f
 its behavior. If the iterator has a `Finite` instance, consider using `IterM.foldM` instead.
 -/
 @[always_inline, inline]
-def Iter.Partial.foldM {n : Type w → Type w} [Monad n]
+def Iter.Partial.foldM {m : Type w → Type w'} [Monad m]
     {α : Type w} {β : Type w} {γ : Type w} [Iterator α Id β]
-    [IteratorLoopPartial α Id n] (f : γ → β → n γ)
-    (init : γ) (it : Iter.Partial (α := α) β) : n γ :=
+    [IteratorLoopPartial α Id m] (f : γ → β → m γ)
+    (init : γ) (it : Iter.Partial (α := α) β) : m γ :=
   ForIn.forIn it init (fun x acc => ForInStep.yield <$> f acc x)
 
 /--

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

@@ -170,8 +170,8 @@ It simply iterates through the iterator using `IterM.step`. For certain iterator
 implementations are possible and should be used instead.
 -/
 @[always_inline, inline]
-def IteratorLoopPartial.defaultImplementation {α : Type w} {m : Type w → Type w'} {n : Type w → Type w'}
-    [Monad m] [Monad n] [Iterator α m β] :
+def IteratorLoopPartial.defaultImplementation {α : Type w} {m : Type w → Type w'}
+    {n : Type w → Type w'} [Monad m] [Monad n] [Iterator α m β] :
     IteratorLoopPartial α m n where
   forInPartial lift := IterM.DefaultConsumers.forInPartial lift _
 
@@ -182,6 +182,19 @@ instance (α : Type w) (m : Type w → Type w') (n : Type w → Type w')
   letI : IteratorLoop α m n := .defaultImplementation
   ⟨rfl⟩
 
+theorem IteratorLoop.wellFounded_of_finite {m : Type w → Type w'}
+    {α β γ : Type w} [Iterator α m β] [Finite α m] :
+    WellFounded α m (γ := γ) fun _ _ _ => True := by
+  apply Subrelation.wf
+    (r := InvImage IterM.TerminationMeasures.Finite.Rel (fun p => p.1.finitelyManySteps))
+  · intro p' p h
+    apply Relation.TransGen.single
+    obtain ⟨b, h, _⟩ | ⟨h, _⟩ := h
+    · exact ⟨.yield p'.fst b, rfl, h⟩
+    · exact ⟨.skip p'.fst, rfl, h⟩
+  · apply InvImage.wf
+    exact WellFoundedRelation.wf
+
 /--
 This `ForIn`-style loop construct traverses a finite iterator using an `IteratorLoop` instance.
 -/
@@ -192,16 +205,7 @@ def IteratorLoop.finiteForIn {m : Type w → Type w'} {n : Type w → Type w''}
     ForIn n (IterM (α := α) m β) β where
   forIn {γ} [Monad n] it init f :=
     IteratorLoop.forIn (α := α) (m := m) lift γ (fun _ _ _ => True)
-      (by
-        apply Subrelation.wf
-          (r := InvImage IterM.TerminationMeasures.Finite.Rel (fun p => p.1.finitelyManySteps))
-        · intro p' p h
-          apply Relation.TransGen.single
-          obtain ⟨b, h, _⟩ | ⟨h, _⟩ := h
-          · exact ⟨.yield p'.fst b, rfl, h⟩
-          · exact ⟨.skip p'.fst, rfl, h⟩
-        · apply InvImage.wf
-          exact WellFoundedRelation.wf)
+      wellFounded_of_finite
       it init ((⟨·, .intro⟩) <$> f · ·)
 
 instance {m : Type w → Type w'} {n : Type w → Type w''}
@@ -227,7 +231,7 @@ number of steps. If the iterator is not finite or such an instance is not availa
 verify the behavior of the partial variant.
 -/
 @[always_inline, inline]
-def IterM.foldM {m : Type w → Type w'} {n : Type w → Type w'} [Monad n]
+def IterM.foldM {m : Type w → Type w'} {n : Type w → Type w''} [Monad n]
     {α : Type w} {β : Type w} {γ : Type w} [Iterator α m β] [Finite α m] [IteratorLoop α m n]
     [MonadLiftT m n]
     (f : γ → β → n γ) (init : γ) (it : IterM (α := α) m β) : n γ :=
@@ -295,7 +299,7 @@ verify the behavior of the partial variant.
 def IterM.drain {α : Type w} {m : Type w → Type w'} [Monad m] {β : Type w}
     [Iterator α m β] [Finite α m] (it : IterM (α := α) m β) [IteratorLoop α m m] :
     m PUnit :=
-  it.foldM (γ := PUnit) (fun _ _ => pure .unit) .unit
+  it.fold (γ := PUnit) (fun _ _ => .unit) .unit
 
 /--
 Iterates over the whole iterator, applying the monadic effects of each step, discarding all
@@ -308,6 +312,6 @@ its behavior. If the iterator has a `Finite` instance, consider using `IterM.dra
 def IterM.Partial.drain {α : Type w} {m : Type w → Type w'} [Monad m] {β : Type w}
     [Iterator α m β] (it : IterM.Partial (α := α) m β) [IteratorLoopPartial α m m] :
     m PUnit :=
-  it.foldM (γ := PUnit) (fun _ _ => pure .unit) .unit
+  it.fold (γ := PUnit) (fun _ _ => .unit) .unit
 
 end Std.Iterators

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

@@ -6,3 +6,4 @@ Authors: Paul Reichert
 prelude
 import Std.Data.Iterators.Lemmas.Consumers.Monadic
 import Std.Data.Iterators.Lemmas.Consumers.Collect
+import Std.Data.Iterators.Lemmas.Consumers.Loop

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

@@ -0,0 +1,169 @@
+/-
+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 Init.Data.List.Control
+import Std.Data.Iterators.Lemmas.Basic
+import Std.Data.Iterators.Lemmas.Consumers.Collect
+import Std.Data.Iterators.Lemmas.Consumers.Monadic.Loop
+import Std.Data.Iterators.Consumers.Collect
+import Std.Data.Iterators.Consumers.Loop
+
+namespace Std.Iterators
+
+theorem Iter.forIn_eq {α β : Type w} [Iterator α Id β] [Finite α Id]
+    {m : Type w → Type w''} [Monad m] [IteratorLoop α Id m] [hl : LawfulIteratorLoop α Id m]
+    {γ : Type w} {it : Iter (α := α) β} {init : γ}
+    {f : β → γ → m (ForInStep γ)} :
+    ForIn.forIn it init f =
+      IterM.DefaultConsumers.forIn (fun _ c => pure c.run) γ (fun _ _ _ => True)
+        IteratorLoop.wellFounded_of_finite it.toIterM init ((⟨·, .intro⟩) <$> f · ·) := by
+  cases hl.lawful; rfl
+
+theorem Iter.forIn_eq_forIn_toIterM {α β : Type w} [Iterator α Id β]
+    [Finite α Id] {m : Type w → Type w''} [Monad m] [LawfulMonad m]
+    [IteratorLoop α Id m] [LawfulIteratorLoop α Id m]
+    {γ : Type w} {it : Iter (α := α) β} {init : γ}
+    {f : β → γ → m (ForInStep γ)} :
+    ForIn.forIn it init f = letI : MonadLift Id m := ⟨pure⟩; ForIn.forIn it.toIterM init f := by
+  rfl
+
+theorem Iter.forIn_eq_match_step {α β : Type w} [Iterator α Id β]
+    [Finite α Id] {m : Type w → Type w''} [Monad m] [LawfulMonad m]
+    [IteratorLoop α Id m] [LawfulIteratorLoop α Id m]
+    {γ : Type w} {it : Iter (α := α) β} {init : γ}
+    {f : β → γ → m (ForInStep γ)} :
+    ForIn.forIn it init f = (do
+      match it.step with
+      | .yield it' out _ =>
+        match ← f out init with
+        | .yield c => ForIn.forIn it' c f
+        | .done c => return c
+      | .skip it' _ => ForIn.forIn it' init f
+      | .done _ => return init) := by
+  rw [Iter.forIn_eq_forIn_toIterM, @IterM.forIn_eq_match_step, Iter.step]
+  simp only [liftM, monadLift, pure_bind]
+  generalize it.toIterM.step = step
+  cases step using PlausibleIterStep.casesOn
+  · apply bind_congr
+    intro forInStep
+    rfl
+  · rfl
+  · rfl
+
+theorem Iter.forIn_toList {α β : Type w} [Iterator α Id β]
+    [Finite α Id] {m : Type w → Type w''} [Monad m] [LawfulMonad m]
+    [IteratorLoop α Id m] [LawfulIteratorLoop α Id m]
+    [IteratorCollect α Id] [LawfulIteratorCollect α Id]
+    {γ : Type w} {it : Iter (α := α) β} {init : γ}
+    {f : β → γ → m (ForInStep γ)} :
+    ForIn.forIn it.toList init f = ForIn.forIn it init f := by
+  rw [List.forIn_eq_foldlM]
+  induction it using Iter.inductSteps generalizing init with case step it ihy ihs =>
+  rw [forIn_eq_match_step, Iter.toList_eq_match_step]
+  simp only [map_eq_pure_bind]
+  generalize it.step = step
+  cases step using PlausibleIterStep.casesOn
+  · rename_i it' out h
+    simp only [List.foldlM_cons, bind_pure_comp, map_bind]
+    apply bind_congr
+    intro forInStep
+    cases forInStep
+    · induction it'.toList <;> simp [*]
+    · simp only [ForIn.forIn, forIn', List.forIn'] at ihy
+      simp [ihy h, forIn_eq_forIn_toIterM]
+  · rename_i it' h
+    simp only [bind_pure_comp]
+    rw [ihs h]
+  · simp
+
+theorem Iter.foldM_eq_forIn {α β γ : Type w} [Iterator α Id β] [Finite α Id] {m : Type w → Type w'}
+    [Monad m] [IteratorLoop α Id m] {f : γ → β → m γ}
+    {init : γ} {it : Iter (α := α) β} :
+    it.foldM (init := init) f = ForIn.forIn it init (fun x acc => ForInStep.yield <$> f acc x) :=
+  rfl
+
+theorem Iter.forIn_yield_eq_foldM {α β γ δ : Type w} [Iterator α Id β]
+    [Finite α Id] {m : Type w → Type w''} [Monad m] [LawfulMonad m] [IteratorLoop α Id m]
+    [LawfulIteratorLoop α Id m] {f : β → γ → m δ} {g : β → γ → δ → γ} {init : γ}
+    {it : Iter (α := α) β} :
+    ForIn.forIn it init (fun c b => (fun d => .yield (g c b d)) <$> f c b) =
+      it.foldM (fun b c => g c b <$> f c b) init := by
+  simp [Iter.foldM_eq_forIn]
+
+theorem Iter.foldM_eq_match_step {α β γ : Type w} [Iterator α Id β] [Finite α Id]
+    {m : Type w → Type w'} [Monad m] [LawfulMonad m] [IteratorLoop α Id m]
+    [LawfulIteratorLoop α Id m] {f : γ → β → m γ} {init : γ} {it : Iter (α := α) β} :
+    it.foldM (init := init) f = (do
+      match it.step with
+      | .yield it' out _ => it'.foldM (init := ← f init out) f
+      | .skip it' _ => it'.foldM (init := init) f
+      | .done _ => return init) := by
+  rw [Iter.foldM_eq_forIn, Iter.forIn_eq_match_step]
+  generalize it.step = step
+  cases step using PlausibleIterStep.casesOn <;> simp [foldM_eq_forIn]
+
+theorem Iter.foldlM_toList {α β γ : Type w} [Iterator α Id β] [Finite α Id] {m : Type w → Type w'}
+    [Monad m] [LawfulMonad m] [IteratorLoop α Id m] [LawfulIteratorLoop α Id m]
+    [IteratorCollect α Id] [LawfulIteratorCollect α Id]
+    {f : γ → β → m γ}
+    {init : γ} {it : Iter (α := α) β} :
+    it.toList.foldlM (init := init) f = it.foldM (init := init) f := by
+  rw [Iter.foldM_eq_forIn, ← Iter.forIn_toList]
+  simp only [List.forIn_yield_eq_foldlM, id_map']
+
+theorem IterM.forIn_eq_foldM {α β : Type w} [Iterator α Id β]
+    [Finite α Id] {m : Type w → Type w''} [Monad m] [LawfulMonad m]
+    [IteratorLoop α Id m] [LawfulIteratorLoop α Id m]
+    [IteratorCollect α Id] [LawfulIteratorCollect α Id]
+    {γ : Type w} {it : Iter (α := α) β} {init : γ}
+    {f : β → γ → m (ForInStep γ)} :
+    forIn it init f = ForInStep.value <$>
+      it.foldM (fun c b => match c with
+        | .yield c => f b c
+        | .done c => pure (.done c)) (ForInStep.yield init) := by
+  simp only [← Iter.forIn_toList, List.forIn_eq_foldlM, ← Iter.foldlM_toList]; rfl
+
+theorem Iter.fold_eq_forIn {α β γ : Type w} [Iterator α Id β]
+    [Finite α Id] [IteratorLoop α Id Id] {f : γ → β → γ} {init : γ} {it : Iter (α := α) β} :
+    it.fold (init := init) f =
+      (ForIn.forIn (m := Id) it init (fun x acc => pure (ForInStep.yield (f acc x)))).run := by
+  rfl
+
+theorem Iter.fold_eq_foldM {α β γ : Type w} [Iterator α Id β]
+    [Finite α Id] [IteratorLoop α Id Id] {f : γ → β → γ} {init : γ}
+    {it : Iter (α := α) β} :
+    it.fold (init := init) f = (it.foldM (m := Id) (init := init) (pure <| f · ·)).run := by
+  simp [foldM_eq_forIn, fold_eq_forIn]
+
+theorem Iter.forIn_yield_eq_fold {α β γ : Type w} [Iterator α Id β]
+    [Finite α Id] [IteratorLoop α Id Id]
+    [LawfulIteratorLoop α Id Id] {f : β → γ → γ} {init : γ}
+    {it : Iter (α := α) β} :
+    ForIn.forIn (m := Id) it init (fun c b => .yield (f c b)) =
+      it.fold (fun b c => f c b) init := by
+  simp only [fold_eq_forIn]
+  rfl
+
+theorem Iter.fold_eq_match_step {α β γ : Type w} [Iterator α Id β] [Finite α Id]
+    [IteratorLoop α Id Id] [LawfulIteratorLoop α Id Id]
+    {f : γ → β → γ} {init : γ} {it : Iter (α := α) β} :
+    it.fold (init := init) f = (match it.step with
+      | .yield it' out _ => it'.fold (init := f init out) f
+      | .skip it' _ => it'.fold (init := init) f
+      | .done _ => init) := by
+  rw [fold_eq_foldM, foldM_eq_match_step]
+  simp only [fold_eq_foldM]
+  generalize it.step = step
+  cases step using PlausibleIterStep.casesOn <;> simp
+
+theorem Iter.foldl_toList {α β γ : Type w} [Iterator α Id β] [Finite α Id]
+    [IteratorLoop α Id Id] [LawfulIteratorLoop α Id Id]
+    [IteratorCollect α Id] [LawfulIteratorCollect α Id]
+    {f : γ → β → γ} {init : γ} {it : Iter (α := α) β} :
+    it.toList.foldl (init := init) f = it.fold (init := init) f := by
+  rw [fold_eq_foldM, List.foldl_eq_foldlM, ← Iter.foldlM_toList]
+
+end Std.Iterators