feat: ac grind action

Author: leodemoura

src/Lean/Elab/Tactic/Grind/BuiltinTactic.lean

@@ -101,7 +101,7 @@ def evalCheck (tacticName : Name) (k : GoalM Bool)
   evalCheck `ring Arith.CommRing.check Arith.CommRing.pp?
 
 @[builtin_grind_tactic ac] def evalAC : GrindTactic := fun _ => do
-  evalCheck `ac AC.check AC.pp?
+  evalCheck `ac AC.check' AC.pp?
 
 def logTheoremAnchor (proof : Expr) : TermElabM Unit := do
   let stx ← getRef

src/Lean/Meta/Tactic/Grind/AC.lean

@@ -17,6 +17,7 @@ public import Lean.Meta.Tactic.Grind.AC.ToExpr
 public import Lean.Meta.Tactic.Grind.AC.VarRename
 public import Lean.Meta.Tactic.Grind.AC.PP
 public import Lean.Meta.Tactic.Grind.AC.Inv
+public import Lean.Meta.Tactic.Grind.AC.Action
 public section
 namespace Lean.Meta.Grind.AC
 builtin_initialize registerTraceClass `grind.ac
@@ -36,7 +37,7 @@ builtin_initialize
     (internalize := AC.internalize)
     (newEq       := AC.processNewEq)
     (newDiseq    := AC.processNewDiseq)
