From bfb73c4a5e5d6050c65c86b5b5e810f9e9b424ef Mon Sep 17 00:00:00 2001
From: Tobias Grosser <github@grosser.es>
Date: Tue, 1 Oct 2024 06:37:51 +0100
Subject: [PATCH] =?UTF-8?q?feat:=20add=20BitVec.getElem=5F[add|add=5Fadd?=
 =?UTF-8?q?=5Fbool|mul|rotateLeft|rotateRight=E2=80=A6=20(#5508)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

…|twoPow|one|replicate]

... and mark `getElem_setWidth` as `@[simp]`.

`getElem_rotateLeft` and `getElem_rotateRight` have a non-trivial rhs
but we follow `getLsbD_[rotateLeft|rotateRight]`for consistency.

---------

Co-authored-by: Kim Morrison <scott@tqft.net>
---
 src/Init/Data/BitVec/Bitblast.lean | 15 +++++++++++++++
 src/Init/Data/BitVec/Lemmas.lean   | 29 +++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/src/Init/Data/BitVec/Bitblast.lean b/src/Init/Data/BitVec/Bitblast.lean
index 710605d27d..7472e49272 100644
--- a/src/Init/Data/BitVec/Bitblast.lean
+++ b/src/Init/Data/BitVec/Bitblast.lean
@@ -164,6 +164,17 @@ theorem getLsbD_add {i : Nat} (i_lt : i < w) (x y : BitVec w) :
       (getLsbD x i ^^ (getLsbD y i ^^ carry i x y false)) := by
   simpa using getLsbD_add_add_bool i_lt x y false
 
+theorem getElem_add_add_bool {i : Nat} (i_lt : i < w) (x y : BitVec w) (c : Bool) :
+    (x + y + setWidth w (ofBool c))[i] =
+      (x[i] ^^ (y[i] ^^ carry i x y c)) := by
+  simp only [← getLsbD_eq_getElem]
+  rw [getLsbD_add_add_bool]
+  omega
+
+theorem getElem_add {i : Nat} (i_lt : i < w) (x y : BitVec w) :
+    (x + y)[i] = (x[i] ^^ (y[i] ^^ carry i x y false)) := by
+  simpa using getElem_add_add_bool i_lt x y false
+
 theorem adc_spec (x y : BitVec w) (c : Bool) :
     adc x y c = (carry w x y c, x + y + setWidth w (ofBool c)) := by
   simp only [adc]
@@ -368,6 +379,10 @@ theorem getLsbD_mul (x y : BitVec w) (i : Nat) :
   · simp
   · omega
 
+theorem getElem_mul {x y : BitVec w} {i : Nat} (h : i < w) :
+    (x * y)[i] = (mulRec x y w)[i] := by
+  simp [mulRec_eq_mul_signExtend_setWidth]
+
 /-! ## shiftLeft recurrence for bitblasting -/
 
 /--
diff --git a/src/Init/Data/BitVec/Lemmas.lean b/src/Init/Data/BitVec/Lemmas.lean
index c9f479ffda..a4e5818029 100644
--- a/src/Init/Data/BitVec/Lemmas.lean
+++ b/src/Init/Data/BitVec/Lemmas.lean
@@ -525,6 +525,7 @@ theorem getElem_setWidth' (x : BitVec w) (i : Nat) (h : w ≤ v) (hi : i < v) :
     (setWidth' h x)[i] = x.getLsbD i := by
   rw [getElem_eq_testBit_toNat, toNat_setWidth', getLsbD]
 
+@[simp]
 theorem getElem_setWidth (m : Nat) (x : BitVec n) (i : Nat) (h : i < m) :
     (setWidth m x)[i] = x.getLsbD i := by
   rw [setWidth]
@@ -2260,6 +2261,12 @@ theorem getLsbD_rotateLeft {x : BitVec w} {r i : Nat}  :
   · simp
   · rw [← rotateLeft_mod_eq_rotateLeft, getLsbD_rotateLeft_of_le (Nat.mod_lt _ (by omega))]
 
+@[simp]
+theorem getElem_rotateLeft {x : BitVec w} {r i : Nat} (h : i < w) :
+    (x.rotateLeft r)[i] =
+      if h' : i < r % w then x[(w - (r % w) + i)] else x[i - (r % w)] := by
+  simp [← BitVec.getLsbD_eq_getElem, h]
+
 /-! ## Rotate Right -/
 
 /--
@@ -2341,6 +2348,12 @@ theorem getLsbD_rotateRight {x : BitVec w} {r i : Nat} :
   · simp
   · rw [← rotateRight_mod_eq_rotateRight, getLsbD_rotateRight_of_le (Nat.mod_lt _ (by omega))]
 
+@[simp]
+theorem getElem_rotateRight {x : BitVec w} {r i : Nat} (h : i < w) :
+    (x.rotateRight r)[i] = if h' : i < w - (r % w) then x[(r % w) + i] else x[(i - (w - (r % w)))] := by
+  simp only [← BitVec.getLsbD_eq_getElem]
+  simp [getLsbD_rotateRight, h]
+
 /- ## twoPow -/
 
 @[simp, bv_toNat]
@@ -2369,6 +2382,12 @@ theorem getLsbD_twoPow (i j : Nat) : (twoPow w i).getLsbD j = ((i < w) && (i = j
           simp at hi
         simp_all
 
+@[simp]
+theorem getElem_twoPow {i j : Nat} (h : j < w) : (twoPow w i)[j] = decide (j = i) := by
+  rw [←getLsbD_eq_getElem, getLsbD_twoPow]
+  simp [eq_comm]
+  omega
+
 theorem and_twoPow (x : BitVec w) (i : Nat) :
     x &&& (twoPow w i) = if x.getLsbD i then twoPow w i else 0#w := by
   ext j
@@ -2400,6 +2419,10 @@ theorem twoPow_zero {w : Nat} : twoPow w 0 = 1#w := by
 theorem getLsbD_one {w i : Nat} : (1#w).getLsbD i = (decide (0 < w) && decide (0 = i)) := by
   rw [← twoPow_zero, getLsbD_twoPow]
 
+@[simp]
+theorem getElem_one {w i : Nat} (h : i < w) : (1#w)[i] = decide (i = 0) := by
+  rw [← twoPow_zero, getElem_twoPow]
+
 theorem shiftLeft_eq_mul_twoPow (x : BitVec w) (n : Nat) :
     x <<< n = x * (BitVec.twoPow w n) := by
   ext i
@@ -2504,6 +2527,12 @@ theorem getLsbD_replicate {n w : Nat} (x : BitVec w) :
       simp only [show ¬i < w * n by omega, decide_False, cond_false, hi, Bool.false_and]
       apply BitVec.getLsbD_ge (x := x) (i := i - w * n) (ge := by omega)
 
+@[simp]
+theorem getElem_replicate {n w : Nat} (x : BitVec w) (h : i < w * n) :
+    (x.replicate n)[i] = if h' : w = 0 then false else x[i % w]'(@Nat.mod_lt i w (by omega)) := by
+  simp only [← getLsbD_eq_getElem, getLsbD_replicate]
+  by_cases h' : w = 0 <;> simp [h'] <;> omega
+
 /-! ### intMin -/
 
 /-- The bitvector of width `w` that has the smallest value when interpreted as an integer. -/