From 489f2da711d3b65fe712d108df13f5ffe804f870 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Mon, 19 Feb 2024 15:15:37 -0800
Subject: [PATCH] feat: add simproc for `BitVec.signExtend` (#3409)

---
 .../Tactic/Simp/BuiltinSimprocs/BitVec.lean   | 19 +++++++++++++------
 tests/lean/run/bitvec_simproc.lean            |  4 ++++
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/BitVec.lean b/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/BitVec.lean
index 73d3ea54b6..72946990db 100644
--- a/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/BitVec.lean
+++ b/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/BitVec.lean
@@ -77,6 +77,15 @@ Helper function for reducing homogenous binary bitvector operators.
   else
     return .continue
 
+/-- Simplification procedure for `zeroExtend` and `signExtend` on `BitVec`s. -/
+@[inline] def reduceExtend (declName : Name)
+    (op : {n : Nat} → (m : Nat) → BitVec n → BitVec m) (e : Expr) : SimpM Step := do
+  unless e.isAppOfArity declName 3 do return .continue
+  let some v ← fromExpr? e.appArg! | return .continue
+  let some n ← Nat.fromExpr? e.appFn!.appArg! | return .continue
+  let lit : Literal := { n, value := op n v.value }
+  return .done { expr := lit.toExpr }
+
 /--
 Helper function for reducing bitvector functions such as `getLsb` and `getMsb`.
 -/
@@ -275,12 +284,10 @@ builtin_simproc [simp, seval] reduceReplicate (replicate _ _) := fun e => do
   return .done { expr := lit.toExpr }
 
 /-- Simplification procedure for `zeroExtend` on `BitVec`s. -/
-builtin_simproc [simp, seval] reduceZeroExtend (zeroExtend _ _) := fun e => do
-  unless e.isAppOfArity ``zeroExtend 3 do return .continue
-  let some v ← fromExpr? e.appArg! | return .continue
-  let some n ← Nat.fromExpr? e.appFn!.appArg! | return .continue
-  let lit : Literal := { n, value := v.value.zeroExtend n }
-  return .done { expr := lit.toExpr }
+builtin_simproc [simp, seval] reduceZeroExtend (zeroExtend _ _) := reduceExtend ``zeroExtend zeroExtend
+
+/-- Simplification procedure for `signExtend` on `BitVec`s. -/
+builtin_simproc [simp, seval] reduceSignExtend (signExtend _ _) := reduceExtend ``signExtend signExtend
 
 /-- Simplification procedure for `allOnes` -/
 builtin_simproc [simp, seval] reduceAllOnes (allOnes _) := fun e => do
diff --git a/tests/lean/run/bitvec_simproc.lean b/tests/lean/run/bitvec_simproc.lean
index 613b372924..465a0ed21f 100644
--- a/tests/lean/run/bitvec_simproc.lean
+++ b/tests/lean/run/bitvec_simproc.lean
@@ -109,3 +109,7 @@ example (h : x = (9: Std.BitVec 6)) : x = (1#3).replicate 2 := by
   simp; guard_target =ₛ x = 9#6; assumption
 example (h : x = (5 : Std.BitVec 7)) : x = (5#3).zeroExtend 7 := by
   simp; guard_target =ₛ x = 5#7; assumption
+example (h : -5#10 = x) : signExtend 10 (-5#8) = x := by
+  simp; guard_target =ₛ1019#10 = x; assumption
+example (h : 5#10 = x) : -signExtend 10 (-5#8) = x := by
+  simp; guard_target =ₛ5#10 = x; assumption