-    (check       := AC.check)
+    (check       := AC.check')
     (checkInv    := AC.checkInvariants)
     (mkTactic?   := return some (← `(grind| ac)))
 

src/Lean/Meta/Tactic/Grind/AC/Action.lean

@@ -0,0 +1,21 @@
+/-
+Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+module
+prelude
+public import Lean.Meta.Tactic.Grind.Action
+import Lean.Meta.Tactic.Grind.AC.Eq
+namespace Lean.Meta.Grind.Action
+
+/-- Linear integer arithmetic action. -/
+public def ac : Action := fun goal kna kp => do
+  let (result, goal') ← GoalM.run goal AC.check
+  match result with
+  | .none       => kna goal'
+  | .progress   => kp goal'
+  | .propagated => concatTactic (← kp goal') `(grind| ac)
+  | .closed     => closeWith `(grind| ac)
+
+end Lean.Meta.Grind.Action

src/Lean/Meta/Tactic/Grind/AC/Eq.lean

@@ -6,6 +6,7 @@ Authors: Leonardo de Moura
 module
 prelude
 public import Lean.Meta.Tactic.Grind.AC.Util
+public import Lean.Meta.Tactic.Grind.CheckResult
 import Lean.Meta.Tactic.Grind.AC.DenoteExpr
 import Lean.Meta.Tactic.Grind.AC.Proof
 import Lean.Meta.Tactic.Grind.AC.Seq
@@ -445,60 +446,65 @@ def mkEqData (e : Expr) (r : AC.Expr) : ACM EqData := do
 
 abbrev PropagateEqMap := Std.HashMap AC.Seq EqData
 
-private def propagateEqs : ACM Unit := do
-  if (← isInconsistent) then return ()
+private def propagateEqs : ACM Bool := do
+  if (← isInconsistent) then return false
   /-
   This is a very simple procedure that does not use any indexing data-structure.
   We don't even cache the simplified expressions.
   TODO: optimize
   -/
-  let mut map : PropagateEqMap := {}
-  for e in (← getStruct).vars do
-    if (← checkMaxSteps) then return ()
-    let r ← asACExpr e
-    map ← process map e r
-  for (e, r) in (← getStruct).denoteEntries do
-    if (← checkMaxSteps) then return ()
-    map ← process map e r
+  let go : StateT (Bool × PropagateEqMap) ACM Unit := do
+    for e in (← getStruct).vars do
+      if (← checkMaxSteps) then return
+      let r ← asACExpr e
+      process e r
+    for (e, r) in (← getStruct).denoteEntries do
+      if (← checkMaxSteps) then return
+      process e r
+  let (_, (propagated, _)) ← go.run (false, {})
+  return propagated
 where
-  process (map : PropagateEqMap) (e : Expr) (r : AC.Expr) : ACM PropagateEqMap := do
+  process (e : Expr) (r : AC.Expr) : StateT (Bool × PropagateEqMap) ACM Unit := do
     let d ← mkEqData e r
     let s := d.c.rhs
     trace[grind.debug.ac.eq] "{e}, s: {← s.denoteExpr}"
-    if let some d' := map[s]? then
+    if let some d' := (← get).2[s]? then
       trace[grind.debug.ac.eq] "found [{← isEqv d.e d'.e}]: {d.e}, {d'.e}"
       unless (← isEqv d.e d'.e) do
         propagateEq d.e d'.e d.r d'.r d.c d'.c
-      return map
+        modify fun s => (true, s.2)
     else
-      return map.insert s d
+      modify fun (propagated, map) => (propagated, map.insert s d)
 
-private def checkStruct : ACM Bool := do
-  unless (← needCheck) do return false
+private def checkStruct : ACM CheckResult := do
+  unless (← needCheck) do return .none
   trace_goal[grind.debug.ac.check] "{(← getStruct).op}"
   repeat
     checkSystem "ac"
     let some c ← getNext? | break
     trace_goal[grind.debug.ac.check] "{← c.denoteExpr}"
     c.addToBasis
-    if (← isInconsistent) then return true
-    if (← checkMaxSteps) then return true
+    if (← isInconsistent) then return .closed
+    if (← checkMaxSteps) then return .progress
   checkDiseqs
-  propagateEqs
   modifyStruct fun s => { s with recheck := false }
-  return true
+  if (← propagateEqs) then return .propagated
+  return .progress
 
-def check : GoalM Bool := do profileitM Exception "grind ac" (← getOptions) do
-  if (← checkMaxSteps) then return false
-  let mut progress := false
+def check : GoalM CheckResult := do profileitM Exception "grind ac" (← getOptions) do
+  if (← checkMaxSteps) then return .none
+  let mut result : CheckResult := .none
   checkInvariants
   try
     for opId in *...(← get').structs.size do
       let r ← ACM.run opId checkStruct
-      progress := progress || r
-      if (← isInconsistent) then return true
-    return progress
+      result := result.join r
+      if (← isInconsistent) then return .closed
+    return result
   finally
     checkInvariants
 
+def check' : GoalM Bool :=
+  return (← check) != .none
+
 end Lean.Meta.Grind.AC

src/Lean/Meta/Tactic/Grind/Action.lean

@@ -224,6 +224,13 @@ def concatTactic (r : ActionResult) (mk : GrindM (TSyntax `grind)) : GrindM Acti
   else
     return r
 
+/-- Returns `.closed [← mk]` if tracing is enabled, and `.closed []` otherwise. -/
+def closeWith (mk : GrindM (TSyntax `grind)) : GrindM ActionResult := do
+  if (← getConfig).trace then
+    return .closed [(← mk)]
+  else
+    return .closed []
+
 /--
 A terminal action which closes the goal or not.
 This kind of action may make progress, but we only include `mkTac` into the resulting tactic sequence
@@ -232,16 +239,10 @@ if it closed the goal.
 public def terminalAction (check : GoalM Bool) (mkTac : GrindM (TSyntax `grind)) : Action := fun goal kna kp => do
   let (progress, goal') ← GoalM.run goal check
   if progress then
-    let r ← kp goal'
-    /-
-    **Note**: terminal actions may make progress by computing a valid assignment satisfying all constraints.
-    That said, unless it closed the goal, it is pointless to include the tactic in the resulting
-    sequence because no information has been communicated to the other solvers.
-    -/
     if goal'.inconsistent then
-      concatTactic r mkTac
+      closeWith mkTac
     else
-      return r
+      kp goal'
   else
     kna goal'
 

src/Lean/Meta/Tactic/Grind/CheckResult.lean

@@ -0,0 +1,68 @@
+/-
+Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+module
+prelude
+public import Lean.Meta.Tactic.Grind.Types
+public section
+namespace Lean.Meta.Grind
+/--
+Result type for satisfiability checking procedures.
+-/
+inductive CheckResult where
+  | /-- No progress -/
+    none
+  | /-- Updated basis, simplified equations. -/
+    progress
+  | /-- Propagated equations back to the core. -/
+    propagated
+  | /-- Closed the goal. -/
+    closed
+  deriving BEq, Inhabited, Repr
+
+def CheckResult.lt (r₁ r₂ : CheckResult) : Bool :=
+  match r₁, r₂ with
+  | _,          .none        => false
+  | .none,       _           => true
+  | _,           .progress   => false
+  | .progress,   _           => true
+  | _,           .propagated => false
+  | .propagated, _           => true
+  | .closed,     .closed     => false
+
+def CheckResult.le (r₁ r₂ : CheckResult) : Bool :=
+  r₁ == r₂ || r₁.lt r₂
+
+/--
+Joins two results. It uses the order `.none < .progress < .propagated < .closed`
+-/
+def CheckResult.join (r₁ r₂ : CheckResult) : CheckResult :=
+  match r₁, r₂ with
+  | .none,       _           => r₂
+  | _,           .none       => r₁
+  | .progress,   _           => r₂
+  | _,           .progress   => r₁
+  | .propagated, _           => r₂
+  | _,           .propagated => r₁
+  | .closed,     .closed     => .closed
+
+/-! Sanity check theorems -/
+section
+open CheckResult
+example : lt .none .progress := rfl
+example : lt .progress .propagated := rfl
+example : lt .propagated .closed := rfl
+example {x} : lt x x = false := by cases x <;> rfl
+example {x y z} : lt x y → lt y z → lt x z := by cases x <;> cases y <;> cases z <;> decide
+example {x y z} : le x y → le y z → le x z := by cases x <;> cases y <;> cases z <;> decide
+example {x y} : le x y ↔ x = y ∨ lt x y := by cases x <;> cases y <;> simp +decide
+example {x} : le x x := by cases x <;> rfl
+example {x y} : le x y → le y x → x = y := by cases x <;> cases y <;> simp +decide
+example {x y} : le x (join x y) := by cases x <;> cases y <;> rfl
+example {y x} : le y (join x y) := by cases x <;> cases y <;> rfl
+example {x} : join x x = x := by cases x <;> rfl
+end
+
+end Lean.Meta.Grind