From 7bd12c71c8c901becd7d717859f66001c971d354 Mon Sep 17 00:00:00 2001
From: Vlad Tsyrklevich <vlad@tsyrklevich.net>
Date: Fri, 31 Jan 2025 14:27:43 +0100
Subject: [PATCH] feat: add or/and/xor lemmas for BitVec/bv_normalize (#6872)

This PR adds lemmas for xor injectivity and when and/or/xor equal
allOnes or zero. Then I plumb support for the new lemmas through to
bv_normalize.
---
 src/Init/Data/BitVec/Lemmas.lean             | 52 ++++++++++++++++++++
 src/Std/Tactic/BVDecide/Normalize/Equal.lean | 40 +++++++++++++++
 tests/lean/run/bv_axiom_check.lean           |  3 +-
 tests/lean/run/bv_decide_rewriter.lean       | 14 ++++++
 4 files changed, 107 insertions(+), 2 deletions(-)

diff --git a/src/Init/Data/BitVec/Lemmas.lean b/src/Init/Data/BitVec/Lemmas.lean
index 8d7920111f79..cab1faa00004 100644
--- a/src/Init/Data/BitVec/Lemmas.lean
+++ b/src/Init/Data/BitVec/Lemmas.lean
@@ -937,6 +937,19 @@ instance : Std.LawfulCommIdentity (α := BitVec n) (· ||| · ) (0#n) where
   ext i h
   simp [h]
 
+@[simp]
+theorem or_eq_zero_iff {x y : BitVec w} : (x ||| y) = 0#w ↔ x = 0#w ∧ y = 0#w := by
+  constructor
+  · intro h
+    constructor
+    all_goals
+    · ext i ih
+      have := BitVec.eq_of_getLsbD_eq_iff.mp h i ih
+      simp only [getLsbD_or, getLsbD_zero, Bool.or_eq_false_iff] at this
+      simp [this]
+  · intro h
+    simp [h]
+
 theorem extractLsb'_or {x y : BitVec w} {start len : Nat} :
    (x ||| y).extractLsb' start len = (x.extractLsb' start len) ||| (y.extractLsb' start len) := by
   ext i hi
@@ -1020,6 +1033,20 @@ instance : Std.LawfulCommIdentity (α := BitVec n) (· &&& · ) (allOnes n) wher
   ext i h
   simp [h]
 
+@[simp]
+theorem and_eq_allOnes_iff {x y : BitVec w} :
+    x &&& y = allOnes w ↔ x = allOnes w ∧ y = allOnes w := by
+  constructor
+  · intro h
+    constructor
+    all_goals
+    · ext i ih
+      have := BitVec.eq_of_getLsbD_eq_iff.mp h i ih
+      simp only [getLsbD_and, getLsbD_allOnes, ih, decide_true, Bool.and_eq_true] at this
+      simp [this, ih]
+  · intro h
+    simp [h]
+
 theorem extractLsb'_and {x y : BitVec w} {start len : Nat} :
    (x &&& y).extractLsb' start len = (x.extractLsb' start len) &&& (y.extractLsb' start len) := by
   ext i hi
@@ -1095,6 +1122,31 @@ instance : Std.LawfulCommIdentity (α := BitVec n) (· ^^^ · ) (0#n) where
   ext i
   simp
 
+@[simp]
+theorem xor_left_inj {x y : BitVec w} (z : BitVec w) : (x ^^^ z = y ^^^ z) ↔ x = y := by
+  constructor
+  · intro h
+    ext i ih
+    have := BitVec.eq_of_getLsbD_eq_iff.mp h i
+    simp only [getLsbD_xor, Bool.xor_left_inj] at this
+    exact this ih
+  · intro h
+    rw [h]
+
+@[simp]
+theorem xor_right_inj {x y : BitVec w} (z : BitVec w) : (z ^^^ x = z ^^^ y) ↔ x = y := by
+  rw [xor_comm z x, xor_comm z y]
+  exact xor_left_inj _
+
+@[simp]
+theorem xor_eq_zero_iff {x y : BitVec w} : (x ^^^ y = 0#w) ↔ x = y := by
+  constructor
+  · intro h
+    apply (xor_left_inj y).mp
+    rwa [xor_self]
+  · intro h
+    simp [h]
+
 theorem extractLsb'_xor {x y : BitVec w} {start len : Nat} :
    (x ^^^ y).extractLsb' start len = (x.extractLsb' start len) ^^^ (y.extractLsb' start len) := by
   ext i hi
diff --git a/src/Std/Tactic/BVDecide/Normalize/Equal.lean b/src/Std/Tactic/BVDecide/Normalize/Equal.lean
index 792e7a21ad30..c38adc93a72d 100644
--- a/src/Std/Tactic/BVDecide/Normalize/Equal.lean
+++ b/src/Std/Tactic/BVDecide/Normalize/Equal.lean
@@ -27,6 +27,46 @@ theorem BitVec.not_beq_not (a b : BitVec w) : (~~~a == ~~~b) = (a == b) := by
   rw [Bool.eq_iff_iff]
   simp
 
+@[bv_normalize]
+theorem BitVec.or_beq_zero_iff (a b : BitVec w) : (a ||| b == 0#w) = (a == 0#w && b == 0#w) := by
+  rw [Bool.eq_iff_iff]
+  simp
+
+@[bv_normalize]
+theorem BitVec.zero_beq_or_iff (a b : BitVec w) : (0#w == a ||| b) = (a == 0#w && b == 0#w) := by
+  rw [Bool.eq_iff_iff, beq_iff_eq, Eq.comm]
+  simp
+
+@[bv_normalize]
+theorem BitVec.xor_beq_zero_iff (a b : BitVec w) : (a ^^^ b == 0#w) = (a == b) := by
+  rw [Bool.eq_iff_iff]
+  simp
+
+@[bv_normalize]
+theorem BitVec.zero_beq_xor_iff (a b : BitVec w) : (0#w == a ^^^ b) = (a == b) := by
+  rw [Bool.eq_iff_iff, beq_iff_eq, Eq.comm]
+  simp
+
+@[bv_normalize]
+theorem BitVec.xor_left_inj (a b c : BitVec w) : (a ^^^ c == b ^^^ c) = (a == b) := by
+  rw [Bool.eq_iff_iff]
+  simp
+
+@[bv_normalize]
+theorem BitVec.xor_left_inj' (a b c : BitVec w) : (a ^^^ c == c ^^^ b) = (a == b) := by
+  rw [Bool.eq_iff_iff, BitVec.xor_comm c]
+  simp
+
+@[bv_normalize]
+theorem BitVec.xor_right_inj (a b c : BitVec w) : (c ^^^ a == c ^^^ b) = (a == b) := by
+  rw [Bool.eq_iff_iff]
+  simp
+
+@[bv_normalize]
+theorem BitVec.xor_right_inj' (a b c : BitVec w) : (c ^^^ a == b ^^^ c) = (a == b) := by
+  rw [Bool.eq_iff_iff, BitVec.xor_comm c]
+  simp
+
 @[bv_normalize]
 theorem BitVec.add_left_inj (a b c : BitVec w) : (a + c == b + c) = (a == b) := by
   rw [Bool.eq_iff_iff]
diff --git a/tests/lean/run/bv_axiom_check.lean b/tests/lean/run/bv_axiom_check.lean
index 2ef0d1732f2a..7e17946821c9 100644
--- a/tests/lean/run/bv_axiom_check.lean
+++ b/tests/lean/run/bv_axiom_check.lean
@@ -2,8 +2,7 @@ import Std.Tactic.BVDecide
 
 open BitVec
 
-theorem bv_axiomCheck (x y : BitVec 1) : x + y = y + x := by
-  bv_decide
+theorem bv_axiomCheck (x y z : BitVec 1) : x < y → y < z → x < z := by bv_decide
 
 /--
 info: 'bv_axiomCheck' depends on axioms: [propext, Classical.choice, Lean.ofReduceBool, Quot.sound]
diff --git a/tests/lean/run/bv_decide_rewriter.lean b/tests/lean/run/bv_decide_rewriter.lean
index 52ff640f2e1b..71bb9d7c8e00 100644
--- a/tests/lean/run/bv_decide_rewriter.lean
+++ b/tests/lean/run/bv_decide_rewriter.lean
@@ -131,6 +131,20 @@ example (x y : BitVec 16) : (x + y == y) = (x == 0) := by bv_normalize
 example (x y : BitVec 16) : (x == x + y) = (y == 0) := by bv_normalize
 example (x y : BitVec 16) : (x == y + x) = (y == 0) := by bv_normalize
 
+-- or_beq_zero_iff
+example (x y : BitVec 16) : (x ||| y == 0) = (x == 0 && y == 0) := by bv_normalize
+example (x y : BitVec 16) : (0 == x ||| y) = (x == 0 && y == 0) := by bv_normalize
+
+-- xor_beq_zero_iff
+example (x y : BitVec 16) : (x ^^^ y == 0) = (x == y) := by bv_normalize
+example (x y : BitVec 16) : (0 == x ^^^ y) = (x == y) := by bv_normalize
+
+-- xor_left_inj / xor_right_inj
+example (x y z : BitVec 16) : (x ^^^ z == y ^^^ z) = (x == y) := by bv_normalize
+example (x y z : BitVec 16) : (x ^^^ z == z ^^^ y) = (x == y) := by bv_normalize
+example (x y z : BitVec 16) : (z ^^^ x == y ^^^ z) = (x == y) := by bv_normalize
+example (x y z : BitVec 16) : (z ^^^ x == z ^^^ y) = (x == y) := by bv_normalize
+
 section
 
 example (x y : BitVec 256) : x * y = y * x := by