From bb47469d1a9debf5e451e1799b8a151b26914d04 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Henrik=20B=C3=B6ving?= <hargonix@gmail.com>
Date: Tue, 11 Mar 2025 11:09:16 +0100
Subject: [PATCH] feat: add simprocs for turning shifts by constants into
 extracts to bv_decide (#7436)

This PR adds simprocs that turn left and right shifts by constants into
extracts to bv_decide.
---
 .../BVDecide/Frontend/Normalize/Simproc.lean  | 38 +++++++++++++++++++
 tests/lean/run/bv_decide_rewriter.lean        | 12 ++++++
 2 files changed, 50 insertions(+)

diff --git a/src/Lean/Elab/Tactic/BVDecide/Frontend/Normalize/Simproc.lean b/src/Lean/Elab/Tactic/BVDecide/Frontend/Normalize/Simproc.lean
index 298be7982f..19694033d5 100644
--- a/src/Lean/Elab/Tactic/BVDecide/Frontend/Normalize/Simproc.lean
+++ b/src/Lean/Elab/Tactic/BVDecide/Frontend/Normalize/Simproc.lean
@@ -225,5 +225,43 @@ builtin_simproc [bv_normalize] bv_mul_ones ((_ : BitVec _) * (BitVec.ofNat _ _))
   let proof := mkApp2 (mkConst ``BitVec.mul_ones) (toExpr w) lhs
   return .visit { expr := expr, proof? := some proof }
 
+builtin_simproc [bv_normalize] bv_elim_ushiftRight_const ((_ : BitVec _) >>> (_ : Nat)) := fun e => do
+  let_expr HShiftRight.hShiftRight bvType _ _ _ lhsExpr rhsExpr := e | return .continue
+  let some rhs ← getNatValue? rhsExpr | return .continue
+  let_expr BitVec wExpr := bvType | return .continue
+  let some w ← getNatValue? wExpr | return .continue
+  if rhs < w then
+    let zero := toExpr 0#rhs
+    let extract := mkApp4 (mkConst ``BitVec.extractLsb') wExpr rhsExpr (toExpr (w - rhs)) lhsExpr
+    let concat ← mkAppM ``HAppend.hAppend #[zero, extract]
+    let expr := mkApp4 (mkConst ``BitVec.cast) wExpr wExpr (← mkEqRefl wExpr) concat
+    let h ← mkDecideProof (← mkLt rhsExpr wExpr)
+    let proof := mkApp4 (mkConst ``BitVec.ushiftRight_eq_extractLsb'_of_lt) wExpr lhsExpr rhsExpr h
+    return .done { expr := expr, proof? := some proof }
+  else
+    let expr := toExpr 0#w
+    let h ← mkDecideProof (← mkLe wExpr rhsExpr)
+    let proof := mkApp4 (mkConst ``BitVec.ushiftRight_eq_zero) wExpr lhsExpr rhsExpr h
+    return .done { expr := expr, proof? := some proof }
+
+builtin_simproc [bv_normalize] bv_elim_shiftLeft_const ((_ : BitVec _) <<< (_ : Nat)) := fun e => do
+  let_expr HShiftLeft.hShiftLeft bvType _ _ _ lhsExpr rhsExpr := e | return .continue
+  let some rhs ← getNatValue? rhsExpr | return .continue
+  let_expr BitVec wExpr := bvType | return .continue
+  let some w ← getNatValue? wExpr | return .continue
+  if rhs < w then
+    let zero := toExpr 0#rhs
+    let extract := mkApp4 (mkConst ``BitVec.extractLsb') wExpr (toExpr 0) (toExpr (w - rhs)) lhsExpr
+    let concat ← mkAppM ``HAppend.hAppend #[extract, zero]
+    let expr := mkApp4 (mkConst ``BitVec.cast) wExpr wExpr (← mkEqRefl wExpr) concat
+    let h ← mkDecideProof (← mkLt rhsExpr wExpr)
+    let proof := mkApp4 (mkConst ``BitVec.shiftLeft_eq_concat_of_lt) wExpr lhsExpr rhsExpr h
+    return .done { expr := expr, proof? := some proof }
+  else
+    let expr := toExpr 0#w
+    let h ← mkDecideProof (← mkLe wExpr rhsExpr)
+    let proof := mkApp4 (mkConst ``BitVec.shiftLeft_eq_zero) wExpr lhsExpr rhsExpr h
+    return .done { expr := expr, proof? := some proof }
+
 end Frontend.Normalize
 end Lean.Elab.Tactic.BVDecide
diff --git a/tests/lean/run/bv_decide_rewriter.lean b/tests/lean/run/bv_decide_rewriter.lean
index be057b3c27..e57fbe3de5 100644
--- a/tests/lean/run/bv_decide_rewriter.lean
+++ b/tests/lean/run/bv_decide_rewriter.lean
@@ -549,6 +549,18 @@ example {a : BitVec 8} : -1 * a = -a := by bv_normalize
 example {a : BitVec 8} : -1 * a + a = 0 := by bv_normalize
 example {a : BitVec 8} : a + -1 * a = 0 := by bv_normalize
 
+-- SHR_CONST
+example {a : BitVec 8} : a >>> 1 = 0#1 ++ BitVec.extractLsb' 1 7 a := by bv_normalize
+example {a : BitVec 8} : a >>> 3 = 0#3 ++ BitVec.extractLsb' 3 5 a := by bv_normalize
+example {a : BitVec 8} : a >>> 8 = 0 := by bv_normalize
+example {a : BitVec 8} : a >>> 12 = 0 := by bv_normalize
+
+-- SHL_CONST
+example {a : BitVec 8} : a <<< 1 = BitVec.extractLsb' 0 7 a ++ 0#1 := by bv_normalize
+example {a : BitVec 8} : a <<< 3 = BitVec.extractLsb' 0 5 a ++ 0#3 := by bv_normalize
+example {a : BitVec 8} : a <<< 8 = 0 := by bv_normalize
+example {a : BitVec 8} : a <<< 12 = 0 := by bv_normalize
+
 section
 
 example (x y : BitVec 256) : x * y = y * x := by