feat: add List/Array/Vector.ofFnM (#8389)

This PR adds the `List/Array/Vector.ofFnM`, the monadic analogues of
`ofFn`, along with basic theory.

At the same time we pave some potholes in nearby API.

Author: kim-em

src/Init/Data/Array/Basic.lean

@@ -112,6 +112,8 @@ theorem mem_def {a : α} {as : Array α} : a ∈ as ↔ a ∈ as.toList :=
   rw [Array.mem_def, ← getElem_toList]
   apply List.getElem_mem
 
+@[simp, grind] theorem mkEmpty_eq_empty {n : Nat} : (mkEmpty n : Array α) = #[] := rfl
+
 end Array
 
 namespace List
@@ -334,6 +336,8 @@ def ofFn {n} (f : Fin n → α) : Array α := go 0 (emptyWithCapacity n) where
     if h : i < n then go (i+1) (acc.push (f ⟨i, h⟩)) else acc
   decreasing_by simp_wf; decreasing_trivial_pre_omega
 
+-- See also `Array.ofFnM` defined in `Init.Data.Array.OfFn`.
+
 /--
 Constructs an array that contains all the numbers from `0` to `n`, exclusive.
 

src/Init/Data/Array/Lemmas.lean

@@ -247,6 +247,12 @@ theorem back?_pop {xs : Array α} :
 
 /-! ### push -/
 
+@[simp] theorem push_empty : #[].push x = #[x] := rfl
+
+@[simp] theorem toList_push {xs : Array α} {x : α} : (xs.push x).toList = xs.toList ++ [x] := by
+  rcases xs with ⟨xs⟩
+  simp
+
 @[simp] theorem push_ne_empty {a : α} {xs : Array α} : xs.push a ≠ #[] := by
   cases xs
   simp
@@ -3266,6 +3272,22 @@ rather than `(arr.push a).size` as the argument.
     (xs.push a).foldrM f init start = f a init >>= xs.foldrM f := by
   simp [← foldrM_push, h]
 
+@[simp, grind] theorem _root_.List.foldrM_push_eq_append [Monad m] [LawfulMonad m] {l : List α} {f : α → m β} {xs : Array β} :
+    l.foldrM (fun x xs => xs.push <$> f x) xs = do return xs ++ (← l.reverse.mapM f).toArray := by
+  induction l with
+  | nil => simp
+  | cons a l ih =>
+    simp [ih]
+    congr 1
+    funext l'
+    congr 1
+    funext x
+    simp
+
+@[simp, grind] theorem _root_.List.foldlM_push_eq_append [Monad m] [LawfulMonad m] {l : List α} {f : α → m β} {xs : Array β} :
+    l.foldlM (fun xs x => xs.push <$> f x) xs = do return xs ++ (← l.mapM f).toArray := by
+  induction l generalizing xs <;> simp [*]
+
 /-! ### foldl / foldr -/
 
 @[grind] theorem foldl_empty {f : β → α → β} {init : β} : (#[].foldl f init) = init := rfl
@@ -3362,6 +3384,24 @@ rather than `(arr.push a).size` as the argument.
   rcases as with ⟨as⟩
   simp
 
+@[simp, grind] theorem _root_.List.foldr_push_eq_append {l : List α} {f : α → β} {xs : Array β} :
+    l.foldr (fun x xs => xs.push (f x)) xs = xs ++ (l.reverse.map f).toArray := by
+  induction l <;> simp [*]
+
+/-- Variant of `List.foldr_push_eq_append` specialized to `f = id`. -/
+@[simp, grind] theorem _root_.List.foldr_push_eq_append' {l : List α} {xs : Array α} :
+    l.foldr (fun x xs => xs.push x) xs = xs ++ l.reverse.toArray := by
+  induction l <;> simp [*]
+
+@[simp, grind] theorem _root_.List.foldl_push_eq_append {l : List α} {f : α → β} {xs : Array β} :
+    l.foldl (fun xs x => xs.push (f x)) xs = xs ++ (l.map f).toArray := by
+  induction l generalizing xs <;> simp [*]
+
+/-- Variant of `List.foldl_push_eq_append` specialized to `f = id`. -/
+@[simp, grind] theorem _root_.List.foldl_push_eq_append' {l : List α} {xs : Array α} :
+    l.foldl (fun xs x => xs.push x) xs = xs ++ l.toArray := by
+  induction l <;> simp [*]
+
 @[simp, grind] theorem foldr_append_eq_append {xs : Array α} {f : α → Array β} {ys : Array β} :
     xs.foldr (f · ++ ·) ys = (xs.map f).flatten ++ ys := by
   rcases xs with ⟨xs⟩

src/Init/Data/Array/Monadic.lean

@@ -34,6 +34,11 @@ open Nat
 @[simp] theorem mapM_id {xs : Array α} {f : α → Id β} : xs.mapM f = xs.map f :=
   mapM_pure
 
+@[simp] theorem mapM_map [Monad m] [LawfulMonad m] {f : α → β} {g : β → m γ} {xs : Array α} :
+    (xs.map f).mapM g = xs.mapM (g ∘ f) := by
+  rcases xs with ⟨xs⟩
+  simp
+
 @[simp] theorem mapM_append [Monad m] [LawfulMonad m] {f : α → m β} {xs ys : Array α} :
     (xs ++ ys).mapM f = (return (← xs.mapM f) ++ (← ys.mapM f)) := by
   rcases xs with ⟨xs⟩

src/Init/Data/Array/OfFn.lean

@@ -9,6 +9,7 @@ prelude
 import all Init.Data.Array.Basic
 import Init.Data.Array.Lemmas
 import Init.Data.List.OfFn
+import Init.Data.List.FinRange
 
 /-!
 # Theorems about `Array.ofFn`
@@ -17,8 +18,42 @@ import Init.Data.List.OfFn
 set_option linter.listVariables true -- Enforce naming conventions for `List`/`Array`/`Vector` variables.
 set_option linter.indexVariables true -- Enforce naming conventions for index variables.
 
+namespace List
+
+/-! ### ofFnM
+
+We prove some theorems about `List.ofFnM` here, as they rely on `Array` lemmas.
+-/
+
+theorem ofFnM_succ {n} [Monad m] [LawfulMonad m] {f : Fin (n + 1) → m α} :
+    ofFnM f = (do
+      let a ← f 0
+      let as ← ofFnM fun i => f i.succ
+      pure (a :: as)) := by
+  simp [ofFnM, Fin.foldlM_eq_finRange_foldlM, List.foldlM_push_eq_append, List.finRange_succ, Function.comp_def]
+
+theorem ofFnM_succ_last {n} [Monad m] [LawfulMonad m] {f : Fin (n + 1) → m α} :
+    ofFnM f = (do
+      let as ← ofFnM fun i => f i.castSucc
+      let a  ← f (Fin.last n)
+      pure (as ++ [a])) := by
+  simp [ofFnM, Fin.foldlM_succ_last]
+
+theorem ofFnM_add {n m} [Monad m] [LawfulMonad m] {f : Fin (n + k) → m α} :
+    ofFnM f = (do
+      let as ← ofFnM fun i : Fin n => f (i.castLE (Nat.le_add_right n k))
+      let bs ← ofFnM fun i : Fin k => f (i.natAdd n)
+      pure (as ++ bs)) := by
+  induction k with
+  | zero => simp
+  | succ k ih => simp [ofFnM_succ_last, ih]
+
+end List
+
 namespace Array
 
+/-! ### ofFn -/
+
 @[simp] theorem ofFn_zero {f : Fin 0 → α} : ofFn f = #[] := by
   simp [ofFn, ofFn.go]
 
@@ -32,12 +67,23 @@ theorem ofFn_succ {f : Fin (n+1) → α} :
     intro h₃
     simp only [show i = n by omega]
 
+theorem ofFn_add {n m} {f : Fin (n + m) → α} :
+    ofFn f = (ofFn (fun i => f (i.castLE (Nat.le_add_right n m)))) ++ (ofFn (fun i => f (i.natAdd n))) := by
+  induction m with
+  | zero => simp
+  | succ m ih => simp [ofFn_succ, ih]
+
 @[simp] theorem _root_.List.toArray_ofFn {f : Fin n → α} : (List.ofFn f).toArray = Array.ofFn f := by
   ext <;> simp
 
 @[simp] theorem toList_ofFn {f : Fin n → α} : (Array.ofFn f).toList = List.ofFn f := by
   apply List.ext_getElem <;> simp
 
+theorem ofFn_succ' {f : Fin (n+1) → α} :
+    ofFn f = #[f 0] ++ ofFn (fun i => f i.succ) := by
+  apply Array.toList_inj.mp
+  simp [List.ofFn_succ]
+
 @[simp]
 theorem ofFn_eq_empty_iff {f : Fin n → α} : ofFn f = #[] ↔ n = 0 := by
   rw [← Array.toList_inj]
@@ -52,4 +98,52 @@ theorem mem_ofFn {n} {f : Fin n → α} {a : α} : a ∈ ofFn f ↔ ∃ i, f i =
   · rintro ⟨i, rfl⟩
     apply mem_of_getElem (i := i) <;> simp
 
+/-! ### ofFnM -/
+
+/-- Construct (in a monadic context) an array by applying a monadic function to each index. -/
+def ofFnM {n} [Monad m] (f : Fin n → m α) : m (Array α) :=
+  Fin.foldlM n (fun xs i => xs.push <$> f i) (Array.mkEmpty n)
+
+@[simp]
+theorem ofFnM_zero [Monad m] {f : Fin 0 → m α} : ofFnM f = pure #[] := by
+  simp [ofFnM]
+
+theorem ofFnM_succ' {n} [Monad m] [LawfulMonad m] {f : Fin (n + 1) → m α} :
+    ofFnM f = (do
+      let a ← f 0
+      let as ← ofFnM fun i => f i.succ
+      pure (#[a] ++ as)) := by
+  simp [ofFnM, Fin.foldlM_eq_finRange_foldlM, List.foldlM_push_eq_append, List.finRange_succ, Function.comp_def]
+
+theorem ofFnM_succ {n} [Monad m] [LawfulMonad m] {f : Fin (n + 1) → m α} :
+    ofFnM f = (do
+      let as ← ofFnM fun i => f i.castSucc
+      let a  ← f (Fin.last n)
+      pure (as.push a)) := by
+  simp [ofFnM, Fin.foldlM_succ_last]
+
+theorem ofFnM_add {n m} [Monad m] [LawfulMonad m] {f : Fin (n + k) → m α} :
+    ofFnM f = (do
+      let as ← ofFnM fun i : Fin n => f (i.castLE (Nat.le_add_right n k))
+      let bs ← ofFnM fun i : Fin k => f (i.natAdd n)
+      pure (as ++ bs)) := by
+  induction k with
+  | zero => simp
+  | succ k ih =>
+    simp only [ofFnM_succ, Nat.add_eq, ih, Fin.castSucc_castLE, Fin.castSucc_natAdd, bind_pure_comp,
+      bind_assoc, bind_map_left, Fin.natAdd_last, map_bind, Functor.map_map]
+    congr 1
+    funext xs
+    congr 1
+    funext ys
+    congr 1
+    funext x
+    simp
+
+@[simp] theorem toList_ofFnM [Monad m] [LawfulMonad m] {f : Fin n → m α} :
+    toList <$> ofFnM f = List.ofFnM f := by
+  induction n with
+  | zero => simp
+  | succ n ih => simp [ofFnM_succ, List.ofFnM_succ_last, ← ih]
+
 end Array

src/Init/Data/Fin/Fold.lean

@@ -100,6 +100,11 @@ Fin.foldrM n f xₙ = do
 
 /-! ### foldlM -/
 
+@[congr] theorem foldlM_congr [Monad m] {n k : Nat} (w : n = k) (f : α → Fin n → m α) :
+    foldlM n f = foldlM k (fun x i => f x (i.cast w.symm)) := by
+  subst w
+  rfl
+
 theorem foldlM_loop_lt [Monad m] (f : α → Fin n → m α) (x) (h : i < n) :
     foldlM.loop n f x i = f x ⟨i, h⟩ >>= (foldlM.loop n f . (i+1)) := by
   rw [foldlM.loop, dif_pos h]
@@ -120,14 +125,49 @@ theorem foldlM_loop [Monad m] (f : α → Fin (n+1) → m α) (x) (h : i < n+1)
     rw [foldlM_loop_eq, foldlM_loop_eq]
 termination_by n - i
 
-@[simp] theorem foldlM_zero [Monad m] (f : α → Fin 0 → m α) (x) : foldlM 0 f x = pure x :=
-  foldlM_loop_eq ..
+@[simp] theorem foldlM_zero [Monad m] (f : α → Fin 0 → m α) : foldlM 0 f = pure := by
+  funext x
+  exact foldlM_loop_eq ..
+
+theorem foldlM_succ [Monad m] (f : α → Fin (n+1) → m α) :
+    foldlM (n+1) f = fun x => f x 0 >>= foldlM n (fun x j => f x j.succ) := by
+  funext x
+  exact foldlM_loop ..
 
-theorem foldlM_succ [Monad m] (f : α → Fin (n+1) → m α) (x) :
-    foldlM (n+1) f x = f x 0 >>= foldlM n (fun x j => f x j.succ) := foldlM_loop ..
+/-- Variant of `foldlM_succ` that splits off `Fin.last n` rather than `0`. -/
+theorem foldlM_succ_last [Monad m] [LawfulMonad m] (f : α → Fin (n+1) → m α) :
+    foldlM (n+1) f = fun x => foldlM n (fun x j => f x j.castSucc) x >>= (f · (Fin.last n)) := by
+  funext x
+  induction n generalizing x with
+  | zero =>
+    simp [foldlM_succ]
+  | succ n ih =>
+    rw [foldlM_succ]
+    conv => rhs; rw [foldlM_succ]
+    simp only [castSucc_zero, castSucc_succ, bind_assoc]
+    congr 1
+    funext x
+    rw [ih]
+    simp
+
+theorem foldlM_add [Monad m] [LawfulMonad m] (f : α → Fin (n + k) → m α) :
+    foldlM (n + k) f =
+      fun x => foldlM n (fun x i => f x (i.castLE (Nat.le_add_right n k))) x >>= foldlM k (fun x i => f x (i.natAdd n)) := by
+  induction k with
+  | zero =>
+    funext x
+    simp
+  | succ k ih =>
+    funext x
+    simp [foldlM_succ_last, ← Nat.add_assoc, ih]
 
 /-! ### foldrM -/
 
+@[congr] theorem foldrM_congr [Monad m] {n k : Nat} (w : n = k) (f : Fin n → α → m α) :
+    foldrM n f = foldrM k (fun i => f (i.cast w.symm)) := by
+  subst w
+  rfl
+
 theorem foldrM_loop_zero [Monad m] (f : Fin n → α → m α) (x) :
     foldrM.loop n f ⟨0, Nat.zero_le _⟩ x = pure x := by
   rw [foldrM.loop]
@@ -145,19 +185,46 @@ theorem foldrM_loop [Monad m] [LawfulMonad m] (f : Fin (n+1) → α → m α) (x
     conv => rhs; rw [←bind_pure (f 0 x)]
     congr
     funext
-    try simp only [foldrM.loop] -- the try makes this proof work with and without opaque wf rec
   | succ i ih =>
     rw [foldrM_loop_succ, foldrM_loop_succ, bind_assoc]
     congr; funext; exact ih ..
 
-@[simp] theorem foldrM_zero [Monad m] (f : Fin 0 → α → m α) (x) : foldrM 0 f x = pure x :=
-  foldrM_loop_zero ..
+@[simp] theorem foldrM_zero [Monad m] (f : Fin 0 → α → m α) : foldrM 0 f = pure := by
+  funext x
+  exact foldrM_loop_zero ..
 
-theorem foldrM_succ [Monad m] [LawfulMonad m] (f : Fin (n+1) → α → m α) (x) :
-    foldrM (n+1) f x = foldrM n (fun i => f i.succ) x >>= f 0 := foldrM_loop ..
+theorem foldrM_succ [Monad m] [LawfulMonad m] (f : Fin (n+1) → α → m α) :
+    foldrM (n+1) f = fun x => foldrM n (fun i => f i.succ) x >>= f 0 := by
+  funext x
+  exact foldrM_loop ..
+
+theorem foldrM_succ_last [Monad m] [LawfulMonad m] (f : Fin (n+1) → α → m α) :
+    foldrM (n+1) f = fun x => f (Fin.last n) x >>= foldrM n (fun i => f i.castSucc) := by
+  funext x
+  induction n generalizing x with
+  | zero => simp [foldrM_succ]
+  | succ n ih =>
+    rw [foldrM_succ]
+    conv => rhs; rw [foldrM_succ]
+    simp [ih]
+
+theorem foldrM_add [Monad m] [LawfulMonad m] (f : Fin (n + k) → α → m α) :
+    foldrM (n + k) f =
+      fun x => foldrM k (fun i => f (i.natAdd n)) x >>= foldrM n (fun i => f (i.castLE (Nat.le_add_right n k))) := by
+  induction k with
+  | zero =>
+    simp
+  | succ k ih =>
+    funext x
+    simp [foldrM_succ_last, ← Nat.add_assoc, ih]
 
 /-! ### foldl -/
 
+@[congr] theorem foldl_congr {n k : Nat} (w : n = k) (f : α → Fin n → α) :
+    foldl n f = foldl k (fun x i => f x (i.cast w.symm)) := by
+  subst w
+  rfl
+
 theorem foldl_loop_lt (f : α → Fin n → α) (x) (h : i < n) :
     foldl.loop n f x i = foldl.loop n f (f x ⟨i, h⟩) (i+1) := by
   rw [foldl.loop, dif_pos h]
@@ -187,14 +254,28 @@ theorem foldl_succ_last (f : α → Fin (n+1) → α) (x) :
   rw [foldl_succ]
   induction n generalizing x with
   | zero => simp [foldl_succ, Fin.last]
-  | succ n ih => rw [foldl_succ, ih (f · ·.succ), foldl_succ]; simp [succ_castSucc]
+  | succ n ih => rw [foldl_succ, ih (f · ·.succ), foldl_succ]; simp
+
+theorem foldl_add (f : α → Fin (n + m) → α) (x) :
+    foldl (n + m) f x =
+      foldl m (fun x i => f x (i.natAdd n))
+        (foldl n (fun x i => f x (i.castLE (Nat.le_add_right n m))) x):= by
+  induction m with
+  | zero => simp
+  | succ m ih => simp [foldl_succ_last, ih, ← Nat.add_assoc]
+
 
 theorem foldl_eq_foldlM (f : α → Fin n → α) (x) :
     foldl n f x = foldlM (m:=Id) n f x := by
   induction n generalizing x <;> simp [foldl_succ, foldlM_succ, *]
 
 /-! ### foldr -/
 
+@[congr] theorem foldr_congr {n k : Nat} (w : n = k) (f : Fin n → α → α) :
+    foldr n f = foldr k (fun i => f (i.cast w.symm)) := by
+  subst w
+  rfl
+
 theorem foldr_loop_zero (f : Fin n → α → α) (x) :
     foldr.loop n f 0 (Nat.zero_le _) x = x := by
   rw [foldr.loop]
@@ -220,7 +301,15 @@ theorem foldr_succ_last (f : Fin (n+1) → α → α) (x) :
     foldr (n+1) f x = foldr n (f ·.castSucc) (f (last n) x) := by
   induction n generalizing x with
   | zero => simp [foldr_succ, Fin.last]
-  | succ n ih => rw [foldr_succ, ih (f ·.succ), foldr_succ]; simp [succ_castSucc]
+  | succ n ih => rw [foldr_succ, ih (f ·.succ), foldr_succ]; simp
+
+theorem foldr_add (f : Fin (n + m) → α → α) (x) :
+    foldr (n + m) f x =
+      foldr n (fun i => f (i.castLE (Nat.le_add_right n m)))
+        (foldr m (fun i => f (i.natAdd n)) x) := by
+  induction m generalizing x with
+  | zero => simp
+  | succ m ih => simp [foldr_succ_last, ih, ← Nat.add_assoc]
 
 theorem foldr_eq_foldrM (f : Fin n → α → α) (x) :
     foldr n f x = foldrM (m:=Id) n f x := by