feat: congruence proofs for grind (#6457)

This PR adds support for generating congruence proofs for congruences
detected by the `grind` tactic.

src/Init/Grind/Util.lean

@@ -11,4 +11,9 @@ namespace Lean.Grind
 /-- A helper gadget for annotating nested proofs in goals. -/
 def nestedProof (p : Prop) (h : p) : p := h
 
+set_option pp.proofs true
+
+theorem nestedProof_congr (p q : Prop) (h : p = q) (hp : p) (hq : q) : HEq (nestedProof p hp) (nestedProof q hq) := by
+  subst h; apply HEq.refl
+
 end Lean.Grind

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

@@ -57,6 +57,7 @@ partial def internalize (e : Expr) (generation : Nat) : GoalM Unit := do
           registerParent e arg
       mkENode e generation
       addCongrTable e
+      propagateUp e
 
 /--
 The fields `target?` and `proof?` in `e`'s `ENode` are encoding a transitivity proof

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

@@ -42,13 +42,16 @@ private def checkParents (e : Expr) : GoalM Unit := do
   if (← isRoot e) then
     for parent in (← getParents e) do
       let mut found := false
+      let checkChild (child : Expr) : GoalM Bool := do
+        let some childRoot ← getRoot? child | return false
+        return isSameExpr childRoot e
       -- There is an argument `arg` s.t. root of `arg` is `e`.
       for arg in parent.getAppArgs do
-        if let some argRoot ← getRoot? arg then
-          if isSameExpr argRoot e then
-            found := true
-            break
-      assert! found
+        if (← checkChild arg) then
+          found := true
+          break
+      unless found do
+        assert! (← checkChild parent.getAppFn)
   else
     -- All the parents are stored in the root of the equivalence class.
     assert! (← getParents e).isEmpty

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

@@ -9,16 +9,6 @@ import Lean.Meta.Tactic.Grind.Types
 
 namespace Lean.Meta.Grind
 
--- TODO: delete after done
-private def mkTodo (a b : Expr) (heq : Bool) : MetaM Expr := do
-  if heq then
-    mkSorry (← mkHEq a b) (synthetic := false)
-  else
-    mkSorry (← mkEq a b) (synthetic := false)
-
-private def isProtoProof (h : Expr) : Bool :=
-  isSameExpr h congrPlaceholderProof
-
 private def isEqProof (h : Expr) : MetaM Bool := do
   return (← whnfD (← inferType h)).isAppOf ``Eq
 
@@ -44,6 +34,13 @@ private def mkTrans' (h₁ : Option Expr) (h₂ : Expr) (heq : Bool) : MetaM Exp
   let some h₁ := h₁ | return h₂
   mkTrans h₁ h₂ heq
 
+/--
+Given `h : HEq a b`, returns a proof `a = b` if `heq == false`.
+Otherwise, it returns `h`.
+-/
+private def mkEqOfHEqIfNeeded (h : Expr) (heq : Bool) : MetaM Expr := do
+  if heq then return h else mkEqOfHEq h
+
 /--
 Given `lhs` and `rhs` that are in the same equivalence class,
 find the common expression that are in the paths from `lhs` and `rhs` to
@@ -71,9 +68,49 @@ private def findCommon (lhs rhs : Expr) : GoalM Expr := do
   unreachable!
 
 mutual
+  private partial def mkNestedProofCongr (lhs rhs : Expr) (heq : Bool) : GoalM Expr := do
+    let p  := lhs.appFn!.appArg!
+    let hp := lhs.appArg!
+    let q  := rhs.appFn!.appArg!
+    let hq := rhs.appArg!
+    let h  := mkApp5 (mkConst ``Lean.Grind.nestedProof_congr) p q (← mkEqProofCore p q false) hp hq
+    mkEqOfHEqIfNeeded h heq
+
   private partial def mkCongrProof (lhs rhs : Expr) (heq : Bool) : GoalM Expr := do
-    -- TODO: implement
-    mkTodo lhs rhs heq
+    let f := lhs.getAppFn
+    let g := rhs.getAppFn
+    let numArgs := lhs.getAppNumArgs
+    assert! rhs.getAppNumArgs == numArgs
+    if f.isConstOf ``Lean.Grind.nestedProof && g.isConstOf ``Lean.Grind.nestedProof && numArgs == 2 then
+      mkNestedProofCongr lhs rhs heq
+    else
+      let thm ← mkHCongrWithArity f numArgs
+      assert! thm.argKinds.size == numArgs
+      let rec loop (lhs rhs : Expr) (i : Nat) : GoalM Expr := do
+        let i := i - 1
+        if lhs.isApp then
+          let proof ← loop lhs.appFn! rhs.appFn! i
+          let a₁  := lhs.appArg!
+          let a₂  := rhs.appArg!
+          let k   := thm.argKinds[i]!
+          return mkApp3 proof a₁ a₂ (← mkEqProofCore a₁ a₂ (k matches .heq))
+        else
+          return thm.proof
+      let proof ← loop lhs rhs numArgs
+      if isSameExpr f g then
+        mkEqOfHEqIfNeeded proof heq
+      else
+        /-
+        `lhs` is of the form `f a_1 ... a_n`
+        `rhs` is of the form `g b_1 ... b_n`
+        `proof : HEq (f a_1 ... a_n) (f b_1 ... b_n)`
+        We construct a proof for `HEq (f a_1 ... a_n) (g b_1 ... b_n)` using `Eq.ndrec`
+        -/
+        let motive ← withLocalDeclD (← mkFreshUserName `x) (← inferType f) fun x => do
+          mkLambdaFVars #[x] (← mkHEq lhs (mkAppN x rhs.getAppArgs))
+        let fEq ← mkEqProofCore f g false
+        let proof ← mkEqNDRec motive proof fEq
+        mkEqOfHEqIfNeeded proof heq
 
   private partial def realizeEqProof (lhs rhs : Expr) (h : Expr) (flipped : Bool) (heq : Bool) : GoalM Expr := do
     let h ← if h == congrPlaceholderProof then
@@ -140,6 +177,9 @@ where
         trace[grind.proof] "{a} ≡ {b}"
         mkEqProofCore a b (heq := true)
 
+def mkHEqProof (a b : Expr) : GoalM Expr :=
+  mkEqProofCore a b (heq := true)
+
 /--
 Returns a proof that `a = True`.
 It assumes `a` and `True` are in the same equivalence class.

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

@@ -120,7 +120,7 @@ builtin_grind_propagator propagateEqUp ↑Eq := fun e => do
   else if (← isEqTrue b) then
     pushEq e a <| mkApp3 (mkConst ``Lean.Grind.eq_eq_of_eq_true_right) a b (← mkEqTrueProof b)
   else if (← isEqv a b) then
-    pushEqTrue e <| mkApp2 (mkConst ``of_eq_true) e (← mkEqProof a b)
+    pushEqTrue e <| mkApp2 (mkConst ``eq_true) e (← mkEqProof a b)
 
 /-- Propagates `Eq` downwards -/
 builtin_grind_propagator propagateEqDown ↓Eq := fun e => do
@@ -134,4 +134,10 @@ builtin_grind_propagator propagateHEqDown ↓HEq := fun e => do
     let_expr HEq _ a _ b := e | return ()
     pushHEq a b <| mkApp2 (mkConst ``of_eq_true) e (← mkEqTrueProof e)
 
+/-- Propagates `HEq` upwards -/
+builtin_grind_propagator propagateHEqUp ↑HEq := fun e => do
+  let_expr HEq _ a _ b := e | return ()
+  if (← isEqv a b) then
+    pushEqTrue e <| mkApp2 (mkConst ``eq_true) e (← mkHEqProof a b)
+
 end Lean.Meta.Grind

tests/lean/run/grind_congr1.lean

@@ -0,0 +1,91 @@
+import Lean
+
+open Lean Meta Elab Tactic Grind in
+elab "grind_test" : tactic => withMainContext do
+  let declName := (← Term.getDeclName?).getD `_main
+  Meta.Grind.preprocessAndProbe (← getMainGoal) declName do
+    unless (← isInconsistent) do
+      throwError "inconsistent state expected"
+
+set_option grind.debug true
+set_option grind.debug.proofs true
+
+example (a b : Nat) (f : Nat → Nat) : (h₁ : a = b) → f a = f b := by
+  grind_test
+  sorry
+
+example (a b : Nat) (f : Nat → Nat) : (h₁ : a = b) → (h₂ : f a ≠ f b) → False := by
+  grind_test
+  sorry
+
+example (a b : Nat) (f : Nat → Nat) : a = b → f (f a) ≠ f (f b) → False := by
+  grind_test
+  sorry
+
+example (a b c : Nat) (f : Nat → Nat) : a = b → c = b → f (f a) ≠ f (f c) → False := by
+  grind_test
+  sorry
+
+example (a b c : Nat) (f : Nat → Nat → Nat) : a = b → c = b → f (f a b) a ≠ f (f c c) c → False := by
+  grind_test
+  sorry
+
+example (a b c : Nat) (f : Nat → Nat → Nat) : a = b → c = b → f (f a b) a = f (f c c) c := by
+  grind_test
+  sorry
+
+example (a b c d : Nat) : a = b → b = c → c = d → a = d := by
+  grind_test
+  sorry
+
+infix:50 "===" => HEq
+
+example (a b c d : Nat) : a === b → b = c → c === d → a === d := by
+  grind_test
+  sorry
+
+example (a b c d : Nat) : a = b → b = c → c === d → a === d := by
+  grind_test
+  sorry
+
+opaque f {α : Type} : α → α → α := fun a _ => a
+opaque g : Nat → Nat
+
+example (a b c : Nat) : a = b → g a === g b := by
+  grind_test
+  sorry
+
+example (a b c : Nat) : a = b → c = b → f (f a b) (g c) = f (f c a) (g b) := by
+  grind_test
+  sorry
+
+example (a b c d e x y : Nat) : a = b → a = x → b = y → c = d → c = e → c = b → a = e := by
+  grind_test
+  sorry
+
+namespace Ex1
+opaque f (a b : Nat) : a > b → Nat
+opaque g : Nat → Nat
+
+example (a₁ a₂ b₁ b₂ c d : Nat)
+        (H₁ : a₁ > b₁)
+        (H₂ : a₂ > b₂) :
+        a₁ = c → a₂ = c →
+        b₁ = d → d  = b₂ →
+        g (g (f a₁ b₁ H₁)) = g (g (f a₂ b₂ H₂)) := by
+  grind_test
+  sorry
+end Ex1
+
+namespace Ex2
+def f (α : Type) (a : α) : α := a
+
+example (a : α) (b : β) : (h₁ : α = β) → (h₂ : HEq a b) → HEq (f α a) (f β b) := by
+  grind_test
+  sorry
+
+end Ex2
+
+example (f g : (α : Type) → α → α) (a : α) (b : β) : (h₁ : α = β) → (h₂ : HEq a b) → (h₃ : f = g) → HEq (f α a) (g β b) := by
+  grind_test
+  sorry

tests/lean/run/grind_pre.lean

@@ -118,8 +118,12 @@ example (p : Prop) (a b c : Nat) : p → a = 0 → a = b → h a = h c → a = c
   grind_pre
   sorry
 
-set_option trace.grind.proof.detail true in
-set_option trace.grind.proof true in
+set_option trace.grind.proof.detail true
+set_option trace.grind.proof true
 example (a : α) (p q r : Prop) : (h₁ : HEq p a) → (h₂ : HEq q a) → (h₃ : p = r) → False := by
   grind_pre
   sorry
+
+example (a b : Nat) (f : Nat → Nat) : (h₁ : a = b) → (h₂ : f a ≠ f b) → False := by
+  grind_pre
+  sorry