feat: bitVec shiftLeft recurrences for bitblasting

src/Init/Data/BitVec/Bitblast.lean

@@ -98,6 +98,37 @@ theorem carry_succ (i : Nat) (x y : BitVec w) (c : Bool) :
     exact mod_two_pow_add_mod_two_pow_add_bool_lt_two_pow_succ ..
   cases x.toNat.testBit i <;> cases y.toNat.testBit i <;> (simp; omega)
 
+/--
+If `x &&& y = 0`, then the carry bit `(x + y + 0)` is always `false` for any index `i`.
+Intuitively, this is because a carry is only produced when at least two of `x`, `y`, and the
+previous carry are true. However, since `x &&& y = 0`, at most one of `x, y` can be true,
+and thus we never have a previous carry, which means that the sum cannot produce a carry.
+-/
+theorem carry_of_and_eq_zero {x y : BitVec w} (h : x &&& y = 0#w) : carry i x y false = false := by
+  induction i with
+  | zero => simp
+  | succ i ih =>
+    replace h := congrArg (·.getLsb i) h
+    simp_all [carry_succ]
+
+/-- The final carry bit when computing `x + y + c` is `true` iff `x.toNat + y.toNat + c.toNat ≥ 2^w`. -/
+theorem carry_width {x y : BitVec w} :
+    carry w x y c = decide (x.toNat + y.toNat + c.toNat ≥ 2^w) := by
+  simp [carry]
+
+/--
+If `x &&& y = 0`, then addition does not overflow, and thus `(x + y).toNat = x.toNat + y.toNat`.
+-/
+theorem toNat_add_eq_toNat_add_toNat_of_and_eq_zero {x y : BitVec w} (h : x &&& y = 0#w) :
+    (x + y).toNat = x.toNat + y.toNat := by
+  rw [toNat_add]
+  apply Nat.mod_eq_of_lt
+  suffices ¬ decide (x.toNat + y.toNat + false.toNat ≥ 2^w) by
+    simp only [decide_eq_true_eq] at this
+    omega
+  rw [← carry_width]
+  simp [not_eq_true, carry_of_and_eq_zero h]
+
 /-- Carry function for bitwise addition. -/
 def adcb (x y c : Bool) : Bool × Bool := (atLeastTwo x y c, Bool.xor x (Bool.xor y c))
 
@@ -290,7 +321,7 @@ theorem zeroExtend_truncate_succ_eq_zeroExtend_truncate_add_twoPow (x : BitVec w
         simp [hik', hik'']
   · ext k
     simp
-    omega
+    by_cases hi : x.getLsb i <;> simp [hi] <;> omega
 
 /--
 Recurrence lemma: multiplying `l` with the first `s` bits of `r` is the
@@ -314,7 +345,7 @@ theorem mulRec_eq_mul_signExtend_truncate (l r : BitVec w) (s : Nat) :
     have heq :
       (if r.getLsb (s' + 1) = true then l <<< (s' + 1) else 0) =
         (l * (r &&& (BitVec.twoPow w (s' + 1)))) := by
-      simp only [ofNat_eq_ofNat, and_twoPow_eq]
+      simp only [ofNat_eq_ofNat, and_twoPow]
       by_cases hr : r.getLsb (s' + 1) <;> simp [hr]
     rw [heq, ← BitVec.mul_add, ← zeroExtend_truncate_succ_eq_zeroExtend_truncate_add_twoPow]
 
@@ -326,4 +357,78 @@ theorem getLsb_mul (x y : BitVec w) (i : Nat) :
   · simp
   · omega
 
+/-! ## shiftLeft recurrence for bitblasting -/
+
+/--
+`shiftLeftRec x y n` shifts `x` to the left by the first `n` bits of `y`.
+
+The theorem `shiftLeft_eq_shiftLeftRec` proves the equivalence of `(x <<< y)` and `shiftLeftRec`.
+
+This in conjecton with equations `shiftLeftRec_zero`, `shiftLeftRec_succ`
+allows us to unfold `shiftLeft` into a circuit for bitblasting.
+ -/
+def shiftLeftRec (x : BitVec w₁) (y : BitVec w₂) (n : Nat) : BitVec w₁ :=
+  let shiftAmt := (y &&& (twoPow w₂ n))
+  match n with
+  | 0 => x <<< shiftAmt
+  | n + 1 => (shiftLeftRec x y n) <<< shiftAmt
+
+@[simp]
+theorem shiftLeftRec_zero {x : BitVec w₁} {y : BitVec w₂} :
+    shiftLeftRec x y 0 = x <<< (y &&& twoPow w₂ 0)  := by
+  simp [shiftLeftRec]
+
+@[simp]
+theorem shiftLeftRec_succ {x : BitVec w₁} {y : BitVec w₂} :
+    shiftLeftRec x y (n + 1) = (shiftLeftRec x y n) <<< (y &&& twoPow w₂ (n + 1)) := by
+  simp [shiftLeftRec]
+
+/--
+If `y &&& z = 0`, `x <<< (y ||| z) = x <<< y <<< z`.
+This follows as `y &&& z = 0` implies `y ||| z = y + z`,
+and thus `x <<< (y ||| z) = x <<< (y + z) = x <<< y <<< z`.
+-/
+theorem shiftLeft_or_eq_shiftLeft_shiftLeft_of_and_eq_zero {x : BitVec w₁} {y z : BitVec w₂}
+    (h : y &&& z = 0#w₂) :
+    x <<< (y ||| z) = x <<< y <<< z := by
+  rw [← add_eq_or_of_and_eq_zero _ _ h,
+    shiftLeft_eq', toNat_add_eq_toNat_add_toNat_of_and_eq_zero h]
+  simp [shiftLeft_add]
+
+/--
+`shiftLeftRec x y n` shifts `x` to the left by the first `n` bits of `y`.
+-/
+theorem shiftLeftRec_eq {x : BitVec w₁} {y : BitVec w₂} {n : Nat} :
+    shiftLeftRec x y n = x <<< (y.truncate (n + 1)).zeroExtend w₂ := by
+  induction n generalizing x y
+  case zero =>
+    ext i
+    simp only [shiftLeftRec_zero, twoPow_zero, Nat.reduceAdd, truncate_one_eq_ofBool_getLsb]
+    suffices (y &&& 1#w₂) = zeroExtend w₂ (ofBool (y.getLsb 0)) by simp [this]
+    ext i
+    by_cases h : (↑i : Nat) = 0
+    · simp [h, Bool.and_comm]
+    · simp [h]; omega
+  case succ n ih =>
+    simp only [shiftLeftRec_succ, and_twoPow]
+    rw [ih]
+    by_cases h : y.getLsb (n + 1)
+    · simp only [h, ↓reduceIte]
+      rw [zeroExtend_truncate_succ_eq_zeroExtend_truncate_or_twoPow_of_getLsb_true h,
+        shiftLeft_or_eq_shiftLeft_shiftLeft_of_and_eq_zero]
+      simp
+    · simp only [h, false_eq_true, ↓reduceIte, shiftLeft_zero']
+      rw [zeroExtend_truncate_succ_eq_zeroExtend_truncate_of_getLsb_false (i := n + 1)]
+      simp [h]
+
+/--
+Show that `x <<< y` can be written in terms of `shiftLeftRec`.
+This can be unfolded in terms of `shiftLeftRec_zero`, `shiftLeftRec_succ` for bitblasting.
+-/
+theorem shiftLeft_eq_shiftLeftRec (x : BitVec w₁) (y : BitVec w₂) :
+    x <<< y = shiftLeftRec x y (w₂ - 1) := by
+  rcases w₂ with rfl | w₂
+  · simp [of_length_zero]
+  · simp [shiftLeftRec_eq]
+
 end BitVec

src/Init/Data/BitVec/Lemmas.lean

@@ -436,6 +436,12 @@ theorem zeroExtend_ofNat_one_eq_ofNat_one_of_lt {v w : Nat} (hv : 0 < v) :
   have hv := Nat.testBit_one_eq_true_iff_self_eq_zero.mp hi₁
   omega
 
+/-- Truncating to width 1 produces a bitvector equal to the least significant bit. -/
+theorem truncate_one_eq_ofBool_getLsb {x : BitVec w} :
+    x.truncate 1 = ofBool (x.getLsb 0) := by
+  ext i
+  simp [show i = 0 by omega]
+
 /-! ## extractLsb -/
 
 @[simp]
@@ -531,6 +537,11 @@ theorem and_assoc (x y z : BitVec w) :
   ext i
   simp [Bool.and_assoc]
 
+theorem and_comm (x y : BitVec w) :
+    x &&& y = y &&& x := by
+  ext i
+  simp [Bool.and_comm]
+
 /-! ### xor -/
 
 @[simp] theorem toNat_xor (x y : BitVec v) :
@@ -626,6 +637,10 @@ theorem shiftLeft_zero_eq (x : BitVec w) : x <<< 0 = x := by
   apply eq_of_toNat_eq
   simp
 
+@[simp]
+theorem zero_shiftLeft (n : Nat) : 0#w <<< n = 0#w := by
+  simp [bv_toNat]
+
 @[simp] theorem getLsb_shiftLeft (x : BitVec m) (n) :
     getLsb (x <<< n) i = (decide (i < m) && !decide (i < n) && getLsb x (i - n)) := by
   rw [← testBit_toNat, getLsb]
@@ -691,6 +706,22 @@ theorem shiftLeft_shiftLeft {w : Nat} (x : BitVec w) (n m : Nat) :
     (x <<< n) <<< m = x <<< (n + m) := by
   rw [shiftLeft_add]
 
+/-! ### shiftLeft reductions from BitVec to Nat -/
+
+@[simp]
+theorem shiftLeft_eq' {x : BitVec w₁} {y : BitVec w₂} : x <<< y = x <<< y.toNat := by rfl
+
+@[simp]
+theorem shiftLeft_zero' {x : BitVec w₁} : x <<< 0#w₂ = x := by simp
+
+theorem shiftLeft_shiftLeft' {x : BitVec w₁} {y : BitVec w₂} {z : BitVec w₃} :
+    x <<< y <<< z = x <<< (y.toNat + z.toNat) := by
+  simp [shiftLeft_add]
+
+theorem getLsb_shiftLeft' {x : BitVec w₁} {y : BitVec w₂} {i : Nat} :
+    (x <<< y).getLsb i = (decide (i < w₁) && !decide (i < y.toNat) && x.getLsb (i - y.toNat)) := by
+  simp [shiftLeft_eq', getLsb_shiftLeft]
+
 /-! ### ushiftRight -/
 
 @[simp, bv_toNat] theorem toNat_ushiftRight (x : BitVec n) (i : Nat) :
@@ -1452,12 +1483,18 @@ theorem getLsb_twoPow (i j : Nat) : (twoPow w i).getLsb j = ((i < w) && (i = j))
           simp at hi
         simp_all
 
-theorem and_twoPow_eq (x : BitVec w) (i : Nat) :
+@[simp]
+theorem and_twoPow (x : BitVec w) (i : Nat) :
     x &&& (twoPow w i) = if x.getLsb i then twoPow w i else 0#w := by
   ext j
   simp only [getLsb_and, getLsb_twoPow]
   by_cases hj : i = j <;> by_cases hx : x.getLsb i <;> simp_all
 
+@[simp]
+theorem twoPow_and (x : BitVec w) (i : Nat) :
+    (twoPow w i) &&& x = if x.getLsb i then twoPow w i else 0#w := by
+  rw [BitVec.and_comm, and_twoPow]
+
 @[simp]
 theorem mul_twoPow_eq_shiftLeft (x : BitVec w) (i : Nat) :
     x * (twoPow w i) = x <<< i := by
@@ -1471,6 +1508,14 @@ theorem mul_twoPow_eq_shiftLeft (x : BitVec w) (i : Nat) :
       apply Nat.pow_dvd_pow 2 (by omega)
     simp [Nat.mul_mod, hpow]
 
+theorem twoPow_zero {w : Nat} : twoPow w 0 = 1#w := by
+  apply eq_of_toNat_eq
+  simp
+
+@[simp]
+theorem getLsb_one {w i : Nat} : (1#w).getLsb i = (decide (0 < w) && decide (0 = i)) := by
+  rw [← twoPow_zero, getLsb_twoPow]
+
 /- ### zeroExtend, truncate, and bitwise operations -/
 
 /--