From 1b8f1f140c56d214b2171abe8675fb8582670fef Mon Sep 17 00:00:00 2001
From: Kim Morrison <477956+kim-em@users.noreply.github.com>
Date: Sat, 16 May 2026 16:28:43 +1000
Subject: [PATCH] feat: add ByteArray indexing and set! lemmas (#13457)

This PR adds the missing `ByteArray` push and `set!` lemmas that are
still carried locally in `ZipForStd.ByteArray` downstream.

These lemmas provide one coherent API layer for reasoning about
`ByteArray` pushes and `set!` updates on top of the existing upstream
append, extract, and size lemmas. Previous downstream code had to keep
local compatibility lemmas for these common proof steps.

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 src/Init/Data/ByteArray/Lemmas.lean | 77 +++++++++++++++++++++++++++++
 1 file changed, 77 insertions(+)

diff --git a/src/Init/Data/ByteArray/Lemmas.lean b/src/Init/Data/ByteArray/Lemmas.lean
index 02cc3efb3d..f8db4b648c 100644
--- a/src/Init/Data/ByteArray/Lemmas.lean
+++ b/src/Init/Data/ByteArray/Lemmas.lean
@@ -328,4 +328,81 @@ theorem extract_eq_extract_iff_getElem {as bs : ByteArray} {i j len : Nat}
     · exact h k hk'
     · exact (by omega : k = len) ▸ h'
 
+theorem getElem!_push_lt (data : ByteArray) (b : UInt8) (i : Nat) (hi : i < data.size) :
+    (data.push b)[i]! = data[i]! := by
+  have hi' : i < (data.push b).size := by
+    simp only [ByteArray.size_push]
+    omega
+  rw [getElem!_pos (data.push b) i hi', getElem!_pos data i hi]
+  exact Array.getElem_push_lt hi
+
+@[simp] theorem getElem!_push_eq (data : ByteArray) (b : UInt8) :
+    (data.push b)[data.size]! = b := by
+  have h : data.size < (data.push b).size := by
+    simp only [ByteArray.size_push]
+    omega
+  rw [getElem!_pos (data.push b) data.size h]
+  exact Array.getElem_push_eq
+
+@[grind =] theorem getElem!_push (data : ByteArray) (b : UInt8) (i : Nat) :
+    (data.push b)[i]! = if i = data.size then b else data[i]! := by
+  split
+  · subst i
+    exact getElem!_push_eq data b
+  · by_cases hi : i < data.size
+    · exact getElem!_push_lt data b i hi
+    · rw [getElem!_neg data i hi,
+        getElem!_neg (data.push b) i (by simp only [ByteArray.size_push]; omega)]
+
+private theorem getElem!_eq_data_getElem! (data : ByteArray) (i : Nat) :
+    data[i]! = data.data[i]! := by
+  by_cases h : i < data.size
+  · rw [getElem!_pos data i h, getElem!_pos data.data i h]
+    rfl
+  · rw [getElem!_neg data i h, getElem!_neg data.data i h]
+
+@[simp, grind =] theorem size_set! (data : ByteArray) (i : Nat) (v : UInt8) :
+    (data.set! i v).size = data.size := by
+  show (data.data.setIfInBounds i v).size = data.data.size
+  exact Array.size_setIfInBounds ..
+
+@[simp] theorem getElem!_set!_self (data : ByteArray) (i : Nat) (v : UInt8) (h : i < data.size) :
+    (data.set! i v)[i]! = v := by
+  rw [getElem!_eq_data_getElem!]
+  show (data.data.set! i v)[i]! = v
+  simp only [Array.set!_eq_setIfInBounds, Array.getElem!_eq_getD, Array.getD_eq_getD_getElem?,
+    Array.getElem?_setIfInBounds_self_of_lt h, Option.getD_some]
+
+@[simp] theorem getElem!_set!_ne (data : ByteArray) (i j : Nat) (v : UInt8) (hij : i ≠ j) :
+    (data.set! i v)[j]! = data[j]! := by
+  rw [getElem!_eq_data_getElem!, getElem!_eq_data_getElem!]
+  show (data.data.set! i v)[j]! = data.data[j]!
+  simp only [Array.set!_eq_setIfInBounds, Array.getElem!_eq_getD, Array.getD_eq_getD_getElem?,
+    Array.getElem?_setIfInBounds_ne hij]
+
+@[grind =] theorem getElem!_set! (data : ByteArray) (i : Nat) (v : UInt8) (j : Nat) (h : i < data.size) :
+    (data.set! i v)[j]! = if i = j then v else data[j]! := by
+  split
+  · next hij => subst hij; exact getElem!_set!_self data i v h
+  · next hij => exact getElem!_set!_ne data i j v hij
+
+@[simp] theorem getElem_set!_ne (data : ByteArray) (i j : Nat) (v : UInt8) (hij : i ≠ j)
+    (hj : j < data.size) :
+    (data.set! i v)[j]'(by rw [size_set!]; exact hj) = data[j] := by
+  rw [← getElem!_pos (data.set! i v) j (by rw [size_set!]; exact hj),
+    ← getElem!_pos data j hj,
+    getElem!_set!_ne _ _ _ _ hij]
+
+@[simp] theorem getElem_set!_self (data : ByteArray) (i : Nat) (v : UInt8) (h : i < data.size) :
+    (data.set! i v)[i]'(by rw [size_set!]; exact h) = v := by
+  rw [← getElem!_pos (data.set! i v) i (by rw [size_set!]; exact h),
+    getElem!_set!_self _ _ _ h]
+
+@[grind =] theorem getElem_set! (data : ByteArray) (i j : Nat) (v : UInt8) (h : i < data.size)
+    (hj : j < data.size) :
+    (data.set! i v)[j]'(by rw [size_set!]; exact hj) = if i = j then v else data[j] := by
+  split
+  · next hij => subst hij; exact getElem_set!_self data i v h
+  · next hij => exact getElem_set!_ne data i j v hij hj
+
 end ByteArray