chore: upstream List.modify, add lemmas, relate to Array.modify

src/Init/Data/Array/Lemmas.lean

@@ -9,6 +9,7 @@ import Init.Data.List.Impl
 import Init.Data.List.Monadic
 import Init.Data.List.Range
 import Init.Data.List.Nat.TakeDrop
+import Init.Data.List.Nat.Modify
 import Init.Data.Array.Mem
 import Init.TacticsExtra
 
@@ -851,6 +852,12 @@ theorem getElem_modify {as : Array α} {x i} (h : i < (as.modify x f).size) :
   · simp only [Id.bind_eq, get_set _ _ _ (by simpa using h)]; split <;> simp [*]
   · rw [if_neg (mt (by rintro rfl; exact h) (by simp_all))]
 
+@[simp] theorem toList_modify (as : Array α) (f : α → α) :
+    (as.modify x f).toList = as.toList.modify f x := by
+  apply List.ext_getElem
+  · simp
+  · simp [getElem_modify, List.getElem_modify]
+
 theorem getElem_modify_self {as : Array α} {i : Nat} (f : α → α) (h : i < (as.modify i f).size) :
     (as.modify i f)[i] = f (as[i]'(by simpa using h)) := by
   simp [getElem_modify h]
@@ -1433,6 +1440,11 @@ theorem all_toArray (p : α → Bool) (l : List α) : l.toArray.all p = l.all p
   apply ext'
   simp
 
+@[simp] theorem modify_toArray (f : α → α) (l : List α) :
+    l.toArray.modify i f = (l.modify f i).toArray := by
+  apply ext'
+  simp
+
 @[simp] theorem filter_toArray' (p : α → Bool) (l : List α) (h : stop = l.toArray.size) :
     l.toArray.filter p 0 stop = (l.filter p).toArray := by
   subst h

src/Init/Data/List/Basic.lean

@@ -38,7 +38,7 @@ The operations are organized as follow:
 * Sublists: `take`, `drop`, `takeWhile`, `dropWhile`, `partition`, `dropLast`,
   `isPrefixOf`, `isPrefixOf?`, `isSuffixOf`, `isSuffixOf?`, `Subset`, `Sublist`,
   `rotateLeft` and `rotateRight`.
-* Manipulating elements: `replace`, `insert`, `erase`, `eraseP`, `eraseIdx`.
+* Manipulating elements: `replace`, `insert`, `modify`, `erase`, `eraseP`, `eraseIdx`.
 * Finding elements: `find?`, `findSome?`, `findIdx`, `indexOf`, `findIdx?`, `indexOf?`,
  `countP`, `count`, and `lookup`.
 * Logic: `any`, `all`, `or`, and `and`.
@@ -1119,6 +1119,35 @@ theorem replace_cons [BEq α] {a : α} :
 @[inline] protected def insert [BEq α] (a : α) (l : List α) : List α :=
   if l.elem a then l else a :: l
 
+/-! ### modify -/
+
+/--
+Apply a function to the nth tail of `l`. Returns the input without
+using `f` if the index is larger than the length of the List.
+```
+modifyTailIdx f 2 [a, b, c] = [a, b] ++ f [c]
+```
+-/
+@[simp] def modifyTailIdx (f : List α → List α) : Nat → List α → List α
+  | 0, l => f l
+  | _+1, [] => []
+  | n+1, a :: l => a :: modifyTailIdx f n l
+
+/-- Apply `f` to the head of the list, if it exists. -/
+@[inline] def modifyHead (f : α → α) : List α → List α
+  | [] => []
+  | a :: l => f a :: l
+
+@[simp] theorem modifyHead_nil (f : α → α) : [].modifyHead f = [] := by rw [modifyHead]
+@[simp] theorem modifyHead_cons (a : α) (l : List α) (f : α → α) :
+    (a :: l).modifyHead f = f a :: l := by rw [modifyHead]
+
+/--
+Apply `f` to the nth element of the list, if it exists, replacing that element with the result.
+-/
+def modify (f : α → α) : Nat → List α → List α :=
+  modifyTailIdx (modifyHead f)
+
 /-! ### erase -/
 
 /--

src/Init/Data/List/Impl.lean

@@ -38,7 +38,7 @@ The following operations were already given `@[csimp]` replacements in `Init/Dat
 
 The following operations are given `@[csimp]` replacements below:
 `set`, `filterMap`, `foldr`, `append`, `bind`, `join`,
-`take`, `takeWhile`, `dropLast`, `replace`, `erase`, `eraseIdx`, `zipWith`,
+`take`, `takeWhile`, `dropLast`, `replace`, `modify`, `erase`, `eraseIdx`, `zipWith`,
 `enumFrom`, and `intercalate`.
 
 -/
@@ -197,6 +197,24 @@ The following operations are given `@[csimp]` replacements below:
     · simp [*]
     · intro h; rw [IH] <;> simp_all
 
+/-! ### modify -/
+
+/-- Tail-recursive version of `modify`. -/
+def modifyTR (f : α → α) (n : Nat) (l : List α) : List α := go l n #[] where
+  /-- Auxiliary for `modifyTR`: `modifyTR.go f l n acc = acc.toList ++ modify f n l`. -/
+  go : List α → Nat → Array α → List α
+  | [], _, acc => acc.toList
+  | a :: l, 0, acc => acc.toListAppend (f a :: l)
+  | a :: l, n+1, acc => go l n (acc.push a)
+
+theorem modifyTR_go_eq : ∀ l n, modifyTR.go f l n acc = acc.toList ++ modify f n l
+  | [], n => by cases n <;> simp [modifyTR.go, modify]
+  | a :: l, 0 => by simp [modifyTR.go, modify]
+  | a :: l, n+1 => by simp [modifyTR.go, modify, modifyTR_go_eq l]
+
+@[csimp] theorem modify_eq_modifyTR : @modify = @modifyTR := by
+  funext α f n l; simp [modifyTR, modifyTR_go_eq]
+
 /-! ### erase -/
 
 /-- Tail recursive version of `List.erase`. -/

src/Init/Data/List/Nat.lean

@@ -13,3 +13,4 @@ import Init.Data.List.Nat.Count
 import Init.Data.List.Nat.Erase
 import Init.Data.List.Nat.Find
 import Init.Data.List.Nat.BEq
+import Init.Data.List.Nat.Modify

src/Init/Data/List/Nat/Modify.lean

@@ -0,0 +1,102 @@
+/-
+Copyright (c) 2014 Parikshit Khanna. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Parikshit Khanna, Jeremy Avigad, Leonardo de Moura, Floris van Doorn, Mario Carneiro
+-/
+
+prelude
+import Init.Data.List.Nat.TakeDrop
+
+namespace List
+
+/-! ### modifyHead -/
+
+@[simp] theorem modifyHead_modifyHead (l : List α) (f g : α → α) :
+    (l.modifyHead f).modifyHead g = l.modifyHead (g ∘ f) := by cases l <;> simp [modifyHead]
+
+/-! ### modify -/
+
+@[simp] theorem modify_nil (f : α → α) (n) : [].modify f n = [] := by cases n <;> rfl
+
+@[simp] theorem modify_zero_cons (f : α → α) (a : α) (l : List α) :
+    (a :: l).modify f 0 = f a :: l := rfl
+
+@[simp] theorem modify_succ_cons (f : α → α) (a : α) (l : List α) (n) :
+    (a :: l).modify f (n + 1) = a :: l.modify f n := by rfl
+
+theorem modifyTailIdx_id : ∀ n (l : List α), l.modifyTailIdx id n = l
+  | 0, _ => rfl
+  | _+1, [] => rfl
+  | n+1, a :: l => congrArg (cons a) (modifyTailIdx_id n l)
+
+theorem eraseIdx_eq_modifyTailIdx : ∀ n (l : List α), eraseIdx l n = modifyTailIdx tail n l
+  | 0, l => by cases l <;> rfl
+  | _+1, [] => rfl
+  | _+1, _ :: _ => congrArg (cons _) (eraseIdx_eq_modifyTailIdx _ _)
+
+theorem getElem?_modify (f : α → α) :
+    ∀ n (l : List α) m, (modify f n l)[m]? = (fun a => if n = m then f a else a) <$> l[m]?
+  | n, l, 0 => by cases l <;> cases n <;> simp
+  | n, [], _+1 => by cases n <;> rfl
+  | 0, _ :: l, m+1 => by cases h : l[m]? <;> simp [h, modify, m.succ_ne_zero.symm]
+  | n+1, a :: l, m+1 => by
+    simp only [modify_succ_cons, getElem?_cons_succ, Nat.reduceEqDiff, Option.map_eq_map]
+    refine (getElem?_modify f n l m).trans ?_
+    cases h' : l[m]? <;> by_cases h : n = m <;>
+      simp [h, if_pos, if_neg, Option.map, mt Nat.succ.inj, not_false_iff, h']
+
+@[simp] theorem length_modifyTailIdx (f : List α → List α) (H : ∀ l, length (f l) = length l) :
+    ∀ n l, length (modifyTailIdx f n l) = length l
+  | 0, _ => H _
+  | _+1, [] => rfl
+  | _+1, _ :: _ => congrArg (·+1) (length_modifyTailIdx _ H _ _)
+
+theorem modifyTailIdx_add (f : List α → List α) (n) (l₁ l₂ : List α) :
+    modifyTailIdx f (l₁.length + n) (l₁ ++ l₂) = l₁ ++ modifyTailIdx f n l₂ := by
+  induction l₁ <;> simp [*, Nat.succ_add]
+
+@[simp] theorem length_modify (f : α → α) : ∀ n l, length (modify f n l) = length l :=
+  length_modifyTailIdx _ fun l => by cases l <;> rfl
+
+@[simp] theorem getElem?_modify_eq (f : α → α) (n) (l : List α) :
+    (modify f n l)[n]? = f <$> l[n]? := by
+  simp only [getElem?_modify, if_pos]
+
+@[simp] theorem getElem?_modify_ne (f : α → α) {m n} (l : List α) (h : m ≠ n) :
+    (modify f m l)[n]? = l[n]? := by
+  simp only [getElem?_modify, if_neg h, id_map']
+
+theorem getElem_modify (f : α → α) (n) (l : List α) (m) (h : m < (modify f n l).length) :
+    (modify f n l)[m] =
+      if n = m then f (l[m]'(by simp at h; omega)) else l[m]'(by simp at h; omega) := by
+  rw [getElem_eq_iff, getElem?_modify]
+  simp at h
+  simp [h]
+
+theorem modifyTailIdx_eq_take_drop (f : List α → List α) (H : f [] = []) :
+    ∀ n l, modifyTailIdx f n l = take n l ++ f (drop n l)
+  | 0, _ => rfl
+  | _ + 1, [] => H.symm
+  | n + 1, b :: l => congrArg (cons b) (modifyTailIdx_eq_take_drop f H n l)
+
+theorem modify_eq_take_drop (f : α → α) :
+    ∀ n l, modify f n l = take n l ++ modifyHead f (drop n l) :=
+  modifyTailIdx_eq_take_drop _ rfl
+
+theorem modify_eq_take_cons_drop (f : α → α) {n l} (h : n < length l) :
+    modify f n l = take n l ++ f l[n] :: drop (n + 1) l := by
+  rw [modify_eq_take_drop, drop_eq_getElem_cons h]; rfl
+
+theorem exists_of_modifyTailIdx (f : List α → List α) {n} {l : List α} (h : n ≤ l.length) :
+    ∃ l₁ l₂, l = l₁ ++ l₂ ∧ l₁.length = n ∧ modifyTailIdx f n l = l₁ ++ f l₂ :=
+  have ⟨_, _, eq, hl⟩ : ∃ l₁ l₂, l = l₁ ++ l₂ ∧ l₁.length = n :=
+    ⟨_, _, (take_append_drop n l).symm, length_take_of_le h⟩
+  ⟨_, _, eq, hl, hl ▸ eq ▸ modifyTailIdx_add (n := 0) ..⟩
+
+theorem exists_of_modify (f : α → α) {n} {l : List α} (h : n < l.length) :
+    ∃ l₁ a l₂, l = l₁ ++ a :: l₂ ∧ l₁.length = n ∧ modify f n l = l₁ ++ f a :: l₂ :=
+  match exists_of_modifyTailIdx _ (Nat.le_of_lt h) with
+  | ⟨_, _::_, eq, hl, H⟩ => ⟨_, _, _, eq, hl, H⟩
+  | ⟨_, [], eq, hl, _⟩ => nomatch Nat.ne_of_gt h (eq ▸ append_nil _ ▸ hl)
+
+end List