diff --git a/Mathlib/Data/Seq/Seq.lean b/Mathlib/Data/Seq/Seq.lean index 27908e64a27817..2aad9b2ca9b360 100644 --- a/Mathlib/Data/Seq/Seq.lean +++ b/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