feat: add Nat.[shiftLeft_or_distrib, shiftLeft_xor_distrib, shiftLeft_and_distrib, testBit_mul_two_pow, bitwise_mul_two_pow, shiftLeft_bitwise_distrib]

src/Init/Data/Nat/Bitwise/Lemmas.lean

@@ -115,6 +115,33 @@ theorem testBit_add (x i n : Nat) : testBit x (i + n) = testBit (x / 2 ^ n) i :=
     rw [← Nat.add_assoc, testBit_add_one, ih (x / 2),
       Nat.pow_succ, Nat.div_div_eq_div_mul, Nat.mul_comm]
 
+theorem testBit_mul_two_pow_le {x i n : Nat} (h : n ≤ i) :
+    testBit (x * 2 ^ n) i = testBit x (i - n) := by
+  simp only [testBit, one_and_eq_mod_two, mod_two_bne_zero]
+  let j := i - n
+  have hj : (x * 2 ^ n) >>> i = (x * 2 ^ n) >>> (j + n) :=  by simp only [j]; rw [Nat.sub_add_cancel (by omega)]
+  have hj' : x >>> (i - n) = x >>> j :=  by simp only [j]
+  rw [hj, hj', ← shiftLeft_eq, Nat.add_comm, shiftRight_add, shiftLeft_shiftRight]
+
+theorem testBit_mul_two_pow_gt {x i n : Nat} (h : i < n) :
+    testBit (x * 2 ^ n) i = false := by
+  simp only [testBit, ← shiftLeft_eq, one_and_eq_mod_two, mod_two_bne_zero, beq_eq_false_iff_ne,
+    ne_eq]
+  let k := n - i
+  have hk : x <<< n >>> i = x <<< (k + i) >>> i := by simp only [k]; rw [Nat.sub_add_cancel (by omega)]
+  rw [hk, Nat.shiftLeft_add, Nat.shiftLeft_shiftRight, shiftLeft_eq]
+  have hx : 2 * (x * 2 ^ k / 2 ^ 1) = x * 2 ^ k := by
+    rw [Nat.mul_div_assoc, Nat.pow_div (by omega) (by omega), Nat.mul_comm, Nat.mul_assoc, ← Nat.pow_one (a := 2), ← Nat.pow_add, Nat.sub_add_cancel (by omega)]
+    simp only [Nat.pow_one, Nat.pow_dvd_pow (m := 1) (n := k) (a := 2) (by omega), k]
+  omega
+
+theorem testBit_mul_two_pow (x i n : Nat) :
+    testBit (x * 2 ^ n) i = if n ≤ i then testBit x (i - n) else false := by
+  by_cases hni : n ≤ i
+  · simp [hni, testBit_mul_two_pow_le hni]
+  · have : i < n := by omega
+    simp [hni, testBit_mul_two_pow_gt this]
+
 theorem testBit_div_two (x i : Nat) : testBit (x / 2) i = testBit x (i + 1) := by
   simp
 
@@ -452,6 +479,15 @@ theorem bitwise_lt_two_pow (left : x < 2^n) (right : y < 2^n) : (Nat.bitwise f x
       case neg =>
         apply Nat.add_lt_add <;> exact hyp1
 
+theorem bitwise_mul_two_pow (of_false_false : f false false = false := by rfl) :
+    (bitwise f x y) * 2 ^ n = bitwise f (x * 2 ^ n) (y * 2 ^ n) := by
+  apply Nat.eq_of_testBit_eq
+  simp only [testBit_mul_two_pow, testBit_bitwise of_false_false, Bool.if_false_right]
+  intro i
+  by_cases hn : n ≤ i
+  · simp [hn]
+  · simp [hn, of_false_false]
+
 theorem bitwise_div_two_pow (of_false_false : f false false = false := by rfl) :
   (bitwise f x y) / 2 ^ n = bitwise f (x / 2 ^ n) (y / 2 ^ n) := by
   apply Nat.eq_of_testBit_eq
@@ -711,6 +747,19 @@ theorem mul_add_lt_is_or {b : Nat} (b_lt : b < 2^i) (a : Nat) : 2^i * a + b = 2^
   rw [mod_two_eq_one_iff_testBit_zero, testBit_shiftLeft]
   simp
 
+theorem shiftLeft_bitwise_distrib {a b : Nat} (of_false_false : f false false = false := by rfl) :
+    (bitwise f a b) <<< i = bitwise f (a <<< i) (b <<< i) := by
+  simp [shiftLeft_eq, bitwise_mul_two_pow of_false_false]
+
+theorem shiftLeft_and_distrib {a b : Nat} : (a &&& b) <<< i = a <<< i &&& b <<< i :=
+  shiftLeft_bitwise_distrib
+
+theorem shiftLeft_or_distrib {a b : Nat} : (a ||| b) <<< i = a <<< i ||| b <<< i :=
+  shiftLeft_bitwise_distrib
+
+theorem shiftLeft_xor_distrib {a b : Nat} : (a ^^^ b) <<< i = a <<< i ^^^ b <<< i :=
+  shiftLeft_bitwise_distrib
+
 @[simp] theorem decide_shiftRight_mod_two_eq_one :
     decide (x >>> i % 2 = 1) = x.testBit i := by
   simp only [testBit, one_and_eq_mod_two, mod_two_bne_zero]