feat(Data/Seq): modify and set operations for Seq

Mathlib/Data/Seq/Seq.lean

@@ -714,6 +714,21 @@ def fold (s : Seq α) (init : β) (f : β → α → β) : Seq β :=
     | some (x, s) => .some (f acc x, f acc x, s)
   cons init <| corec f (init, s)
 
+/-- Applies `f` to the `n`th element of the sequence, if it exists, replacing that element
+with the result. -/
+def modify (s : Seq α) (n : ℕ) (f : α → α) : Seq α where
+  val := Function.update s.val n ((s.val n).map f)
+  property := by
+    have (i : ℕ) : (Function.update s.val n ((s.get? n).map f)) i = none ↔ s.get? i = none := by
+      by_cases hi : i = n <;> simp [Function.update, hi]
+    simp [IsSeq, this]
+    exact @s.prop
+
+/-- Sets the value of sequence `s` at index `n` to `a`. If the `n`th element does not exist
+(`s` terminates earlier), the sequence is left unchanged. -/
+def set (s : Seq α) (n : ℕ) (a : α) : Seq α :=
+  modify s n (fun _ ↦ a)
+
 section OfStream
 
 @[simp]
@@ -1243,6 +1258,60 @@ theorem fold_head (init : β) (f : β → α → β) (s : Seq α) :
 
 end Fold
 
+section Modify
+
+@[simp]
+theorem modify_nil {f : α → α} {n} : modify nil n f = nil := by
+  ext1 m
+  simp [modify, Function.update]
+
+@[simp]
+theorem set_nil {n : ℕ} {x : α} : set nil n x = nil := modify_nil
+
+@[simp]
+theorem modify_cons_zero {f : α → α} {hd : α} {tl : Seq α} :
+    (cons hd tl).modify 0 f = cons (f hd) tl := by
+  ext1 n
+  cases n <;> simp [modify]
+  rfl
+
+@[simp]
+theorem set_cons_zero {hd hd' : α} {tl : Seq α} :
+    (cons hd tl).set 0 hd' = cons hd' tl := modify_cons_zero
+
+@[simp]
+theorem modify_cons_succ {hd : α} {f : α → α} {n : ℕ} {tl : Seq α} :
+    (cons hd tl).modify (n + 1) f = cons hd (tl.modify n f) := by
+  ext1 n
+  cases n <;> simp [modify, Function.update] <;> rfl
+
+@[simp]
+theorem set_cons_succ {hd x : α} {n : ℕ} {tl : Seq α} :
+    (cons hd tl).set (n + 1) x = cons hd (tl.set n x) := modify_cons_succ
+
+theorem set_get_of_not_terminated {s : Seq α} {x : α} {n : ℕ}
+    (h_not_terminated : ¬ s.TerminatedAt n) :
+    (s.set n x).get? n = x := by
+  simpa [set, modify, ← Option.ne_none_iff_exists'] using h_not_terminated
+
+theorem set_get_of_terminated {s : Seq α} {x : α} {n : ℕ}
+    (h_terminated : s.TerminatedAt n) :
+    (s.set n x).get? n = .none := by
+  simpa [set, modify] using h_terminated
+
+theorem set_get_stable {s : Seq α} {x : α} {n m : ℕ}
+    (h : n ≠ m) :
+    (s.set m x).get? n = s.get? n := by
+  simp [set, modify, Function.update, h]
+
+theorem set_dropn_stable_of_lt {s : Seq α} {m n : ℕ} {x : α}
+    (h : m < n) :
+    (s.set m x).drop n = s.drop n := by
+  ext1 i
+  simp [set_get_stable (show n + i ≠ m by omega)]
+
+end Modify
+
 instance : Functor Seq where map := @map
 
 instance : LawfulFunctor Seq where