From ffac974dba799956a97d63ffcb13a774f700149c Mon Sep 17 00:00:00 2001
From: Mac Malone <tydeu@hatpress.net>
Date: Sun, 1 Dec 2024 17:38:49 -0500
Subject: [PATCH] feat: more UInt bitwise theorems (#6188)

This PR completes the `toNat` theorems for the bitwise operations
(`and`, `or`, `xor`, `shiftLeft`, `shiftRight`) of the UInt types and
adds `toBitVec` theorems as well. It also renames `and_toNat` to
`toNat_and` to fit with the current naming convention.
---
 src/Init/Data/UInt/Bitwise.lean           | 32 ++++++++++++++++-------
 src/Std/Data/DHashMap/Internal/Index.lean | 16 +++---------
 2 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/src/Init/Data/UInt/Bitwise.lean b/src/Init/Data/UInt/Bitwise.lean
index b9282fcf4f74..25912f31d814 100644
--- a/src/Init/Data/UInt/Bitwise.lean
+++ b/src/Init/Data/UInt/Bitwise.lean
@@ -1,25 +1,39 @@
 /-
 Copyright (c) 2024 Lean FRO, LLC. All Rights Reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
-Authors: Markus Himmel
+Authors: Markus Himmel, Mac Malone
 -/
 prelude
-import Init.Data.UInt.Basic
+import Init.Data.UInt.Lemmas
 import Init.Data.Fin.Bitwise
 import Init.Data.BitVec.Lemmas
 
 set_option hygiene false in
-macro "declare_bitwise_uint_theorems" typeName:ident : command =>
+macro "declare_bitwise_uint_theorems" typeName:ident bits:term:arg : command =>
 `(
 namespace $typeName
 
-@[simp] protected theorem and_toNat (a b : $typeName) : (a &&& b).toNat = a.toNat &&& b.toNat := BitVec.toNat_and ..
+@[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]
+
+open $typeName (toNat_and) in
+@[deprecated toNat_and (since := "2024-11-28")]
+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
diff --git a/src/Std/Data/DHashMap/Internal/Index.lean b/src/Std/Data/DHashMap/Internal/Index.lean
index 60e14072e767..0bfca7684458 100644
--- a/src/Std/Data/DHashMap/Internal/Index.lean
+++ b/src/Std/Data/DHashMap/Internal/Index.lean
@@ -46,19 +46,11 @@ cf. https://github.com/leanprover/lean4/issues/4157
 @[irreducible, inline] def mkIdx (sz : Nat) (h : 0 < sz) (hash : UInt64) :
     { u : USize // u.toNat < sz } :=
   ⟨(scrambleHash hash).toUSize &&& (sz.toUSize - 1), by
-    -- Making this proof significantly less painful will be a good test for our USize API
+    -- This proof is a good test for our USize API
     by_cases h' : sz < USize.size
-    · rw [USize.and_toNat, ← USize.ofNat_one, USize.toNat_sub_of_le, USize.toNat_ofNat_of_lt]
-      · refine Nat.lt_of_le_of_lt Nat.and_le_right (Nat.sub_lt h ?_)
-        rw [USize.toNat_ofNat_of_lt]
-        · exact Nat.one_pos
-        · exact Nat.lt_of_le_of_lt h h'
-      · exact h'
-      · rw [USize.le_def, BitVec.le_def]
-        change _ ≤ (_ % _)
-        rw [Nat.mod_eq_of_lt h', USize.ofNat, BitVec.toNat_ofNat, Nat.mod_eq_of_lt]
-        · exact h
-        · exact Nat.lt_of_le_of_lt h h'
+    · rw [USize.toNat_and, USize.toNat_sub_of_le, USize.toNat_ofNat_of_lt h']
+      · exact Nat.lt_of_le_of_lt Nat.and_le_right (Nat.sub_lt h (by simp))
+      · simp [USize.le_iff_toNat_le, Nat.mod_eq_of_lt h', Nat.succ_le_of_lt h]
     · exact Nat.lt_of_lt_of_le (USize.toNat_lt_size _) (Nat.le_of_not_lt h')⟩
 
 end Std.DHashMap.Internal