From d0f4060336468be427c1478c9a4001b8612d530b Mon Sep 17 00:00:00 2001
From: Mac Malone <tydeu@hatpress.net>
Date: Sat, 23 Nov 2024 13:09:15 -0500
Subject: [PATCH] feat: more UInt bitwise theorems

---
 src/Init/Data/UInt/Bitwise.lean | 49 +++++++++++++++++++++++++++------
 1 file changed, 41 insertions(+), 8 deletions(-)

diff --git a/src/Init/Data/UInt/Bitwise.lean b/src/Init/Data/UInt/Bitwise.lean
index b9282fcf4f74..86a31feca122 100644
--- a/src/Init/Data/UInt/Bitwise.lean
+++ b/src/Init/Data/UInt/Bitwise.lean
@@ -4,22 +4,55 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Markus Himmel
 -/
 prelude
-import Init.Data.UInt.Basic
+import Init.Data.UInt.Lemmas
 import Init.Data.Fin.Bitwise
 import Init.Data.BitVec.Lemmas
 
+protected theorem Nat.lt_pow_self (h : a > 1) : n < a ^ n := by
+  induction n with
+  | zero => exact Nat.zero_lt_one
+  | succ n ih =>
+    cases a with
+    | zero => contradiction
+    | succ a =>
+      rw [Nat.pow_succ, Nat.mul_succ]
+      conv => rhs; rw [Nat.add_comm]
+      apply Nat.add_lt_add_of_lt_of_le ih
+      exact Nat.mul_pos (Nat.pow_pos (Nat.succ_pos a)) (Nat.lt_of_succ_lt_succ h)
+
+protected theorem Nat.lt_two_pow_self : n < 2 ^ n :=
+  Nat.lt_pow_self Nat.one_lt_two
+
+@[simp] protected theorem Nat.mod_two_pow_self (n : Nat) : n % 2 ^ n = n :=
+  Nat.mod_eq_of_lt Nat.lt_two_pow_self
+
 set_option hygiene false in
-macro "declare_bitwise_uint_theorems" typeName:ident : command =>
+macro "declare_bitwise_uint_theorems" typeName:ident bits:term:arg : command =>
+let toNat_and := Lean.mkIdentFrom typeName (typeName.getId.modifyBase (·.str "toNat_and"))
 `(
 namespace $typeName
 
-@[simp] protected theorem and_toNat (a b : $typeName) : (a &&& b).toNat = a.toNat &&& b.toNat := BitVec.toNat_and ..
+@[simp] protected theorem toNat_ofNat : (USize.ofNat n).toNat = n % 2 ^ System.Platform.numBits := BitVec.toNat_ofNat ..
+
+@[simp] protected theorem toBitVec_and (a b : $typeName) : (a &&& b).toBitVec = a.toBitVec &&& b.toBitVec := rfl
+@[simp] protected theorem toBitVec_or (a b : $typeName) : (a ||| b).toBitVec = a.toBitVec ||| b.toBitVec := rfl
+@[simp] protected theorem toBitVec_xor (a b : $typeName) : (a ^^^ b).toBitVec = a.toBitVec ^^^ b.toBitVec := rfl
+@[simp] protected theorem toBitVec_shiftLeft (a b : $typeName) : (a <<< b).toBitVec = a.toBitVec <<< (b.toBitVec % $bits) := rfl
+@[simp] protected theorem toBitVec_shiftRight (a b : $typeName) : (a >>> b).toBitVec = a.toBitVec >>> (b.toBitVec % $bits) := rfl
+
+@[simp] protected theorem toNat_and (a b : $typeName) : (a &&& b).toNat = a.toNat &&& b.toNat := by simp [toNat]
+@[simp] protected theorem toNat_or (a b : $typeName) : (a ||| b).toNat = a.toNat ||| b.toNat := by simp [toNat]
+@[simp] protected theorem toNat_xor (a b : $typeName) : (a ^^^ b).toNat = a.toNat ^^^ b.toNat := by simp [toNat]
+@[simp] protected theorem toNat_shiftLeft (a b : $typeName) : (a <<< b).toNat = a.toNat <<< (b.toNat % $bits) % 2 ^ $bits := by simp [toNat]
+@[simp] protected theorem toNat_shiftRight (a b : $typeName) : (a >>> b).toNat = a.toNat >>> (b.toNat % $bits) := by simp [toNat]
+
+@[deprecated $toNat_and (since := "2024-11-23")] protected theorem and_toNat (a b : $typeName) : (a &&& b).toNat = a.toNat &&& b.toNat := BitVec.toNat_and ..
 
 end $typeName
 )
 
-declare_bitwise_uint_theorems UInt8
-declare_bitwise_uint_theorems UInt16
-declare_bitwise_uint_theorems UInt32
-declare_bitwise_uint_theorems UInt64
-declare_bitwise_uint_theorems USize
+declare_bitwise_uint_theorems UInt8 8
+declare_bitwise_uint_theorems UInt16 16
+declare_bitwise_uint_theorems UInt32 32
+declare_bitwise_uint_theorems UInt64 64
+declare_bitwise_uint_theorems USize System.Platform.numBits