From 5c97ad326f57cff694c45ae52e999e1452c589f7 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Wed, 21 Feb 2024 09:52:07 -0800
Subject: [PATCH] fix: support for `Fin` and `BitVec` literal normalization

---
 .../Tactic/Simp/BuiltinSimprocs/BitVec.lean   |  9 ++++++++
 .../Meta/Tactic/Simp/BuiltinSimprocs/Fin.lean | 13 ++++++++----
 tests/lean/run/bitvec_fin_literal_norm.lean   | 21 +++++++++++++++++++
 3 files changed, 39 insertions(+), 4 deletions(-)
 create mode 100644 tests/lean/run/bitvec_fin_literal_norm.lean

diff --git a/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/BitVec.lean b/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/BitVec.lean
index 72946990dbc5..13e80a619b18 100644
--- a/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/BitVec.lean
+++ b/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/BitVec.lean
@@ -225,6 +225,15 @@ builtin_simproc [simp, seval] reduceOfInt (BitVec.ofInt _ _) := fun e => do
   let lit : Literal := { n, value := BitVec.ofInt n i }
   return .done { expr := lit.toExpr }
 
+/-- Simplification procedure for ensuring `BitVec.ofNat` literals are normalized. -/
+builtin_simproc [simp, seval] reduceOfNat (BitVec.ofNat _ _) := fun e => do
+  unless e.isAppOfArity ``BitVec.ofNat 2 do return .continue
+  let some n ← Nat.fromExpr? e.appFn!.appArg! | return .continue
+  let some v ← Nat.fromExpr? e.appArg! | return .continue
+  let bv := BitVec.ofNat n v
+  if bv.toNat == v then return .continue -- already normalized
+  return .done { expr := { n, value := BitVec.ofNat n v : Literal }.toExpr }
+
 /-- Simplification procedure for `<` on `BitVec`s. -/
 builtin_simproc [simp, seval] reduceLT (( _ : BitVec _) < _)  := reduceBinPred ``LT.lt 4 (· < ·)
 /-- Simplification procedure for `≤` on `BitVec`s. -/
diff --git a/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/Fin.lean b/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/Fin.lean
index 9165bc5b35b7..80cd362891ef 100644
--- a/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/Fin.lean
+++ b/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/Fin.lean
@@ -69,9 +69,14 @@ builtin_simproc [simp, seval] reduceNe  (( _ : Fin _) ≠ _)  := reduceBinPred `
 builtin_simproc [simp, seval] reduceBEq  (( _ : Fin _) == _)  := reduceBoolPred ``BEq.beq 4 (. == .)
 builtin_simproc [simp, seval] reduceBNe  (( _ : Fin _) != _)  := reduceBoolPred ``bne 4 (. != .)
 
-/-- Return `.done` for Fin values. We don't want to unfold in the symbolic evaluator. -/
-builtin_simproc [seval] isValue ((OfNat.ofNat _ : Fin _)) := fun e => do
-  unless e.isAppOfArity ``OfNat.ofNat 3 do return .continue
-  return .done { expr := e }
+/-- Simplification procedure for ensuring `Fin` literals are normalized. -/
+builtin_simproc [simp, seval] isValue ((OfNat.ofNat _ : Fin _)) := fun e => do
+  let some v ← fromExpr? e | return .continue
+  let v' ← Nat.fromExpr? e.appFn!.appArg!
+  if v.value == v' then
+    -- Design decision: should we return `.continue` instead of `.done` when simplifying.
+    -- In the symbolic evaluator, we must return `.done`, otherwise it will unfold the `OfNat.ofNat`
+    return .done { expr := e }
+  return .done { expr := v.toExpr }
 
 end Fin
diff --git a/tests/lean/run/bitvec_fin_literal_norm.lean b/tests/lean/run/bitvec_fin_literal_norm.lean
new file mode 100644
index 000000000000..662c595ce50e
--- /dev/null
+++ b/tests/lean/run/bitvec_fin_literal_norm.lean
@@ -0,0 +1,21 @@
+open Std BitVec
+
+example : (5 : Fin 4) = x := by
+  simp
+  guard_target =ₛ 1 = x
+  sorry
+
+example : (1 : Fin 4) = x := by
+  fail_if_success simp
+  guard_target =ₛ 1 = x
+  sorry
+
+example : 17#4 = x := by
+  simp
+  guard_target =ₛ  1#4 = x
+  sorry
+
+example : (17 : BitVec 4) = x := by
+  simp
+  guard_target =ₛ  1#4 = x
+  sorry