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.

---------

Co-authored-by: Eric Wieser <[email protected]>

Author: kim-em

src/Init/Data/Array/Basic.lean

@@ -112,6 +112,10 @@ theorem mem_def {a : α} {as : Array α} : a ∈ as ↔ a ∈ as.toList :=
   rw [Array.mem_def, ← getElem_toList]
   apply List.getElem_mem
 
+@[simp] theorem emptyWithCapacity_eq {α n} : @emptyWithCapacity α n = #[] := rfl
+
+@[simp] theorem mkEmpty_eq {α n} : @mkEmpty α n = #[] := rfl
+
 end Array
 
 namespace List
@@ -334,6 +338,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

@@ -61,11 +61,6 @@ theorem toArray_eq : List.toArray as = xs ↔ as = xs.toList := by
 
 @[grind] theorem size_empty : (#[] : Array α).size = 0 := rfl
 
-@[simp] theorem emptyWithCapacity_eq {α n} : @emptyWithCapacity α n = #[] := rfl
-
-@[deprecated emptyWithCapacity_eq (since := "2025-03-12")]
-theorem mkEmpty_eq {α n} : @mkEmpty α n = #[] := rfl
-
 /-! ### size -/
 
 @[grind →] theorem eq_empty_of_size_eq_zero (h : xs.size = 0) : xs = #[] := by
@@ -247,6 +242,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 +3267,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 +3379,32 @@ 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
+  simpa using List.foldl_push_eq_append (f := id)
+
+@[deprecated _root_.List.foldl_push_eq_append' (since := "2025-05-18")]
+theorem _root_.List.foldl_push {l : List α} {as : Array α} : l.foldl Array.push as = as ++ l.toArray := by
+  induction l generalizing as <;> simp [*]
+
+@[deprecated _root_.List.foldr_push_eq_append' (since := "2025-05-18")]
+theorem _root_.List.foldr_push {l : List α} {as : Array α} : l.foldr (fun a bs => push bs a) as = as ++ l.reverse.toArray := by
+  rw [List.foldr_eq_foldl_reverse, List.foldl_push_eq_append']
+
 @[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

@@ -25,6 +25,11 @@ open Nat
 
 /-! ## Monadic operations -/
 
+@[simp] theorem map_toList_inj [Monad m] [LawfulMonad m]
+    {xs : m (Array α)} {ys : m (Array α)} :
+   toList <$> xs = toList <$> ys ↔ xs = ys :=
+  _root_.map_inj_right (by simp)
+
 /-! ### mapM -/
 
 @[simp] theorem mapM_pure [Monad m] [LawfulMonad m] {xs : Array α} {f : α → β} :
@@ -34,6 +39,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

@@ -8,7 +8,9 @@ module
 prelude
 import all Init.Data.Array.Basic
 import Init.Data.Array.Lemmas
+import Init.Data.Array.Monadic
 import Init.Data.List.OfFn
+import Init.Data.List.FinRange
 
 /-!
 # Theorems about `Array.ofFn`
@@ -19,6 +21,8 @@ set_option linter.indexVariables true -- Enforce naming conventions for index va
 
 namespace Array
 
+/-! ### ofFn -/
+
 @[simp] theorem ofFn_zero {f : Fin 0 → α} : ofFn f = #[] := by
   simp [ofFn, ofFn.go]
 
@@ -32,12 +36,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 +67,71 @@ 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.emptyWithCapacity 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]
+
+@[simp]
+theorem ofFnM_pure_comp [Monad m] [LawfulMonad m] {n} {f : Fin n → α} :
+    ofFnM (pure ∘ f) = (pure (ofFn f) : m (Array α)) := by
+  apply Array.map_toList_inj.mp
+  simp
+
+-- Variant of `ofFnM_pure_comp` using a lambda.
+-- This is not marked a `@[simp]` as it would match on every occurrence of `ofFnM`.
+theorem ofFnM_pure [Monad m] [LawfulMonad m] {n} {f : Fin n → α} :
+    ofFnM (fun i => pure (f i)) = (pure (ofFn f) : m (Array α)) :=
+  ofFnM_pure_comp
+
+@[simp, grind =] theorem idRun_ofFnM {f : Fin n → Id α} :
+    Id.run (ofFnM f) = ofFn (fun i => Id.run (f i)) := by
+  unfold Id.run
+  induction n with
+  | zero => simp
+  | succ n ih => simp [ofFnM_succ', ofFn_succ', ih]
+
 end Array