From c0b3791156c4a2eb0b311e650f9eb86a4d31e366 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois=20G=2E=20Dorais?= <fgdorais@gmail.com>
Date: Tue, 15 Oct 2024 20:39:32 -0400
Subject: [PATCH] feat: add `Array.eraseIdx!` and some lemmas (#988)

---
 Batteries/Data/Array/Basic.lean  | 10 ++++++++++
 Batteries/Data/Array/Lemmas.lean |  7 +++++++
 2 files changed, 17 insertions(+)

diff --git a/Batteries/Data/Array/Basic.lean b/Batteries/Data/Array/Basic.lean
index 7fb0cf8cf8..c74b8f4750 100644
--- a/Batteries/Data/Array/Basic.lean
+++ b/Batteries/Data/Array/Basic.lean
@@ -172,6 +172,16 @@ greater than i.
 abbrev eraseIdxN (a : Array α) (i : Nat) (h : i < a.size := by get_elem_tactic) : Array α :=
   a.feraseIdx ⟨i, h⟩
 
+/--
+Remove the element at a given index from an array, panics if index is out of bounds.
+-/
+def eraseIdx! (a : Array α) (i : Nat) : Array α :=
+  if h : i < a.size then
+    a.feraseIdx ⟨i, h⟩
+  else
+    have : Inhabited (Array α) := ⟨a⟩
+    panic! s!"index {i} out of bounds"
+
 end Array
 
 
diff --git a/Batteries/Data/Array/Lemmas.lean b/Batteries/Data/Array/Lemmas.lean
index e6b1cf385a..124f8d412c 100644
--- a/Batteries/Data/Array/Lemmas.lean
+++ b/Batteries/Data/Array/Lemmas.lean
@@ -147,6 +147,13 @@ where
 @[simp] proof_wanted toList_erase [BEq α] {l : Array α} {a : α} :
     (l.erase a).toList = l.toList.erase a
 
+@[simp] theorem eraseIdx!_eq_eraseIdx (a : Array α) (i : Nat) :
+    a.eraseIdx! i = a.eraseIdx i := rfl
+
+@[simp] theorem size_eraseIdx (a : Array α) (i : Nat) :
+    (a.eraseIdx i).size = if i < a.size then a.size-1 else a.size := by
+  simp only [eraseIdx]; split; simp; rfl
+
 /-! ### shrink -/
 
 theorem size_shrink_loop (a : Array α) (n) : (shrink.loop n a).size = a.size - n := by