feat: proofs for Batteries.Data.Array/Vector.Monadic

Batteries/Data/Array/Lemmas.lean

@@ -26,45 +26,164 @@ theorem forIn_eq_forIn_toList [Monad m]
 @[deprecated (since := "2024-09-09")] alias data_zipWith := toList_zipWith
 @[deprecated (since := "2024-08-13")] alias zipWith_eq_zipWith_data := data_zipWith
 
-/-! ### filter -/
-
-theorem size_filter_le (p : α → Bool) (l : Array α) :
-    (l.filter p).size ≤ l.size := by
-  simp only [← length_toList, toList_filter]
-  apply List.length_filter_le
-
 /-! ### flatten -/
 
 @[deprecated (since := "2024-09-09")] alias data_join := toList_flatten
 @[deprecated (since := "2024-08-13")] alias join_data := toList_flatten
 @[deprecated (since := "2024-10-15")] alias mem_join := mem_flatten
 
+end Array
+
 /-! ### indexOf? -/
 
-theorem indexOf?_toList [BEq α] {a : α} {l : Array α} :
-    l.toList.indexOf? a = (l.indexOf? a).map Fin.val := by
-  simpa using aux l 0
-where
-  aux (l : Array α) (i : Nat) :
-       ((l.toList.drop i).indexOf? a).map (·+i) = (indexOfAux l a i).map Fin.val := by
-    rw [indexOfAux]
-    if h : i < l.size then
-      rw [List.drop_eq_getElem_cons h, getElem_toList, List.indexOf?_cons]
-      if h' : l[i] == a then
-        simp [h, h']
-      else
-        simp [h, h', ←aux l (i+1), Function.comp_def, ←Nat.add_assoc, Nat.add_right_comm]
-    else
-      have h' : l.size ≤ i := Nat.le_of_not_lt h
-      simp [h, List.drop_of_length_le h', List.indexOf?]
-  termination_by l.size - i
+namespace Array
+
+theorem findIdx?_loop_eq_map_findFinIdx?_loop_val {xs : Array α} {p : α → Bool} {j} :
+    findIdx?.loop p xs j = (findFinIdx?.loop p xs j).map (·.val) := by
+  unfold findIdx?.loop
+  unfold findFinIdx?.loop
+  split <;> rename_i h
+  case isTrue =>
+    split
+    case isTrue => simp
+    case isFalse =>
+      have : xs.size - (j + 1) < xs.size - j := Nat.sub_succ_lt_self xs.size j h
+      rw [findIdx?_loop_eq_map_findFinIdx?_loop_val (j := j + 1)]
+  case isFalse => simp
+termination_by xs.size - j
+
+theorem findIdx?_eq_map_findFinIdx?_val {xs : Array α} {p : α → Bool} :
+    xs.findIdx? p = (xs.findFinIdx? p).map (·.val) := by
+  simp [findIdx?, findFinIdx?, findIdx?_loop_eq_map_findFinIdx?_loop_val]
+
+end Array
+
+namespace List
+
+theorem findFinIdx?_go_beq_eq_idxOfAux_toArray [BEq α]
+    {xs as : List α} {a : α} {i : Nat} {h} (w : as = xs.drop i) :
+    findFinIdx?.go (fun x => x == a) xs as i h =
+      xs.toArray.idxOfAux a i := by
+  unfold findFinIdx?.go
+  unfold Array.idxOfAux
+  split <;> rename_i b as
+  · simp at h
+    simp [h]
+  · simp at h
+    rw [dif_pos (by simp; omega)]
+    simp only [getElem_toArray]
+    erw [getElem_drop' (j := 0)]
+    simp only [← w, getElem_cons_zero]
+    have : xs.length - (i + 1) < xs.length - i := by omega
+    rw [findFinIdx?_go_beq_eq_idxOfAux_toArray]
+    rw [← drop_drop, ← w]
+    simp
+termination_by xs.length - i
+
+@[simp] theorem finIdxOf?_toArray [BEq α] {as : List α} {a : α} :
+    as.toArray.finIdxOf? a = as.finIdxOf? a := by
+  unfold Array.finIdxOf?
+  unfold finIdxOf?
+  unfold findFinIdx?
+  rw [findFinIdx?_go_beq_eq_idxOfAux_toArray]
+  simp
+
+theorem findIdx?_go_eq_map_findFinIdx?_go_val {xs : List α} {p : α → Bool} {i : Nat} {h} :
+    List.findIdx?.go p xs i =
+      (List.findFinIdx?.go p l xs i h).map (·.val) := by
+  unfold findIdx?.go
+  unfold findFinIdx?.go
+  split <;> rename_i a xs
+  · simp_all
+  · simp only
+    split
+    · simp
+    · rw [findIdx?_go_eq_map_findFinIdx?_go_val]
+
+theorem findIdx?_eq_map_findFinIdx?_val {xs : List α} {p : α → Bool} :
+    xs.findIdx? p = (xs.findFinIdx? p).map (·.val) := by
+  simp [findIdx?, findFinIdx?]
+  rw [findIdx?_go_eq_map_findFinIdx?_go_val]
+
+theorem idxOf?_eq_map_finIdxOf?_val [BEq α] {xs : List α} {a : α} :
+    xs.idxOf? a = (xs.finIdxOf? a).map (·.val) := by
+  simp [idxOf?, finIdxOf?, findIdx?_eq_map_findFinIdx?_val]
+
+@[simp] theorem idxOf?_toArray [BEq α] {as : List α} {a : α} :
+    as.toArray.idxOf? a = as.idxOf? a := by
+  rw [Array.idxOf?, finIdxOf?_toArray, idxOf?_eq_map_finIdxOf?_val]
+
+open Array in
+@[simp] theorem findIdx?_toArray {as : List α} {p : α → Bool} :
+    as.toArray.findIdx? p = as.findIdx? p := by
+  unfold Array.findIdx?
+  suffices ∀ i, i ≤ as.length →
+      Array.findIdx?.loop p as.toArray (as.length - i) =
+        (findIdx? p (as.drop  (as.length - i))).map fun j => j +  (as.length - i) by
+    specialize this as.length
+    simpa
+  intro i
+  induction i with
+  | zero => simp [findIdx?.loop]
+  | succ i ih =>
+    unfold findIdx?.loop
+    simp only [size_toArray, getElem_toArray]
+    split <;> rename_i h
+    · rw [drop_eq_getElem_cons h]
+      rw [findIdx?_cons]
+      split <;> rename_i h'
+      · simp
+      · intro w
+        have : as.length - (i + 1) + 1 = as.length - i := by omega
+        specialize ih (by omega)
+        simp only [Option.map_map, this, ih]
+        congr
+        ext
+        simp
+        omega
+    · have : as.length = 0 := by omega
+      simp_all
+
+@[simp] theorem findFinIdx?_toArray {as : List α} {p : α → Bool} :
+    as.toArray.findFinIdx? p = as.findFinIdx? p := by
+  have h := findIdx?_toArray (as := as) (p := p)
+  rw [findIdx?_eq_map_findFinIdx?_val, Array.findIdx?_eq_map_findFinIdx?_val] at h
+  rwa [Option.map_inj_right] at h
+  rintro ⟨x, hx⟩ ⟨y, hy⟩ rfl
+  simp
+
+end List
+
+open List
+
+namespace Array
+
+theorem idxOf?_toList [BEq α] {a : α} {l : Array α} :
+    l.toList.idxOf? a = l.idxOf? a := by
+  rcases l with ⟨l⟩
+  simp
 
 /-! ### erase -/
 
-@[simp] theorem toList_erase [BEq α] (l : Array α) (a : α) :
+@[simp] theorem eraseP_toArray {as : List α} {p : α → Bool} :
+    as.toArray.eraseP p = (as.eraseP p).toArray := by
+  rw [Array.eraseP, List.eraseP_eq_eraseIdx, findFinIdx?_toArray]
+  split
+  · simp only [List.findIdx?_eq_map_findFinIdx?_val, Option.map_none', *]
+  · simp only [eraseIdx_toArray, List.findIdx?_eq_map_findFinIdx?_val, Option.map_some', *]
+
+-- Adaptation note: We can remove the `LawfulBEq α` assumption again on nightly-2025-01-31.
+@[simp] theorem erase_toArray [BEq α] [LawfulBEq α] {as : List α} {a : α} :
+    as.toArray.erase a = (as.erase a).toArray := by
+  rw [Array.erase, finIdxOf?_toArray, List.erase_eq_eraseIdx]
+  rw [idxOf?_eq_map_finIdxOf?_val]
+  split <;> simp_all
+
+-- Adaptation note: We can remove the `LawfulBEq α` assumption again on nightly-2025-01-31.
+@[simp] theorem toList_erase [BEq α] [LawfulBEq α] (l : Array α) (a : α) :
     (l.erase a).toList = l.toList.erase a := by
-  simp only [erase, ← List.eraseIdx_indexOf_eq_erase, List.indexOf_eq_indexOf?, length_toList]
-  cases h : l.indexOf? a <;> simp [Array.indexOf?_toList, List.eraseIdx_of_length_le, *]
+  rcases l with ⟨l⟩
+  simp
 
 @[simp] theorem size_eraseIdxIfInBounds (a : Array α) (i : Nat) :
     (a.eraseIdxIfInBounds i).size = if i < a.size then a.size-1 else a.size := by

Batteries/Data/Array/Monadic.lean

@@ -5,6 +5,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Mario Carneiro, Gabriel Ebner
 -/
 import Batteries.Classes.SatisfiesM
+import Batteries.Util.ProofWanted
 
 /-!
 # Results about monadic operations on `Array`, in terms of `SatisfiesM`.
@@ -125,39 +126,158 @@ theorem SatisfiesM_foldrM [Monad m] [LawfulMonad m]
   simp [foldrM]; split; {exact go _ h0}
   · next h => exact .pure (Nat.eq_zero_of_not_pos h ▸ h0)
 
-theorem SatisfiesM_mapFinIdxM [Monad m] [LawfulMonad m] (as : Array α) (f : Fin as.size → α → m β)
+theorem SatisfiesM_mapFinIdxM [Monad m] [LawfulMonad m] (as : Array α)
+    (f : (i : Nat) → α → (h : i < as.size) → m β)
     (motive : Nat → Prop) (h0 : motive 0)
-    (p : Fin as.size → β → Prop)
-    (hs : ∀ i, motive i.1 → SatisfiesM (p i · ∧ motive (i + 1)) (f i as[i])) :
+    (p : (i : Nat) → β → (h : i < as.size) → Prop)
+    (hs : ∀ i h, motive i → SatisfiesM (p i · h ∧ motive (i + 1)) (f i as[i] h)) :
     SatisfiesM
-      (fun arr => motive as.size ∧ ∃ eq : arr.size = as.size, ∀ i h, p ⟨i, h⟩ arr[i])
+      (fun arr => motive as.size ∧ ∃ eq : arr.size = as.size, ∀ i h, p i arr[i] h)
       (Array.mapFinIdxM as f) := by
-  let rec go {bs i j h} (h₁ : j = bs.size) (h₂ : ∀ i h h', p ⟨i, h⟩ bs[i]) (hm : motive j) :
+  let rec go {bs i j h} (h₁ : j = bs.size) (h₂ : ∀ i h h', p i bs[i] h) (hm : motive j) :
     SatisfiesM
-      (fun arr => motive as.size ∧ ∃ eq : arr.size = as.size, ∀ i h, p ⟨i, h⟩ arr[i])
+      (fun arr => motive as.size ∧ ∃ eq : arr.size = as.size, ∀ i h, p i arr[i] h)
       (Array.mapFinIdxM.map as f i j h bs) := by
     induction i generalizing j bs with simp [mapFinIdxM.map]
     | zero =>
       have := (Nat.zero_add _).symm.trans h
       exact .pure ⟨this ▸ hm, h₁ ▸ this, fun _ _ => h₂ ..⟩
     | succ i ih =>
-      refine (hs _ (by exact hm)).bind fun b hb => ih (by simp [h₁]) (fun i hi hi' => ?_) hb.2
+      refine (hs _ _ (by exact hm)).bind fun b hb => ih (by simp [h₁]) (fun i hi hi' => ?_) hb.2
       simp at hi'; simp [getElem_push]; split
       · next h => exact h₂ _ _ h
       · next h => cases h₁.symm ▸ (Nat.le_or_eq_of_le_succ hi').resolve_left h; exact hb.1
   simp [mapFinIdxM]; exact go rfl nofun h0
 
 theorem SatisfiesM_mapIdxM [Monad m] [LawfulMonad m] (as : Array α) (f : Nat → α → m β)
     (motive : Nat → Prop) (h0 : motive 0)
-    (p : Fin as.size → β → Prop)
-    (hs : ∀ i, motive i.1 → SatisfiesM (p i · ∧ motive (i + 1)) (f i as[i])) :
+    (p : (i : Nat) → β → (h : i < as.size) → Prop)
+    (hs : ∀ i h, motive i → SatisfiesM (p i · h ∧ motive (i + 1)) (f i as[i])) :
     SatisfiesM
-      (fun arr => motive as.size ∧ ∃ eq : arr.size = as.size, ∀ i h, p ⟨i, h⟩ arr[i])
+      (fun arr => motive as.size ∧ ∃ eq : arr.size = as.size, ∀ i h, p i arr[i] h)
       (as.mapIdxM f) :=
-  SatisfiesM_mapFinIdxM as (fun i => f i) motive h0 p hs
+  SatisfiesM_mapFinIdxM as (fun i a _ => f i a) motive h0 p hs
+
+theorem size_mapFinIdxM [Monad m] [LawfulMonad m]
+    (as : Array α) (f : (i : Nat) → α → (h : i < as.size) → m β) :
+    SatisfiesM (fun arr => arr.size = as.size) (Array.mapFinIdxM as f) :=
+  (SatisfiesM_mapFinIdxM _ _ (fun _ => True) trivial (fun _ _ _ => True)
+    (fun _ _ _ => .of_true fun _ => ⟨trivial, trivial⟩)).imp (·.2.1)
+
+theorem size_mapIdxM [Monad m] [LawfulMonad m] (as : Array α) (f : Nat → α → m β) :
+    SatisfiesM (fun arr => arr.size = as.size) (Array.mapIdxM f as) :=
+  size_mapFinIdxM _ _
 
 theorem size_modifyM [Monad m] [LawfulMonad m] (a : Array α) (i : Nat) (f : α → m α) :
     SatisfiesM (·.size = a.size) (a.modifyM i f) := by
   unfold modifyM; split
   · exact .bind_pre <| .of_true fun _ => .pure <| by simp only [size_set]
   · exact .pure rfl
+
+end Array
+
+namespace List
+
+theorem filterM_toArray [Monad m] [LawfulMonad m] (l : List α) (p : α → m Bool) :
+    l.toArray.filterM p = toArray <$> l.filterM p := by
+  simp [Array.filterM, filterM]
+  conv => lhs; rw [← reverse_nil]
+  generalize [] = acc
+  induction l generalizing acc with simp
+  | cons x xs ih =>
+    congr; funext b
+    cases b
+    · simp; exact ih acc
+    · simp [← List.reverse_cons]; exact ih (x :: acc)
+
+theorem filterRevM_toArray [Monad m] [LawfulMonad m] (l : List α) (p : α → m Bool) :
+    l.toArray.filterRevM p = toArray <$> l.filterRevM p := by
+  simp [Array.filterRevM, filterRevM]
+  rw [← foldlM_reverse, ← foldlM_toArray, ← Array.filterM, filterM_toArray]
+  simp [filterM]
+
+theorem filterMapM_toArray [Monad m] [LawfulMonad m] (l : List α) (f : α → m (Option β)) :
+    l.toArray.filterMapM f = toArray <$> l.filterMapM f := by
+  simp [Array.filterMapM, filterMapM]
+  conv => lhs; rw [← reverse_nil]
+  generalize [] = acc
+  induction l generalizing acc with simp [filterMapM.loop]
+  | cons x xs ih =>
+    congr; funext o
+    cases o
+    · simp; exact ih acc
+    · simp; rw [← List.reverse_cons]; exact ih _
+
+-- This needs to wait until the definition of `flatMapM` is updated in `nightly-2025-02-01`.
+proof_wanted flatMapM_toArray [Monad m] [LawfulMonad m] (l : List α) (f : α → m (Array β)) :
+    l.toArray.flatMapM f = toArray <$> l.flatMapM (fun a => toList <$> f a)
+
+theorem mapFinIdxM_toArray [Monad m] [LawfulMonad m] (l : List α)
+    (f : (i : Nat) → α → (h : i < l.length) → m β) :
+    l.toArray.mapFinIdxM f = toArray <$> l.mapFinIdxM f := by
+  let rec go (i : Nat) (acc : Array β) (inv : i + acc.size = l.length) :
+      Array.mapFinIdxM.map l.toArray f i acc.size inv acc
+      = toArray <$> mapFinIdxM.go l f (l.drop acc.size) acc
+        (by simp [Nat.sub_add_cancel (Nat.le.intro (Nat.add_comm _ _ ▸ inv))]) := by
+    match i with
+    | 0 =>
+      simp at inv
+      simp [Array.mapFinIdxM.map, inv, mapFinIdxM.go]
+    | k + 1 =>
+      conv => enter [2, 2, 3]; rw [← getElem_cons_drop l acc.size (by omega)]
+      simp [Array.mapFinIdxM.map, mapFinIdxM.go]
+      congr; funext x
+      conv => enter [1, 4]; rw [← Array.size_push _ x]
+      conv => enter [2, 2, 3]; rw [← Array.size_push _ x]
+      refine go k (acc.push x) _
+  simp [Array.mapFinIdxM, mapFinIdxM]
+  exact go _ #[] _
+
+theorem mapIdxM_toArray [Monad m] [LawfulMonad m] (l : List α)
+    (f : Nat → α → m β) :
+    l.toArray.mapIdxM f = toArray <$> l.mapIdxM f := by
+  let rec go (bs : List α) (acc : Array β) (inv : bs.length + acc.size = l.length) :
+      mapFinIdxM.go l (fun i a h => f i a) bs acc inv = mapIdxM.go f bs acc := by
+    match bs with
+    | [] => simp [mapFinIdxM.go, mapIdxM.go]
+    | x :: xs => simp [mapFinIdxM.go, mapIdxM.go, go]
+  unfold Array.mapIdxM
+  rw [mapFinIdxM_toArray]
+  simp [mapFinIdxM, mapIdxM]
+  rw [go]
+
+end List
+
+namespace Array
+
+theorem toList_filterM [Monad m] [LawfulMonad m] (a : Array α) (p : α → m Bool) :
+    toList <$> a.filterM p = a.toList.filterM p := by
+  rw [List.filterM_toArray]
+  simp
+
+theorem toList_filterRevM [Monad m] [LawfulMonad m] (a : Array α) (p : α → m Bool) :
+    toList <$> a.filterRevM p = a.toList.filterRevM p := by
+  rw [List.filterRevM_toArray]
+  simp
+
+theorem toList_filterMapM [Monad m] [LawfulMonad m] (a : Array α) (f : α → m (Option β)) :
+    toList <$> a.filterMapM f = a.toList.filterMapM f := by
+  rw [List.filterMapM_toArray]
+  simp
+
+proof_wanted toList_flatMapM [Monad m] [LawfulMonad m] (a : Array α) (f : α → m (Array β)) :
+    toList <$> a.flatMapM f = a.toList.flatMapM (fun a => toList <$> f a)
+
+theorem toList_mapFinIdxM [Monad m] [LawfulMonad m] (l : Array α)
+    (f : (i : Nat) → α → (h : i < l.size) → m β) :
+    toList <$> l.mapFinIdxM f = l.toList.mapFinIdxM f := by
+  rw [List.mapFinIdxM_toArray]
+  simp
+
+theorem toList_mapIdxM [Monad m] [LawfulMonad m] (l : Array α)
+    (f : Nat → α → m β) :
+    toList <$> l.mapIdxM f = l.toList.mapIdxM f := by
+  rw [List.mapIdxM_toArray]
+  simp
+
+end Array

Batteries/Data/ByteArray.lean

@@ -53,7 +53,7 @@ theorem get_push_lt (a : ByteArray) (x : UInt8) (i : Nat) (h : i < a.size) :
   Array.size_set ..
 
 @[simp] theorem get_set_eq (a : ByteArray) (i : Fin a.size) (v : UInt8) : (a.set i v)[i.val] = v :=
-  Array.get_set_eq ..
+  Array.getElem_set_self _ _ _ _ (eq := rfl) _
 
 theorem get_set_ne (a : ByteArray) (i : Fin a.size) (v : UInt8) (hj : j < a.size) (h : i.val ≠ j) :
     (a.set i v)[j]'(a.size_set .. ▸ hj) = a[j] :=
@@ -76,7 +76,7 @@ theorem set_set (a : ByteArray) (i : Fin a.size) (v v' : UInt8) :
 
 @[simp] theorem data_append (a b : ByteArray) : (a ++ b).data = a.data ++ b.data := by
   rw [←append_eq]; simp [ByteArray.append, size]
-  rw [Array.extract_empty_of_stop_le_start (h:=Nat.le_add_right ..), Array.append_nil]
+  rw [Array.extract_empty_of_stop_le_start (h:=Nat.le_add_right ..), Array.append_empty]
 @[deprecated (since := "2024-08-13")] alias append_data := data_append
 
 theorem size_append (a b : ByteArray) : (a ++ b).size = a.size + b.size := by