feat: recurrence for BitVec.mul as repeated shifts for bitblasting

src/Init/Data/BitVec/Basic.lean

@@ -614,6 +614,13 @@ theorem ofBool_append (msb : Bool) (lsbs : BitVec w) :
     ofBool msb ++ lsbs = (cons msb lsbs).cast (Nat.add_comm ..) :=
   rfl
 
+/--
+`twoPow i` is the bitvector `2^i` if `i < w`, and `0` otherwise.
+That is, 2 to the power `i`.
+For the bitwise point of view, it has the `i`th bit as `1` and all other bits as `0`.
+-/
+def twoPow {w : Nat} (i : Nat) : BitVec w := (1#w) <<< i
+
 end bitwise
 
 section normalization_eqs

src/Init/Data/BitVec/Bitblast.lean

@@ -159,6 +159,20 @@ theorem add_eq_adc (w : Nat) (x y : BitVec w) : x + y = (adc x y false).snd := b
 theorem allOnes_sub_eq_not (x : BitVec w) : allOnes w - x = ~~~x := by
   rw [← add_not_self x, BitVec.add_comm, add_sub_cancel]
 
+/-- Adding two bitvectors equals or-ing them if they are 1 in mutually exclusive locations -/
+theorem add_eq_or_of_and_eq_zero {w : Nat} (x y : BitVec w)
+    (h : x &&& y = 0#w) : x + y = x ||| y := by
+  rw [add_eq_adc, adc, iunfoldr_replace (fun _ => false) (x ||| y)]
+  · rfl
+  · simp [adcb, atLeastTwo, h]
+    intros i
+    replace h : (x &&& y).getLsb i = (0#w).getLsb i := by rw [h]
+    simp only [getLsb_and, getLsb_zero, and_eq_false_imp] at h
+    constructor
+    · intros hx
+      simp_all [hx]
+    · by_cases hx : x.getLsb i <;> simp_all [hx]
+
 /-! ### Negation -/
 
 theorem bit_not_testBit (x : BitVec w) (i : Fin w) :
@@ -235,4 +249,83 @@ theorem sle_eq_carry (x y : BitVec w) :
     x.sle y = !((x.msb == y.msb).xor (carry w y (~~~x) true)) := by
   rw [sle_eq_not_slt, slt_eq_not_carry, beq_comm]
 
+/-! ### mul recurrence for bitblasting -/
+
+def mulRec (l r : BitVec w) (s : Nat) : BitVec w :=
+  let cur := if r.getLsb s then (l <<< s) else 0
+  match s with
+  | 0 => cur
+  | s + 1 => mulRec l r s + cur
+
+theorem mulRec_zero_eq (l r : BitVec w) :
+    mulRec l r 0 = if r.getLsb 0 then l else 0 := by
+  simp [mulRec]
+
+theorem mulRec_succ_eq (l r : BitVec w) (s : Nat) :
+    mulRec l r (s + 1) = mulRec l r s + if r.getLsb (s + 1) then (l <<< (s + 1)) else 0 := by
+  simp [mulRec]
+
+/-- Recurrence lemma: truncating to `i+1` bits and then zero extending to `w`
+equals truncating upto `i` bits `[0..i-1]`, and then adding the `i`th bit of `x`. -/
+theorem zeroExtend_truncate_succ_eq_zeroExtend_truncate_add_twoPow (x : BitVec w) (i : Nat) :
+    zeroExtend w (x.truncate (i + 1)) =
+      zeroExtend w (x.truncate i) + (x &&& twoPow i) := by
+  rw [add_eq_or_of_and_eq_zero]
+  · ext k
+    simp only [getLsb_zeroExtend, Fin.is_lt, decide_True, Bool.true_and, getLsb_or, getLsb_and]
+    by_cases hik:i = k
+    · subst hik
+      simp
+    · simp [hik]
+      /- Really, 'omega' should be able to do this-/
+      by_cases hik' : k < (i + 1)
+      · have hik'' : k < i := by omega
+        simp [hik', hik'']
+      · have hik'' : ¬ (k < i) := by omega
+        simp [hik', hik'']
+  · ext k
+    simp
+    intros h₁ _ _ _
+    omega
+
+theorem mulRec_eq_mul_signExtend_truncate (l r : BitVec w) (s : Nat) :
+    (mulRec l r s) = l * ((r.truncate (s + 1)).zeroExtend w) := by
+  induction s
+  case zero =>
+    simp [mulRec_zero_eq]
+    by_cases r.getLsb 0
+    case pos hr =>
+      simp only [hr, ↓reduceIte, truncate, zeroExtend_one_eq_ofBool_getLsb_zero,
+        hr, ofBool_true, ofNat_eq_ofNat]
+      rw [zeroExtend_ofNat_one_eq_ofNat_one_of_lt (by omega)]; simp
+    case neg hr =>
+      simp [hr, zeroExtend_one_eq_ofBool_getLsb_zero]
+  case succ s' hs =>
+    rw [mulRec_succ_eq, hs]
+    have heq :
+      (if r.getLsb (s' + 1) = true then l <<< (s' + 1) else 0) =
+        (l * (r &&& (BitVec.twoPow (s' + 1)))) := by
+      simp only [ofNat_eq_ofNat, and_twoPow_eq_getLsb]
+      by_cases hr : r.getLsb (s' + 1) <;> simp [hr]
+    rw [heq, ← BitVec.mul_add, ← zeroExtend_truncate_succ_eq_zeroExtend_truncate_add_twoPow]
+
+/-- Zero extending by number of bits larger than the bitwidth has no effect. -/
+theorem zeroExtend_of_ge {x : BitVec w} {i j : Nat} (hi : i ≥ w) :
+    (x.zeroExtend i).zeroExtend j = x.zeroExtend j := by
+  ext k
+  simp
+  intros hx;
+  have hi' : k < w := BitVec.lt_of_getLsb _ _ hx
+  omega
+
+/-- Zero extending by the bitwidth has no effect. -/
+theorem zeroExtend_eq_self {x : BitVec w} : x.zeroExtend w = x := by
+  ext i
+  simp [getLsb_zeroExtend]
+
+theorem getLsb_mul (x y : BitVec w) (i : Nat) :
+    (x * y).getLsb i = (mulRec x y w).getLsb i := by
+  simp [mulRec_eq_mul_signExtend_truncate]
+  rw [truncate, zeroExtend_of_ge (by omega), zeroExtend_eq_self]
+
 end BitVec

src/Init/Data/BitVec/Lemmas.lean

@@ -163,6 +163,16 @@ theorem toNat_zero (n : Nat) : (0#n).toNat = 0 := by trivial
 private theorem lt_two_pow_of_le {x m n : Nat} (lt : x < 2 ^ m) (le : m ≤ n) : x < 2 ^ n :=
   Nat.lt_of_lt_of_le lt (Nat.pow_le_pow_of_le_right (by trivial : 0 < 2) le)
 
+@[simp]
+theorem getLsb_ofBool (b : Bool) (i : Nat) : (BitVec.ofBool b).getLsb i = ((i = 0) && b) := by
+  rcases b with rfl | rfl
+  · simp [ofBool]
+  · simp [ofBool, getLsb_ofNat]
+    by_cases hi : (i = 0)
+    · simp [hi]
+    · simp [hi]
+      omega
+
 /-! ### msb -/
 
 @[simp] theorem msb_zero : (0#w).msb = false := by simp [BitVec.msb, getMsb]
@@ -408,6 +418,29 @@ theorem msb_zeroExtend (x : BitVec w) : (x.zeroExtend v).msb = (decide (0 < v) &
 theorem msb_zeroExtend' (x : BitVec w) (h : w ≤ v) : (x.zeroExtend' h).msb = (decide (0 < v) && x.getLsb (v - 1)) := by
   rw [zeroExtend'_eq, msb_zeroExtend]
 
+/-- zero extending a bitvector to width 1 equals the boolean of the lsb. -/
+theorem zeroExtend_one_eq_ofBool_getLsb_zero (x : BitVec w) :
+    x.zeroExtend 1 = BitVec.ofBool (x.getLsb 0) := by
+  ext i
+  simp [getLsb_zeroExtend, Fin.fin_one_eq_zero i]
+
+/-- `testBit 1 i` is true iff the index `i` equals 0. -/
+private theorem Nat.testBit_one_eq_true_iff_self_eq_zero {i : Nat} :
+    Nat.testBit 1 i = true ↔ i = 0 := by
+  cases i <;> simp
+
+/-- Zero extending `1#v` to `1#w` equals `1#w` when `v > 0`. -/
+theorem zeroExtend_ofNat_one_eq_ofNat_one_of_lt {v w : Nat} (hv : 0 < v):
+    (BitVec.ofNat v 1).zeroExtend w = BitVec.ofNat w 1 := by
+  ext i
+  obtain ⟨i, hilt⟩ := i
+  simp only [getLsb_zeroExtend, hilt, decide_True, getLsb_ofNat, Bool.true_and,
+    Bool.and_iff_right_iff_imp, decide_eq_true_eq]
+  intros hi1
+  have hv := Nat.testBit_one_eq_true_iff_self_eq_zero.mp hi1
+  omega
+
+
 /-! ## extractLsb -/
 
 @[simp]
@@ -593,6 +626,11 @@ theorem not_def {x : BitVec v} : ~~~x = allOnes v ^^^ x := rfl
 @[simp] theorem toFin_shiftLeft {n : Nat} (x : BitVec w) :
     BitVec.toFin (x <<< n) = Fin.ofNat' (x.toNat <<< n) (Nat.two_pow_pos w) := rfl
 
+@[simp]
+theorem shiftLeft_zero_eq (x : BitVec w) : x <<< 0 = x := by
+  apply eq_of_toNat_eq
+  simp
+
 @[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]
@@ -726,6 +764,15 @@ theorem getLsb_sshiftRight (x : BitVec w) (s i : Nat) :
         Nat.not_lt, decide_eq_true_eq]
       omega
 
+/-! ### udiv -/
+
+theorem udiv_eq {x y : BitVec n} :
+    x.udiv y = BitVec.ofNat n (x.toNat / y.toNat) := by
+  apply BitVec.eq_of_toNat_eq
+  simp only [udiv, toNat_ofNatLt, toNat_ofNat]
+  rw [Nat.mod_eq_of_lt]
+  exact Nat.lt_of_le_of_lt (Nat.div_le_self ..) (by omega)
+
 /-! ### append -/
 
 theorem append_def (x : BitVec v) (y : BitVec w) :
@@ -1085,6 +1132,23 @@ theorem ofInt_mul {n} (x y : Int) : BitVec.ofInt n (x * y) =
   apply eq_of_toInt_eq
   simp
 
+@[simp]
+theorem BitVec.mul_one {x : BitVec w} : x * 1#w = x := by
+  apply eq_of_toNat_eq
+  simp [toNat_mul, Nat.mod_eq_of_lt x.isLt]
+
+@[simp]
+theorem BitVec.mul_zero {x : BitVec w} : x * (0#w) = (0#w) := by
+  apply eq_of_toNat_eq
+  simp [toNat_mul]
+
+theorem BitVec.mul_add {x y z : BitVec w} :
+    x * (y + z) = x * y + x * z := by
+  apply eq_of_toNat_eq
+  simp
+  rw [Nat.mul_mod, Nat.mod_mod (y.toNat + z.toNat),
+    ← Nat.mul_mod, Nat.mul_add]
+
 /-! ### le and lt -/
 
 @[bv_toNat] theorem le_def (x y : BitVec n) :
@@ -1311,4 +1375,63 @@ theorem getLsb_rotateRight {x : BitVec w} {r i : Nat} :
   · simp
   · rw [← rotateRight_mod_eq_rotateRight, getLsb_rotateRight_of_le (Nat.mod_lt _ (by omega))]
 
+/- ## twoPow -/
+
+@[simp]
+theorem toNat_twoPow (w : Nat) (i : Nat) : (twoPow i : BitVec w).toNat = 2^i % 2^w := by
+  rcases w with rfl | w
+  · simp [Nat.mod_one]
+  · simp [twoPow, toNat_shiftLeft]
+    have h1 : 1 < 2 ^ (w + 1) := Nat.one_lt_two_pow (by omega)
+    rw [Nat.mod_eq_of_lt h1]
+    rw [Nat.shiftLeft_eq, Nat.one_mul]
+
+@[simp]
+theorem getLsb_twoPow (i j : Nat) : (twoPow i : BitVec w).getLsb j = ((i < w) && (i = j)) := by
+  rcases w with rfl | w
+  · simp only [twoPow, BitVec.reduceOfNat, Nat.zero_le, getLsb_ge, Bool.false_eq,
+    Bool.and_eq_false_imp, decide_eq_true_eq, decide_eq_false_iff_not]
+    omega
+  · simp only [twoPow, getLsb_shiftLeft, getLsb_ofNat]
+    by_cases hj : j < i
+    · simp only [hj, decide_True, Bool.not_true, Bool.and_false, Bool.false_and, Bool.false_eq,
+      Bool.and_eq_false_imp, decide_eq_true_eq, decide_eq_false_iff_not]
+      omega
+    · by_cases hi : Nat.testBit 1 (j - i)
+      · obtain hi' := Nat.testBit_one_eq_true_iff_self_eq_zero.mp hi
+        have hij : j = i := by omega
+        simp_all
+      · have hij : i ≠ j := by
+          intro h; subst h
+          simp at hi
+        simp_all
+
+theorem and_twoPow_eq_getLsb (x : BitVec w) (i : Nat) :
+    x &&& (twoPow i : BitVec w) = if x.getLsb i then twoPow 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 mul_twoPow_eq_shiftLeft (x : BitVec w) (i : Nat) :
+    x * (twoPow i : BitVec w) = x <<< i := by
+  apply eq_of_toNat_eq
+  simp only [toNat_mul, toNat_twoPow, toNat_shiftLeft, Nat.shiftLeft_eq]
+  by_cases hi : i < w
+  · have hpow : 2^i < 2^w := Nat.pow_lt_pow_of_lt (by omega) (by omega)
+    rw [Nat.mod_eq_of_lt hpow]
+  · have hpow : 2 ^ i % 2 ^ w = 0 := by
+      rw [Nat.mod_eq_zero_of_dvd]
+      apply Nat.pow_dvd_pow 2 (by omega)
+    simp [Nat.mul_mod, hpow]
+
+theorem BitVec.toNat_twoPow (w : Nat) (i : Nat) : (twoPow i : BitVec w).toNat = 2^i % 2^w := by
+  rcases w with rfl | w
+  · simp [Nat.mod_one]
+  · simp [twoPow, toNat_shiftLeft]
+    have hone : 1 < 2 ^ (w + 1) := by
+      rw [show 1 = 2^0 by simp[Nat.pow_zero]]
+      exact Nat.pow_lt_pow_of_lt (by omega) (by omega)
+    simp [Nat.mod_eq_of_lt hone, Nat.shiftLeft_eq]
+
 end BitVec