feat: grind annotations, Iter(M).sum and toList+take+(productive combinator) composition lemmas

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

@@ -707,6 +707,9 @@ def Iter.Total.isEmpty {α β : Type w} [Iterator α Id β] [IteratorLoop α Id
 /--
 Steps through the whole iterator, counting the number of outputs emitted.
 
+If the iterator is not finite, this function might run forever. The variant
+`it.ensureTermination.length` always terminates after finitely many steps.
+
 **Performance**:
 
 This function's runtime is linear in the number of steps taken by the iterator.
@@ -746,4 +749,33 @@ def Iter.Partial.size {α : Type w} {β : Type w} [Iterator α Id β] [IteratorL
     (it : Iter.Partial (α := α) β) : Nat :=
   it.it.length
 
+/--
+Steps through the whole iterator, summing up its elements from left to right.
+
+If the iterator is not finite, this function might run forever. The variant
+`it.ensureTermination.sum` always terminates after finitely many steps.
+
+**Performance**:
+
+This function's runtime is linear in the number of steps taken by the iterator.
+-/
+def Iter.sum [Add β] [Zero β] [Iterator α Id β] [IteratorLoop α Id Id]
+    (it : Iter (α := α) β) : β :=
+  it.fold (init := 0) (· + ·)
+
+/--
+Steps through the whole iterator, summing up its elements from left to right.
+
+This variant terminates after finitely many steps and requires a proof that the iterator is
+finite. If such a proof is not available, consider using `Iter.sum`.
+
+**Performance**:
+
+This function's runtime is linear in the number of steps taken by the iterator.
+-/
+@[inline]
+def Iter.Total.sum [Add β] [Zero β] [Iterator α Id β] [IteratorLoop α Id Id] [Finite α Id]
+    (it : Iter.Total (α := α) β) : β :=
+  it.it.sum
+
 end Std

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

@@ -1029,4 +1029,37 @@ def IterM.Partial.size {α : Type w} {m : Type w → Type w'} {β : Type w} [Ite
 
 end Count
 
+section Sum
+
+/--
+Steps through the whole iterator, summing up its elements from left to right.
+
+If the iterator is not finite, this function might run forever. The variant
+`it.ensureTermination.sum` always terminates after finitely many steps.
+
+**Performance**:
+
+This function's runtime is linear in the number of steps taken by the iterator.
+-/
+def IterM.sum [Monad m] [Add β] [Zero β] [Iterator α m β] [IteratorLoop α m m]
+    (it : IterM (α := α) m β) : m β :=
+  it.fold (init := 0) (· + ·)
+
+/--
+Steps through the whole iterator, summing up its elements from left to right.
+
+This variant terminates after finitely many steps and requires a proof that the iterator is
+finite. If such a proof is not available, consider using `Iter.sum`.
+
+**Performance**:
+
+This function's runtime is linear in the number of steps taken by the iterator.
+-/
+@[inline]
+def IterM.Total.sum [Monad m] [Add β] [Zero β] [Iterator α m β] [IteratorLoop α m m] [Finite α m]
+    (it : IterM.Total (α := α) m β) : m β :=
+  it.it.sum
+
+end Sum
+
 end Std

src/Init/Data/Iterators/Lemmas/Combinators/FilterMap.lean

@@ -7,13 +7,17 @@ module
 
 prelude
 public import Init.Data.Iterators.Combinators.FilterMap
+public import Init.Data.Iterators.Combinators.Take
 public import Init.Data.Iterators.Consumers.Collect
 public import Init.Data.Iterators.Consumers.Loop
+public import Init.Data.Iterators.Consumers.Access
+import Init.Data.Iterators.Lemmas.Consumers.Access
 public import Init.Data.List.Control
 import Init.Data.Array.Lemmas
 import Init.Data.Bool
 import Init.Data.Iterators.Lemmas.Basic
 import Init.Data.Iterators.Lemmas.Combinators.Monadic.FilterMap
+import Init.Data.Iterators.Lemmas.Combinators.Take
 import Init.Data.Iterators.Lemmas.Consumers.Collect
 import Init.Data.Iterators.Lemmas.Consumers.Loop
 import Init.Data.Iterators.Lemmas.Consumers.Monadic.Loop
@@ -728,7 +732,7 @@ end Fold
 
 section Count
 
-@[simp]
+@[simp, grind =]
 theorem Iter.length_map {α β β' : Type w} [Iterator α Id β]
     [IteratorLoop α Id Id] [Finite α Id] [LawfulIteratorLoop α Id Id]
     {it : Iter (α := α) β} {f : β → β'} :
@@ -1055,4 +1059,22 @@ theorem Iter.all_map {α β β' : Type w}
   · simp [ihs ‹_›]
   · simp
 
+@[simp, grind =]
+theorem Iter.atIdxSlow?_map [Iterators.Productive α Id]
+    {it : Iter (α := α) β} {f : β → γ} :
+    (it.map f).atIdxSlow? k = (it.atIdxSlow? k).map f := by
+  induction k, it using atIdxSlow?.induct_unfolding
+  all_goals
+    rw [atIdxSlow?_eq_match, step_map]
+    simp [*]
+
+@[simp, grind =]
+theorem Iter.toList_take_map [Iterators.Productive α Id]
+    {it : Iter (α := α) β} {f : β → γ} :
+    (it.map f |>.take k).toList = (it.take k).toList.map f := by
+  apply List.ext_getElem?
+  simp only [getElem?_toList_eq_atIdxSlow?, atIdxSlow?_take, atIdxSlow?_map, List.getElem?_map]
+  intro i
+  split <;> simp
+
 end Std

src/Init/Data/Iterators/Lemmas/Combinators/Take.lean

@@ -49,6 +49,7 @@ theorem Iter.step_take {α β} [Iterator α Id β] {n : Nat}
     cases step.inflate using PlausibleIterStep.casesOn <;>
       simp [PlausibleIterStep.yield, PlausibleIterStep.skip, PlausibleIterStep.done]
 
+@[grind =]
 theorem Iter.atIdxSlow?_take {α β}
     [Iterator α Id β] [Productive α Id] {k l : Nat}
     {it : Iter (α := α) β} :

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

@@ -434,6 +434,12 @@ theorem Iter.fold_hom {γ₁ : Type x₁} {γ₂ : Type x₂} [Iterator α Id β
   · rw [ihs ‹_›]
   · simp
 
+theorem Iter.fold_assoc [Iterator α Id β] [Finite α Id]
+    [IteratorLoop α Id Id] [LawfulIteratorLoop α Id Id]
+    {it : Iter (α := α) β} {op : β → β → β} [Associative op] :
+    it.fold (init := op b₁ b₂) op = op b₁ (it.fold (init := b₂) op) := by
+  simp [Iter.fold_eq_fold_toIterM, IterM.fold_assoc]
+
 theorem Iter.toList_eq_fold {α β : Type w} [Iterator α Id β]
     [Finite α Id] [IteratorLoop α Id Id] [LawfulIteratorLoop α Id Id]
     {it : Iter (α := α) β} :
@@ -521,7 +527,7 @@ def Iter.size_toArray_eq_size := @size_toArray_eq_length
 @[deprecated Iter.size_toArray_eq_length (since := "2026-01-28")]
 def Iter.size_toArray_eq_count := @size_toArray_eq_length
 
-@[simp]
+@[simp, grind =]
 theorem Iter.length_toList_eq_length {α β : Type w} [Iterator α Id β] [Finite α Id]
     [IteratorLoop α Id Id] [LawfulIteratorLoop α Id Id]
     {it : Iter (α := α) β} :
@@ -984,4 +990,50 @@ theorem Iter.isEmpty_toList {α β : Type w} [Iterator α Id β] [IteratorLoop 
   rw [isEmpty_eq_match_step, toList_eq_match_step]
   cases it.step using PlausibleIterStep.casesOn <;> simp [*]
 
+theorem Iter.sum_eq_sum_toIterM
+    [Add β] [Zero β]
+    [Iterator α Id β] [IteratorLoop α Id Id] [LawfulIteratorLoop α Id Id]
+    [Finite α Id] {it : Iter (α := α) β} :
+    it.sum = it.toIterM.sum.run := by
+  simp [Iter.sum, IterM.sum, Iter.fold_eq_fold_toIterM]
+
+theorem Iter.sum_eq_match_step
+    [Add β] [Zero β] [Associative (α := β) (· + ·)] [LawfulIdentity (α := β) (· + ·) 0]
+    [Iterator α Id β] [IteratorLoop α Id Id] [LawfulIteratorLoop α Id Id]
+    [Finite α Id] {it : Iter (α := α) β} :
+    it.sum = (match it.step.val with
+      | .yield it' out => out + it'.sum
+      | .skip it' => it'.sum
+      | .done => 0) := by
+  rw [Iter.sum_eq_sum_toIterM, IterM.sum_eq_match_step]
+  simp only [bind_pure_comp, Id.run_bind, Iter.step]
+  cases it.toIterM.step.run.inflate using PlausibleIterStep.casesOn <;> simp [Iter.sum_eq_sum_toIterM]
+
+@[simp, grind =]
+theorem Iter.sum_toList [Add β] [Zero β]
+    [Associative (α := β) (· + ·)]
+    [LawfulIdentity (· + ·) (0 : β)]
+    [Iterator α Id β] [IteratorLoop α Id Id]
+    [LawfulIteratorLoop α Id Id] [Iterators.Finite α Id] {it : Iter (α := α) β} :
+    it.toList.sum = it.sum := by
+  simp only [Iter.sum, ← Iter.foldl_toList, List.sum_eq_foldl]
+
+@[simp, grind =]
+theorem Iter.sum_toArray [Add β] [Zero β]
+    [Associative (α := β) (· + ·)]
+    [LawfulIdentity (· + ·) (0 : β)]
+    [Iterator α Id β] [IteratorLoop α Id Id]
+    [LawfulIteratorLoop α Id Id] [Iterators.Finite α Id] {it : Iter (α := α) β} :
+    it.toArray.sum = it.sum := by
+  simp [← Iter.toArray_toList, Iter.sum_toList]
+
+@[simp, grind =]
+theorem Iter.sum_toListRev [Add β] [Zero β]
+    [Associative (α := β) (· + ·)] [Commutative (α := β) (· + ·)]
+    [LawfulIdentity (· + ·) (0 : β)]
+    [Iterator α Id β] [IteratorLoop α Id Id]
+    [LawfulIteratorLoop α Id Id] [Iterators.Finite α Id] {it : Iter (α := α) β} :
+    it.toListRev.sum = it.sum := by
+  simp [Iter.toListRev_eq, List.sum_reverse, Iter.sum_toList]
+
 end Std

src/Init/Data/Iterators/Lemmas/Consumers/Monadic/Loop.lean

@@ -368,6 +368,19 @@ theorem IterM.fold_hom {m : Type w → Type w'} [Iterator α m β] [Finite α m]
   · rw [ihs ‹_›]
   · simp
 
+theorem IterM.fold_assoc {m : Type w → Type w'} [Iterator α m β] [Finite α m]
+    [Monad m] [LawfulMonad m] [IteratorLoop α m m] [LawfulIteratorLoop α m m]
+    {it : IterM (α := α) m β} {op : β → β → β} [Associative op] :
+    it.fold (init := op b₁ b₂) op = (op b₁ ·) <$> it.fold (init := b₂) op := by
+  induction it using IterM.inductSteps generalizing b₁ b₂ with | step it ihy ihs
+  rw [IterM.fold_eq_match_step, IterM.fold_eq_match_step]
+  simp only [map_eq_pure_bind, bind_assoc]
+  apply bind_congr; intro step
+  split
+  · simp [ihy ‹_›, Associative.assoc]
+  · simp [ihs ‹_›]
+  · simp
+
 theorem IterM.toList_eq_fold {α β : Type w} {m : Type w → Type w'} [Iterator α m β]
     [Finite α m] [Monad m] [LawfulMonad m] [IteratorLoop α m m] [LawfulIteratorLoop α m m]
     {it : IterM (α := α) m β} :
@@ -906,4 +919,48 @@ theorem IterM.isEmpty_eq_match_step {α β : Type w} {m : Type w → Type w'} [M
   simp only [DefaultConsumers.forIn_eq, *]
   exact IterM.DefaultConsumers.forIn'_eq_forIn' _ this (by simp)
 
+theorem IterM.sum_eq_match_step
+    [Add β] [Zero β] [Associative (α := β) (· + ·)] [LawfulIdentity (α := β) (· + ·) 0]
+    [Monad m] [Iterator α m β] [IteratorLoop α m m] [LawfulIteratorLoop α m m] [LawfulMonad m]
+    [Finite α m] {it : IterM (α := α) m β} :
+    it.sum = (do
+      match (← it.step).inflate.val with
+      | .yield it' out => return out + (← it'.sum)
+      | .skip it' => it'.sum
+      | .done => return 0) := by
+  rw [IterM.sum, IterM.fold_eq_match_step]
+  apply bind_congr; intro step
+  cases step.inflate using PlausibleIterStep.casesOn
+  · have (b : β) : 0 + b = b + 0 := by simp [LawfulLeftIdentity.left_id, LawfulRightIdentity.right_id]
+    simp [IterM.sum, this, IterM.fold_assoc]
+  · simp [IterM.sum]
+  · simp
+
+@[simp, grind =]
+theorem IterM.sum_toList [Monad m] [Add β] [Zero β]
+    [Associative (α := β) (· + ·)]
+    [LawfulIdentity (· + ·) (0 : β)]
+    [Iterator α m β] [IteratorLoop α m m] [LawfulMonad m]
+    [LawfulIteratorLoop α m m] [Iterators.Finite α m] {it : IterM (α := α) m β} :
+    List.sum <$> it.toList = it.sum := by
+  simp only [IterM.sum, ← IterM.foldl_toList, ← List.sum_eq_foldl]
+
+@[simp, grind =]
+theorem IterM.sum_toArray [Monad m] [Add β] [Zero β]
+    [Associative (α := β) (· + ·)]
+    [LawfulIdentity (· + ·) (0 : β)]
+    [Iterator α m β] [IteratorLoop α m m] [LawfulMonad m]
+    [LawfulIteratorLoop α m m] [Iterators.Finite α m] {it : IterM (α := α) m β} :
+    Array.sum <$> it.toArray = it.sum := by
+  simp [← IterM.toArray_toList, IterM.sum_toList]
+
+@[simp, grind =]
+theorem IterM.sum_toListRev [Monad m] [Add β] [Zero β]
+    [Associative (α := β) (· + ·)] [Commutative (α := β) (· + ·)]
+    [LawfulIdentity (· + ·) (0 : β)]
+    [Iterator α m β] [IteratorLoop α m m] [LawfulMonad m]
+    [LawfulIteratorLoop α m m] [Iterators.Finite α m] {it : IterM (α := α) m β} :
+    List.sum <$> it.toListRev = it.sum := by
+  simp [IterM.toListRev_eq, List.sum_reverse, IterM.sum_toList]
+
 end Std

src/Std/Data/Iterators/Lemmas/Combinators/Drop.lean

@@ -40,6 +40,7 @@ theorem Iter.step_drop {α β} [Iterator α Id β] {n : Nat}
   cases step.inflate using PlausibleIterStep.casesOn <;> cases n <;>
     simp [PlausibleIterStep.yield, PlausibleIterStep.skip, PlausibleIterStep.done]
 
+@[simp, grind =]
 theorem Iter.atIdxSlow?_drop {α β}
     [Iterator α Id β] [Productive α Id] {k l : Nat}
     {it : Iter (α := α) β} :
@@ -50,24 +51,32 @@ theorem Iter.atIdxSlow?_drop {α β}
     rw [atIdxSlow?_eq_match, atIdxSlow?_eq_match, step_drop]
     cases it.step using PlausibleIterStep.casesOn <;> simp [*]
 
-@[simp]
+@[simp, grind =]
 theorem Iter.toList_drop {α β} [Iterator α Id β] {n : Nat}
     [Finite α Id] {it : Iter (α := α) β} :
     (it.drop n).toList = it.toList.drop n := by
   ext
   simp only [getElem?_toList_eq_atIdxSlow?, List.getElem?_drop, atIdxSlow?_drop]
   rw [Nat.add_comm]
 
-@[simp]
+@[simp, grind =]
 theorem Iter.toListRev_drop {α β} [Iterator α Id β] {n : Nat}
     [Finite α Id] {it : Iter (α := α) β} :
     (it.drop n).toListRev = (it.toList.reverse.take (it.toList.length - n)) := by
   rw [toListRev_eq, toList_drop, List.reverse_drop]
 
-@[simp]
+@[simp, grind =]
 theorem Iter.toArray_drop {α β} [Iterator α Id β] {n : Nat}
     [Finite α Id] {it : Iter (α := α) β} :
     (it.drop n).toArray = it.toArray.extract n := by
   rw [← toArray_toList, ← toArray_toList, ← List.toArray_drop, toList_drop]
 
+@[simp, grind =]
+theorem Iter.toList_take_drop {α β} [Iterator α Id β] {n : Nat}
+    [Productive α Id] {it : Iter (α := α) β} :
+    (it.drop m |>.take n).toList = (it.take (m + n)).toList.drop m := by
+  ext
+  simp [getElem?_toList_eq_atIdxSlow?, List.getElem?_drop, atIdxSlow?_take, atIdxSlow?_drop,
+    Nat.add_comm]
+
 end Std

src/Std/Data/Iterators/Lemmas/Combinators/TakeWhile.lean

@@ -6,7 +6,9 @@ Authors: Paul Reichert
 module
 
 prelude
+public import Init.Data.Iterators.Combinators.Take
 public import Std.Data.Iterators.Combinators.TakeWhile
+public import Init.Data.Iterators.Lemmas.Combinators.Take
 public import Std.Data.Iterators.Lemmas.Combinators.Monadic.TakeWhile
 public import Std.Data.Iterators.Lemmas.Consumers
 import Init.Data.List.TakeDrop
@@ -157,4 +159,17 @@ theorem Iter.toArray_takeWhile_of_finite {α β} [Iterator α Id β] {P}
     (it.takeWhile P).toArray = it.toArray.takeWhile P := by
   rw [← toArray_toList, ← toArray_toList, List.takeWhile_toArray, toList_takeWhile_of_finite]
 
+@[simp, grind =]
+theorem Iter.toList_take_takeWhile {α β} [Iterator α Id β] {P} [Productive α Id]
+    {it : Iter (α := α) β} :
+    (it.takeWhile P |>.take k).toList = (it.take k).toList.takeWhile P := by
+  ext i b
+  have (k' : Nat) (hk : i < k) (hk' : k' ≤ i) : k' < k := Nat.lt_of_le_of_lt hk' hk
+  simp only [getElem?_toList_eq_atIdxSlow?, atIdxSlow?_take, atIdxSlow?_takeWhile,
+    List.getElem?_takeWhile, Option.ite_none_right_eq_some]
+  apply Iff.intro
+  · simp +contextual [this]
+  · intro h
+    simpa +contextual [*] using h.1
+
 end Std