From c7c1e091c9f07ae6f8e8ff7246eb7650e2740dcb Mon Sep 17 00:00:00 2001
From: Vlad Tsyrklevich <vlad@tsyrklevich.net>
Date: Tue, 28 Jan 2025 09:48:11 +0100
Subject: [PATCH] feat: add BitVec comparison lemmas to bv_normalize (#6799)

This PR adds a number of simple comparison lemmas to the top/bottom
element for BitVec. Then they are applied to teach bv_normalize that
`(a<1) = (a==0)` and to remove an intermediate proof that is no longer
necessary along the way.
---
 src/Init/Data/BitVec/Lemmas.lean              | 34 +++++++++++++++++++
 src/Std/Tactic/BVDecide/Normalize/BitVec.lean | 17 ++++------
 tests/lean/run/bv_decide_rewriter.lean        |  4 +++
 3 files changed, 44 insertions(+), 11 deletions(-)

diff --git a/src/Init/Data/BitVec/Lemmas.lean b/src/Init/Data/BitVec/Lemmas.lean
index 5e511e3d4119..f9ff5a6a7a78 100644
--- a/src/Init/Data/BitVec/Lemmas.lean
+++ b/src/Init/Data/BitVec/Lemmas.lean
@@ -2716,6 +2716,40 @@ theorem not_lt_iff_le {x y : BitVec w} : (¬ x < y) ↔ y ≤ x := by
   constructor <;>
     (intro h; simp only [lt_def, Nat.not_lt, le_def] at h ⊢; omega)
 
+@[simp]
+theorem not_lt_zero {x : BitVec w} : ¬x < 0#w := of_decide_eq_false rfl
+
+@[simp]
+theorem le_zero_iff {x : BitVec w} : x ≤ 0#w ↔ x = 0#w := by
+  constructor
+  · intro h
+    have : x ≥ 0 := not_lt_iff_le.mp not_lt_zero
+    exact Eq.symm (BitVec.le_antisymm this h)
+  · simp_all
+
+@[simp]
+theorem lt_one_iff {x : BitVec w} (h : 0 < w) : x < 1#w ↔ x = 0#w := by
+  constructor
+  · intro h₂
+    rw [lt_def, toNat_ofNat, ← Int.ofNat_lt, Int.ofNat_emod, Int.ofNat_one, Int.natCast_pow,
+      Int.ofNat_two, @Int.emod_eq_of_lt 1 (2^w) (by omega) (by omega)] at h₂
+    simp [toNat_eq, show x.toNat = 0 by omega]
+  · simp_all
+
+@[simp]
+theorem not_allOnes_lt {x : BitVec w} : ¬allOnes w < x := by
+  have : 2^w ≠ 0 := Ne.symm (NeZero.ne' (2^w))
+  rw [BitVec.not_lt, le_def, Nat.le_iff_lt_add_one, toNat_allOnes, Nat.sub_one_add_one this]
+  exact isLt x
+
+@[simp]
+theorem allOnes_le_iff {x : BitVec w} : allOnes w ≤ x ↔ x = allOnes w := by
+  constructor
+  · intro h
+    have : x ≤ allOnes w := not_lt_iff_le.mp not_allOnes_lt
+    exact Eq.symm (BitVec.le_antisymm h this)
+  · simp_all
+
 /-! ### udiv -/
 
 theorem udiv_def {x y : BitVec n} : x / y = BitVec.ofNat n (x.toNat / y.toNat) := by
diff --git a/src/Std/Tactic/BVDecide/Normalize/BitVec.lean b/src/Std/Tactic/BVDecide/Normalize/BitVec.lean
index caaa37084132..6b77a93d1e7b 100644
--- a/src/Std/Tactic/BVDecide/Normalize/BitVec.lean
+++ b/src/Std/Tactic/BVDecide/Normalize/BitVec.lean
@@ -299,20 +299,15 @@ theorem BitVec.lt_irrefl (a : BitVec n) : (BitVec.ult a a) = false := by
 @[bv_normalize]
 theorem BitVec.not_lt_zero (a : BitVec n) : (BitVec.ult a 0#n) = false := by rfl
 
-theorem BitVec.max_ult (a : BitVec w) : ¬ ((-1#w) < a) := by
-  rcases w with rfl | w
-  · simp [bv_toNat, BitVec.toNat_of_zero_length]
-  · simp only [BitVec.lt_def, BitVec.toNat_neg, BitVec.toNat_ofNat, Nat.not_lt]
-    rw [Nat.mod_eq_of_lt (a := 1) (by simp)];
-    rw [Nat.mod_eq_of_lt]
-    · omega
-    · apply Nat.sub_one_lt_of_le (Nat.pow_pos (by omega)) (Nat.le_refl ..)
+@[bv_normalize]
+theorem BitVec.lt_one_iff (a : BitVec n) (h : 0 < n) : (BitVec.ult a 1#n) = (a == 0#n) := by
+  rw [Bool.eq_iff_iff, beq_iff_eq, ← BitVec.lt_ult]
+  exact _root_.BitVec.lt_one_iff h
 
 -- used in simproc because of -1#w normalisation
 theorem BitVec.max_ult' (a : BitVec w) : (BitVec.ult (-1#w) a) = false := by
-  have := BitVec.max_ult a
-  rw [BitVec.lt_ult] at this
-  simp [this]
+  rw [BitVec.negOne_eq_allOnes, ← Bool.not_eq_true, ← @lt_ult]
+  exact BitVec.not_allOnes_lt
 
 attribute [bv_normalize] BitVec.replicate_zero_eq
 attribute [bv_normalize] BitVec.add_eq_xor
diff --git a/tests/lean/run/bv_decide_rewriter.lean b/tests/lean/run/bv_decide_rewriter.lean
index 08717a7954dd..20e7cac9625b 100644
--- a/tests/lean/run/bv_decide_rewriter.lean
+++ b/tests/lean/run/bv_decide_rewriter.lean
@@ -97,6 +97,10 @@ example (x : BitVec 16) : ¬x < 0 := by bv_normalize
 example (x : BitVec 16) : x ≥ 0 := by bv_normalize
 example (x : BitVec 16) : !(x.ult 0) := by bv_normalize
 
+-- lt_one_iff
+example (x : BitVec 16) : (x < 1) ↔ (x = 0) := by bv_normalize
+example (x : BitVec 16) : (x.ult 1) = (x == 0) := by bv_normalize
+
 section
 
 example (x y : BitVec 256) : x * y = y * x := by