feat: variants of List.forIn_eq_foldlM (#6023)

src/Init/Data/List/Lemmas.lean

@@ -999,6 +999,21 @@ theorem foldr_rel {l : List α} {f g : α → β → β} {a b : β} (r : β →
     · simp
     · exact ih h fun a m c c' h => h' _ (by simp_all) _ _ h
 
+@[simp] theorem foldl_add_const (l : List α) (a b : Nat) :
+    l.foldl (fun x _ => x + a) b = b + a * l.length := by
+  induction l generalizing b with
+  | nil => simp
+  | cons y l ih =>
+    simp only [foldl_cons, ih, length_cons, Nat.mul_add, Nat.mul_one, Nat.add_assoc,
+      Nat.add_comm a]
+
+@[simp] theorem foldr_add_const (l : List α) (a b : Nat) :
+    l.foldr (fun _ x => x + a) b = b + a * l.length := by
+  induction l generalizing b with
+  | nil => simp
+  | cons y l ih =>
+    simp only [foldr_cons, ih, length_cons, Nat.mul_add, Nat.mul_one, Nat.add_assoc]
+
 /-! ### getLast -/
 
 theorem getLast_eq_getElem : ∀ (l : List α) (h : l ≠ []),

src/Init/Data/List/Monadic.lean

@@ -208,8 +208,8 @@ in which whenever we reach `.done b` we keep that value through the rest of the
 theorem forIn'_eq_foldlM [Monad m] [LawfulMonad m]
     (l : List α) (f : (a : α) → a ∈ l → β → m (ForInStep β)) (init : β) :
     forIn' l init f = ForInStep.value <$>
-      l.attach.foldlM (fun b a => match b with
-        | .yield b => f a.1 a.2 b
+      l.attach.foldlM (fun b ⟨a, m⟩ => match b with
+        | .yield b => f a m b
         | .done b => pure (.done b)) (ForInStep.yield init) := by
   induction l generalizing init with
   | nil => simp
@@ -234,6 +234,31 @@ theorem forIn'_eq_foldlM [Monad m] [LawfulMonad m]
     | .yield b =>
       simp [ih, List.foldlM_map]
 
+/-- We can express a for loop over a list which always yields as a fold. -/
+@[simp] theorem forIn'_yield_eq_foldlM [Monad m] [LawfulMonad m]
+    (l : List α) (f : (a : α) → a ∈ l → β → m γ) (g : (a : α) → a ∈ l → β → γ → β) (init : β) :
+    forIn' l init (fun a m b => (fun c => .yield (g a m b c)) <$> f a m b) =
+      l.attach.foldlM (fun b ⟨a, m⟩ => g a m b <$> f a m b) init := by
+  simp only [forIn'_eq_foldlM]
+  generalize l.attach = l'
+  induction l' generalizing init <;> simp_all
+
+theorem forIn'_pure_yield_eq_foldl [Monad m] [LawfulMonad m]
+    (l : List α) (f : (a : α) → a ∈ l → β → β) (init : β) :
+    forIn' l init (fun a m b => pure (.yield (f a m b))) =
+      pure (f := m) (l.attach.foldl (fun b ⟨a, h⟩ => f a h b) init) := by
+  simp only [forIn'_eq_foldlM]
+  generalize l.attach = l'
+  induction l' generalizing init <;> simp_all
+
+@[simp] theorem forIn'_yield_eq_foldl
+    (l : List α) (f : (a : α) → a ∈ l → β → β) (init : β) :
+    forIn' (m := Id) l init (fun a m b => .yield (f a m b)) =
+      l.attach.foldl (fun b ⟨a, h⟩ => f a h b) init := by
+  simp only [forIn'_eq_foldlM]
+  generalize l.attach = l'
+  induction l' generalizing init <;> simp_all
+
 /--
 We can express a for loop over a list as a fold,
 in which whenever we reach `.done b` we keep that value through the rest of the fold.
@@ -260,6 +285,28 @@ theorem forIn_eq_foldlM [Monad m] [LawfulMonad m]
     | .yield b =>
       simp [ih]
 
+/-- We can express a for loop over a list which always yields as a fold. -/
+@[simp] theorem forIn_yield_eq_foldlM [Monad m] [LawfulMonad m]
+    (l : List α) (f : α → β → m γ) (g : α → β → γ → β) (init : β) :
+    forIn l init (fun a b => (fun c => .yield (g a b c)) <$> f a b) =
+      l.foldlM (fun b a => g a b <$> f a b) init := by
+  simp only [forIn_eq_foldlM]
+  induction l generalizing init <;> simp_all
+
+theorem forIn_pure_yield_eq_foldl [Monad m] [LawfulMonad m]
+    (l : List α) (f : α → β → β) (init : β) :
+    forIn l init (fun a b => pure (.yield (f a b))) =
+      pure (f := m) (l.foldl (fun b a => f a b) init) := by
+  simp only [forIn_eq_foldlM]
+  induction l generalizing init <;> simp_all
+
+@[simp] theorem forIn_yield_eq_foldl
+    (l : List α) (f : α → β → β) (init : β) :
+    forIn (m := Id) l init (fun a b => .yield (f a b)) =
+      l.foldl (fun b a => f a b) init := by
+  simp only [forIn_eq_foldlM]
+  induction l generalizing init <;> simp_all
+
 /-! ### allM -/
 
 theorem allM_eq_not_anyM_not [Monad m] [LawfulMonad m] (p : α → m Bool) (as : List α) :

src/Std/Data/DHashMap/Internal/AssocList/Lemmas.lean

@@ -40,7 +40,7 @@ theorem foldl_eq {f : δ → (a : α) → β a → δ} {init : δ} {l : AssocLis
 @[simp]
 theorem length_eq {l : AssocList α β} : l.length = l.toList.length := by
   rw [length, foldl_eq]
-  suffices ∀ n, l.toList.foldl (fun d _ => d + 1) n = l.toList.length + n by simpa using this 0
+  suffices ∀ n, l.toList.foldl (fun d _ => d + 1) n = l.toList.length + n by simp
   induction l
   · simp
   · next _ _ t ih =>

tests/lean/run/list_simp.lean

@@ -477,6 +477,15 @@ attribute [local simp] Id.run in
       s := s + i
     pure s) ~> 10
 
+attribute [local simp] Id.run in
+variable (l : List α) (k m : Nat) in
+#check_simp
+  (Id.run do
+    let mut x := m
+    for _ in l do
+      x := x + k
+    pure x) ~> m + k * l.length
+
 /-! ### mapM -/
 
 /-! ### forM -/