feat: extra critical pairs for AC + idempotent

src/Init/Grind/AC.lean

@@ -495,7 +495,7 @@ noncomputable def superpose_ac_cert (r₁ c r₂ lhs₁ rhs₁ lhs₂ rhs₂ lhs
 Given `lhs₁ = rhs₁` and `lhs₂ = rhs₂` where `lhs₁ := union c r₁` and `lhs₂ := union c r₂`,
 `lhs = rhs` where `lhs := union r₂ rhs₁` and `rhs := union r₁ rhs₂`
 -/
-theorem superpose_ac {α} (ctx : Context α) {inst₁ : Std.Associative ctx.op} {inst₂ : Std.Commutative ctx.op}  (r₁ c r₂ lhs₁ rhs₁ lhs₂ rhs₂ lhs rhs : Seq)
+theorem superpose_ac {α} (ctx : Context α) {inst₁ : Std.Associative ctx.op} {inst₂ : Std.Commutative ctx.op} (r₁ c r₂ lhs₁ rhs₁ lhs₂ rhs₂ lhs rhs : Seq)
     : superpose_ac_cert r₁ c r₂ lhs₁ rhs₁ lhs₂ rhs₂ lhs rhs → lhs₁.denote ctx = rhs₁.denote ctx → lhs₂.denote ctx = rhs₂.denote ctx
       → lhs.denote ctx = rhs.denote ctx := by
   simp [superpose_ac_cert]; intro _ _ _ _; subst lhs₁ lhs₂ lhs rhs; simp
@@ -506,6 +506,46 @@ theorem superpose_ac {α} (ctx : Context α) {inst₁ : Std.Associative ctx.op}
   apply congrArg (ctx.op (c.denote ctx))
   rw [Std.Commutative.comm (self := inst₂) (r₂.denote ctx)]
 
+noncomputable def Seq.contains_k (s : Seq) (x : Var) : Bool :=
+  Seq.rec (fun y => Nat.beq x y) (fun y _ ih => Bool.or' (Nat.beq x y) ih) s
+
+theorem Seq.contains_k_var (y x : Var) : Seq.contains_k (.var y) x = (x == y) := by
+  simp [Seq.contains_k]; rw [Bool.eq_iff_iff]; simp
+
+theorem Seq.contains_k_cons (y x : Var) (s : Seq) : Seq.contains_k (.cons y s) x = (x == y || s.contains_k x)  := by
+  show (Nat.beq x y |>.or' (s.contains_k x)) = (x == y || s.contains_k x)
+  simp; rw [Bool.eq_iff_iff]; simp
+
+attribute [local simp] Seq.contains_k_var Seq.contains_k_cons
+
+theorem Seq.denote_insert_of_contains {α} (ctx : Context α) {inst₁ : Std.Associative ctx.op} {inst₂ : Std.Commutative ctx.op} {inst₃ : Std.IdempotentOp ctx.op}
+    (s : Seq) (x : Var) : s.contains_k x → (s.insert x).denote ctx = s.denote ctx := by
+  induction s
+  next => simp; intro; subst x; rw [Std.IdempotentOp.idempotent (self := inst₃)]
+  next y s ih =>
+    simp; intro h; cases h
+    next => subst x; rw [← Std.Associative.assoc (self := inst₁), Std.IdempotentOp.idempotent (self := inst₃)]
+    next h =>
+      replace ih := ih h
+      simp at ih
+      rw [← Std.Associative.assoc (self := inst₁), Std.Commutative.comm (self := inst₂) (x.denote ctx)]
+      rw [Std.Associative.assoc (self := inst₁), ih]
+
+noncomputable def superpose_ac_idempotent_cert (x : Var) (lhs₁ rhs₁ rhs : Seq) : Bool :=
+  lhs₁.contains_k x |>.and' (rhs.beq' (rhs₁.insert x))
+
+/-!
+Remark: see Section 4.1 of the paper "MODULARITY, COMBINATION, AC CONGRUENCE CLOSURE" to understand why
+`superpose_ac_idempotent` is needed.
+-/
+
+theorem superpose_ac_idempotent {α} (ctx : Context α) {inst₁ : Std.Associative ctx.op} {inst₂ : Std.Commutative ctx.op} {inst₃ : Std.IdempotentOp ctx.op}
+    (x : Var) (lhs₁ rhs₁ rhs : Seq) : superpose_ac_idempotent_cert x lhs₁ rhs₁ rhs → lhs₁.denote ctx = rhs₁.denote ctx → lhs₁.denote ctx = rhs.denote ctx := by
+  simp [superpose_ac_idempotent_cert]; intro h₁ _ h₂; subst rhs
+  replace h₂ : Seq.denote ctx (lhs₁.insert x) = Seq.denote ctx (rhs₁.insert x) := by
+    simp [h₂]
+  rw [← h₂, Seq.denote_insert_of_contains ctx lhs₁ x h₁] <;> assumption
+
 noncomputable def eq_norm_a_cert (lhs rhs : Expr) (lhs' rhs' : Seq) : Bool :=
   lhs.toSeq.beq' lhs' |>.and' (rhs.toSeq.beq' rhs')
 

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

@@ -254,8 +254,37 @@ private def EqCnstr.superposeWithA (c₁ : EqCnstr) : ACM Unit := do
       c₁.superposeA c₂
       c₂.superposeA c₁
 
+/--
+If the operator is idempotent, we have to add extra critical pairs.
+See section 4.1 of the paper "MODULARITY, COMBINATION, AC CONGRUENCE CLOSURE"
+The idea is the following, given `c` of the form `lhs = rhs`,
+for each variable `x` in `lhs` s.t. `x` is not in `rhs`, we add the equation
+`lhs = rhs.union {x}`
+Note that the paper does not include `x` is not in `rhs`, but this extra filter is correct
+since after normalization and simplification `lhs = rhs.union {x}` would be discarded.
+-/
+private def EqCnstr.superposeAC_Idempotent (c : EqCnstr) : ACM Unit := do
+  go c.lhs
+where
+  goVar (x : AC.Var) : ACM Unit := do
+    unless c.rhs.contains x do
+      let c := { c with rhs := c.rhs.insert x, h := .superpose_ac_idempotent x c }
+      if (← hasNeutral) && c.rhs.contains 0 then
+        (← c.erase0).addToQueue
+      else
+        c.addToQueue
+
+  go (s : AC.Seq) : ACM Unit := do
+    match s with
+    | .var x => goVar x
+    | .cons x s => goVar x; go s
+
 private def EqCnstr.superposeWith (c : EqCnstr) : ACM Unit := do
-  if (← isCommutative) then c.superposeWithAC else c.superposeWithA
+  if (← isCommutative) then
+    c.superposeWithAC
+    if (← isIdempotent) then c.superposeAC_Idempotent
+  else
+    c.superposeWithA
 
 private def EqCnstr.simplifyBasis (c : EqCnstr) : ACM Unit := do
   let rec go (basis : List EqCnstr) (acc : List EqCnstr) : ACM (List EqCnstr) := do

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

@@ -17,6 +17,7 @@ open Lean.Grind
 
 structure ProofM.State where
   cache      : Std.HashMap UInt64 Expr := {}
+  varDecls   : Std.HashMap AC.Var Expr := {}
   exprDecls  : Std.HashMap AC.Expr Expr := {}
   seqDecls   : Std.HashMap AC.Seq Expr := {}
 
@@ -45,6 +46,9 @@ local macro "declare! " decls:ident a:ident : term =>
        modify fun s => { s with $decls:ident := (s.$decls).insert $a x };
        return x)
 
+private def mkVarDecl (x : AC.Var) : ProofM Expr := do
+  declare! varDecls x
+
 private def mkSeqDecl (s : AC.Seq) : ProofM Expr := do
   declare! seqDecls s
 
@@ -83,6 +87,10 @@ private def mkACIPrefix (declName : Name) : ProofM Expr := do
   let s ← getStruct
   return mkApp5 (mkConst declName [s.u]) s.type (← getContext) s.assocInst (← getCommInst) (← getNeutralInst)
 
+private def mkACDPrefix (declName : Name) : ProofM Expr := do
+  let s ← getStruct
+  return mkApp5 (mkConst declName [s.u]) s.type (← getContext) s.assocInst (← getCommInst) (← getIdempotentInst)
+
 private def mkDupPrefix (declName : Name) : ProofM Expr := do
   let s ← getStruct
   return mkApp4 (mkConst declName [s.u]) s.type (← getContext) s.assocInst (← getIdempotentInst)
@@ -95,13 +103,18 @@ private def toContextExpr (vars : Array Expr) : ACM Expr := do
 
 private def mkContext (h : Expr) : ProofM Expr := do
   let s ← getStruct
+  let varDecls := (← get).varDecls
   let mut usedVars     :=
+    collectMapVars varDecls collectVar >>
     collectMapVars (← get).seqDecls (·.collectVars) >>
     collectMapVars (← get).exprDecls (·.collectVars) >>
     (if (← hasNeutral) then (collectVar 0) else id) <| {}
   let vars'        := usedVars.toArray
   let varRename    := mkVarRename vars'
   let vars         := (← getStruct).vars
+  let varFVars     := vars'.map fun x => varDecls[x]?.getD default
+  let varIdsAsExpr := List.range vars'.size |>.toArray |>.map toExpr
+  let h := h.replaceFVars varFVars varIdsAsExpr
   let up           := mkApp (mkConst ``PLift.up [s.u]) s.type
   let vars         := vars'.map fun x => mkApp up vars[x]!
   let h := mkLetOfMap (← get).seqDecls h `s (mkConst ``Grind.AC.Seq) fun p => toExpr <| p.renameVars varRename
@@ -173,6 +186,10 @@ partial def EqCnstr.toExprProof (c : EqCnstr) : ProofM Expr := do caching c do
     let h := mkApp9 h (← mkSeqDecl p) (← mkSeqDecl common) (← mkSeqDecl s) (← mkSeqDecl c₁.lhs) (← mkSeqDecl c₁.rhs)
         (← mkSeqDecl c₂.lhs) (← mkSeqDecl c₂.rhs) (← mkSeqDecl c.lhs) (← mkSeqDecl c.rhs)
     return mkApp3 h eagerReflBoolTrue (← c₁.toExprProof) (← c₂.toExprProof)
+  | .superpose_ac_idempotent x c₁ =>
+    let h ← mkACDPrefix ``AC.superpose_ac_idempotent
+    let h := mkApp4 h (← mkVarDecl x) (← mkSeqDecl c₁.lhs) (← mkSeqDecl c₁.rhs) (← mkSeqDecl c.rhs)
+    return mkApp2 h eagerReflBoolTrue (← c₁.toExprProof)
 
 partial def DiseqCnstr.toExprProof (c : DiseqCnstr) : ProofM Expr := do caching c do
   match c.h with

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

@@ -37,6 +37,7 @@ inductive EqCnstrProof where
   | simp_middle (lhs : Bool) (s₁ s₂ : AC.Seq) (c₁ : EqCnstr) (c₂ : EqCnstr)
   | superpose_ac (r₁ c r₂ : AC.Seq) (c₁ : EqCnstr) (c₂ : EqCnstr)
   | superpose (p s c : AC.Seq) (c₁ : EqCnstr) (c₂ : EqCnstr)
+  | superpose_ac_idempotent (x : AC.Var) (c₁ : EqCnstr)
 end
 
 instance : Inhabited EqCnstrProof where

tests/lean/run/grind_ac_2.lean

@@ -42,3 +42,21 @@ example {α : Sort u} (op : α → α → α) [Std.Associative op] [Std.Commutat
     (one : α) [Std.LawfulIdentity op one] (a b c d : α)
     : op a (op (op b one) b) = op d c → op (op b a) (op (op b one) c) = op (op c one) (op d c)  := by
   grind only
+
+example {α : Sort u} (op : α → α → α) [Std.Associative op] [Std.Commutative op] (a b : α)
+    : op a (op a b) = op a a →
+      op a (op b b) = op b b →
+      op b (op b b) = op b b  := by
+  grind only
+
+example {α : Sort u} (op : α → α → α) [Std.Associative op] [Std.Commutative op] [Std.IdempotentOp op] (a b : α)
+    : op a (op a b) = op a a →
+      op a (op b b) = op b b →
+      a = b := by
+  grind only
+
+example {α : Sort u} (op : α → α → α) [Std.Associative op] [Std.Commutative op]
+    [Std.IdempotentOp op]
+    (a b c d : α)
+    : op a (op b b) = op d c → op (op b a) (op b c) = op c (op d c)  := by
+  grind only

tests/lean/run/grind_ac_3.lean

@@ -227,6 +227,11 @@ example (a b c d e f g h : Nat) :
     max (max c d) (max f g) = max (max c d) (max h e) := by
   grind -cutsat only
 
+example (a b c d e f g h : Nat) :
+    max a b = max c d → max b e = max (max d 0) f → max b g = max d h →
+    max (max c d) (max f g) = max (max c d) (max (max 0 h) e) := by
+  grind -cutsat only
+
 example (a b c d e f g h : Nat) :
     max a b = max c d → max b e = max d f → max b g = max d h →
     max (max f d) (max c g) = max (max e d) (max h c) := by
@@ -241,3 +246,10 @@ example {α} (op : α → α → α) [Std.Associative op] [Std.Commutative op] (
     : op a b = op b c → op c c = op d c →
       op (op d a) (op b d) = op (op a a) (op b d) := by
   grind only
+
+example {α} (op : α → α → α) [Std.Associative op] [Std.Commutative op] (u : α) [Std.LawfulIdentity op u]
+    (a b c d e : α)
+    : e = u →
+      op (op a u) b = op (op b e) c → op c c = op d (op u c) →
+      op (op d a) (op b (op d e)) = op (op a a) (op (op u b) d) := by
+  grind only

tests/lean/run/grind_ac_4.lean

@@ -100,6 +100,12 @@ example {α} (op : α → α → α) [Std.Associative op] (a b c d : α)
      op (op c a) (op b c) = op (op a d) (op d b) := by
   grind
 
+example {α} (op : α → α → α) [Std.Associative op] (u : α) [Std.LawfulIdentity op u] (a b c d : α)
+   : op u (op a b) = op u c →
+     op b (op a u) = op d u →
+     op (op c a) (op (op b u) c) = op (op a d) (op d b) := by
+  grind
+
 example {α} (a b c d : List α)
    : a ++ b = c →
      b ++ a = d →
@@ -111,3 +117,9 @@ example {α} (a b c d : List α)
      b ++ a = d ++ d →
      c ++ c ++ a ++ b ++ c ++ c = a ++ d ++ d ++ d ++ d ++ b := by
   grind only
+
+example {α} (a b c d : List α)
+   : a ++ b = c ++ [] ++ c  →
+     b ++ a = d ++ d →
+     c ++ c ++ a ++ [] ++ b ++ c ++ c = a ++ d ++ [] ++ d ++ d ++ d ++ b := by
+  grind only