From 70386a12ccfe749a154b5750df1da5b6ca3f8aae Mon Sep 17 00:00:00 2001
From: Tobias Grosser <tobias@grosser.es>
Date: Sun, 29 Sep 2024 20:48:29 +0100
Subject: [PATCH] feat: add ac_nf and add test cases for BitVec

ac_nf is a counterpart to ac_rfl, which normalizes bitvector expressions
with respect to associativity and commutativity.
---
 src/Init/Tactics.lean             |  4 ++
 src/Lean/Meta/Tactic/AC/Main.lean | 41 ++++++++-----
 tests/lean/run/ac.lean            | 97 +++++++++++++++++++++++++++++++
 3 files changed, 129 insertions(+), 13 deletions(-)
 create mode 100644 tests/lean/run/ac.lean

diff --git a/src/Init/Tactics.lean b/src/Init/Tactics.lean
index ef7671ce2759..19ba9620fe9a 100644
--- a/src/Init/Tactics.lean
+++ b/src/Init/Tactics.lean
@@ -399,6 +399,10 @@ example (a b c d : Nat) : a + b + c + d = d + (b + c) + a := by ac_rfl
 -/
 syntax (name := acRfl) "ac_rfl" : tactic
 
+/--
+-/
+syntax (name := acNf) "ac_nf" : tactic
+
 /--
 The `sorry` tactic closes the goal using `sorryAx`. This is intended for stubbing out incomplete
 parts of a proof while still having a syntactically correct proof skeleton. Lean will give
diff --git a/src/Lean/Meta/Tactic/AC/Main.lean b/src/Lean/Meta/Tactic/AC/Main.lean
index eac00e2f5db8..e5d5170c91c1 100644
--- a/src/Lean/Meta/Tactic/AC/Main.lean
+++ b/src/Lean/Meta/Tactic/AC/Main.lean
@@ -140,19 +140,7 @@ where
     | .op l r => mkApp2 preContext.op (convertTarget vars l) (convertTarget vars r)
     | .var x => vars[x]!
 
-def rewriteUnnormalized (mvarId : MVarId) : MetaM Unit := do
-  let simpCtx :=
-    {
-      simpTheorems  := {}
-      congrTheorems := (← getSimpCongrTheorems)
-      config        := Simp.neutralConfig
-    }
-  let tgt ← instantiateMVars (← mvarId.getType)
-  let (res, _) ← Simp.main tgt simpCtx (methods := { post })
-  let newGoal ← applySimpResultToTarget mvarId tgt res
-  newGoal.refl
-where
-  post (e : Expr) : SimpM Simp.Step := do
+def post (e : Expr) : SimpM Simp.Step := do
     let ctx ← Simp.getContext
     match e, ctx.parent? with
     | bin op₁ l r, some (bin op₂ _ _) =>
@@ -171,10 +159,37 @@ where
       | none => return Simp.Step.done { expr := e }
     | e, _ => return Simp.Step.done { expr := e }
 
+def rewriteUnnormalized (mvarId : MVarId) : MetaM Unit := do
+  let simpCtx :=
+    {
+      simpTheorems  := {}
+      congrTheorems := (← getSimpCongrTheorems)
+      config        := Simp.neutralConfig
+    }
+  let tgt ← instantiateMVars (← mvarId.getType)
+  let (res, _) ← Simp.main tgt simpCtx (methods := { post })
+  let newGoal ← applySimpResultToTarget mvarId tgt res
+  newGoal.refl
+
 @[builtin_tactic acRfl] def acRflTactic : Lean.Elab.Tactic.Tactic := fun _ => do
   let goal ← getMainGoal
   goal.withContext <| rewriteUnnormalized goal
 
