chore: generalize Array/Vector.extract_push

src/Init/Data/Array/Extract.lean

@@ -46,15 +46,45 @@ theorem size_extract_of_le {as : Array α} {i j : Nat} (h : j ≤ as.size) :
   simp
   omega
 
-@[simp, grind =]
-theorem extract_push {as : Array α} {b : α} {start stop : Nat} (h : stop ≤ as.size) :
+@[grind =]
+theorem extract_push {as : Array α} {b : α} {start stop : Nat} :
+    (as.push b).extract start stop =
+      if stop ≤ as.size then
+        as.extract start stop
+      else if start ≤ as.size then
+        (as.extract start as.size).push b
+      else #[] := by
+  split
+  · ext i h₁ h₂
+    · simp
+      omega
+    · simp only [size_extract, size_push] at h₁ h₂
+      simp only [getElem_extract, getElem_push]
+      rw [dif_pos (by omega)]
+  · split
+    · ext i h₁ h₂
+      · simp
+        omega
+      · simp only [size_extract, size_push] at h₁ h₂
+        simp only [getElem_extract, getElem_push]
+        split <;> rename_i h₃
+        · split
+          · rfl
+          · simp_all
+            omega
+        · split <;> rename_i h₄
+          · simp at h₄
+            omega
+          · rfl
+    · ext i h₁ h₂
+      · simp
+        omega
+      · simp at h₂
+
+@[simp]
+theorem extract_push_of_le {as : Array α} {b : α} {start stop : Nat} (h : stop ≤ as.size) :
     (as.push b).extract start stop = as.extract start stop := by
-  ext i h₁ h₂
-  · simp
-    omega
-  · simp only [size_extract, size_push] at h₁ h₂
-    simp only [getElem_extract, getElem_push]
-    rw [dif_pos (by omega)]
+  rw [extract_push, if_pos h]
 
 @[simp, grind =]
 theorem extract_eq_pop {as : Array α} {stop : Nat} (h : stop = as.size - 1) :

src/Init/Data/Vector/Extract.lean

@@ -28,8 +28,22 @@ set_option linter.indexVariables false
   rcases xs with ⟨as, rfl⟩
   simp [h]
 
-@[simp, grind =]
-theorem extract_push {xs : Vector α n} {b : α} {start stop : Nat} (h : stop ≤ n) :
+@[grind =]
+theorem extract_push {α} {xs : Vector α n} {b : α} {start stop : Nat} :
+    (xs.push b).extract start stop =
+      if h₁ : stop ≤ n then
+        (xs.extract start stop).cast (by omega)
+      else if h₂ : start ≤ n then
+        ((xs.extract start n).push b).cast (by omega)
+      else #v[].cast (by omega) := by
+  rcases xs with ⟨xs, rfl⟩
+  simp [Array.extract_push]
+  split
+  · simp
+  · split <;> simp
+
+@[simp]
+theorem extract_push_of_le {xs : Vector α n} {b : α} {start stop : Nat} (h : stop ≤ n) :
     (xs.push b).extract start stop = (xs.extract start stop).cast (by omega) := by
   rcases xs with ⟨xs, rfl⟩
   simp [h]