+def canonicalizeUnnormalized (mvarId : MVarId) : TacticM Unit := do
+  let simpCtx :=
+    {
+      simpTheorems  := {}
+      congrTheorems := (← getSimpCongrTheorems)
+      config        := Simp.neutralConfig
+    }
+  let tgt ← instantiateMVars (← mvarId.getType)
+  let (res, _) ← Simp.main tgt simpCtx (methods := { post })
+  replaceMainGoal [← applySimpResultToTarget mvarId tgt res]
+
+@[builtin_tactic acNf] def acNfTactic : Lean.Elab.Tactic.Tactic := fun _ => do
+  let goal ← getMainGoal
+  goal.withContext <| canonicalizeUnnormalized goal
+
 builtin_initialize
   registerTraceClass `Meta.AC
 
diff --git a/tests/lean/run/ac.lean b/tests/lean/run/ac.lean
new file mode 100644
index 000000000000..80f91d889781
--- /dev/null
+++ b/tests/lean/run/ac.lean
@@ -0,0 +1,97 @@
+/- arithmetic | commutativity -/
+
+example (a b c d : BitVec w) :
+    a * b * c * d = d * c * b * a := by
+  ac_nf
+  rfl
+
+example (a b c d : BitVec w) :
+    a * b * c * d = d * c * b * a := by
+  ac_rfl
+
+example (a b c d : BitVec w) :
+    a + b + c + d = d + c + b + a := by
+  ac_nf
+  rfl
+
+example (a b c d : BitVec w) :
+    a + b + c + d = d + c + b + a := by
+  ac_rfl
+
+/- arithmetic | associativity -/
+
+example (a b c d : BitVec w) :
+    a * (b * (c * d)) = ((a * b) * c) * d := by
+  ac_nf
+  rfl
+
+example (a b c d : BitVec w) :
+    a * (b * (c * d)) = ((a * b) * c) * d := by
+  ac_rfl
+
+example (a b c d : BitVec w) :
+    a + (b + (c + d)) = ((a + b) + c) + d := by
+  ac_nf
+  rfl
+
+example (a b c d : BitVec w) :
+    a + (b + (c + d)) = ((a + b) + c) + d := by
+  ac_rfl
+
+/- bitwise operations | commutativity -/
+
+example (a b c d : BitVec w) :
+    a ^^^ b ^^^ c ^^^ d = d ^^^ c ^^^ b ^^^ a := by
+  ac_nf
+  rfl
+
+example (a b c d : BitVec w) :
+    a ^^^ b ^^^ c ^^^ d = d ^^^ c ^^^ b ^^^ a := by
+  ac_rfl
+
+example (a b c d : BitVec w) :
+    a &&& b &&& c &&& d = d &&& c &&& b &&& a := by
+  ac_nf
+  rfl
+
+example (a b c d : BitVec w) :
+    a &&& b &&& c &&& d = d &&& c &&& b &&& a := by
+  ac_rfl
+
+example (a b c d : BitVec w) :
+    a ||| b ||| c ||| d = d ||| c ||| b ||| a := by
+  ac_nf
+  rfl
+
+example (a b c d : BitVec w) :
+    a ||| b ||| c ||| d = d ||| c ||| b ||| a := by
+  ac_rfl
+
+/- bitwise operations | associativity -/
+
+example (a b c d : BitVec w) :
+    a &&& (b &&& (c &&& d)) = ((a &&& b) &&& c) &&& d := by
+  ac_nf
+  rfl
+
+example (a b c d : BitVec w) :
+    a &&& (b &&& (c &&& d)) = ((a &&& b) &&& c) &&& d := by
+  ac_rfl
+
+example (a b c d : BitVec w) :
+    a ||| (b ||| (c ||| d)) = ((a ||| b) ||| c) ||| d := by
+  ac_nf
+  rfl
+
+example (a b c d : BitVec w) :
+    a ||| (b ||| (c ||| d)) = ((a ||| b) ||| c) ||| d := by
+  ac_rfl
+
+example (a b c d : BitVec w) :
+    a ^^^ (b ^^^ (c ^^^ d)) = ((a ^^^ b) ^^^ c) ^^^ d := by
+  ac_nf
+  rfl
+
+example (a b c d : BitVec w) :
+    a ^^^ (b ^^^ (c ^^^ d)) = ((a ^^^ b) ^^^ c) ^^^ d := by
+  ac_rfl