feat: match-expressions with congruence equation theorems (leanprover#8506)

This PR implements `match`-expressions in `grind` using `match`
congruence equations. The goal is to minimize the number of `cast`
operations that need to be inserted, and avoid `cast` over functions.
The new approach support `match`-expressions of the form `match h : ...
with ...`.

Author: leodemoura

src/Init/Grind/Util.lean

@@ -32,6 +32,55 @@ def offset (a b : Nat) : Nat := a + b
 /-- Gadget for representing `a = b` in patterns for backward propagation. -/
 def eqBwdPattern (a b : α) : Prop := a = b
 
+/--
+Gadget for representing generalization steps `h : x = val` in patterns
+This gadget is used to represent patterns in theorems that have been generalized to reduce the
+number of casts introduced during E-matching based instantiation.
+
+For example, consider the theorem
+```
+Option.pbind_some {α1 : Type u_1} {a : α1} {α2 : Type u_2}
+    {f : (a_1 : α1) → some a = some a_1 → Option α2}
+    : (some a).pbind f = f a rfl
+```
+Now, suppose we have a goal containing the term `c.pbind g` and the equivalence class
+`{c, some b}`. The E-matching module generates the instance
+```
+(some b).pbind (cast ⋯ g)
+```
+The `cast` is necessary because `g`'s type contains `c` instead of `some b.
+This `cast` problematic because we don't have a systematic way of pushing casts over functions
+to its arguments. Moreover, heterogeneous equality is not effective because the following theorem
+is not provable in DTT:
+```
+theorem hcongr (h₁ : f ≍ g) (h₂ : a ≍ b)  : f a ≍ g b := ...
+```
+The standard solution is to generalize the theorem above and write it as
+```
+theorem Option.pbind_some'
+        {α1 : Type u_1} {a : α1} {α2 : Type u_2}
+        {x : Option α1}
+        {f : (a_1 : α1) → x = some a_1 → Option α2}
+        (h : x = some a)
+        : x.pbind f = f a h := by
+  subst h
+  apply Option.pbind_some
+```
+Internally, we use this gadget to mark the E-matching pattern as
+```
+(genPattern h x (some a)).pbind f
+```
+This pattern is matched in the same way we match `(some a).pbind f`, but it saves the proof
+for the actual term to the `some`-application in `f`, and the actual term in `x`.
+
+In the example above, `c.pbind g` also matches the pattern `(genPattern h x (some a)).pbind f`,
+and stores `c` in `x`, `b` in `a`, and the proof that `c = some b` in `h`.
+-/
+def genPattern {α : Sort u} (_h : Prop) (x : α) (_val : α) : α := x
+
+/-- Similar to `genPattern` but for the heterogenous case -/
+def genHEqPattern {α β : Sort u} (_h : Prop) (x : α) (_val : β) : α := x
+
 /--
 Gadget for annotating the equalities in `match`-equations conclusions.
 `_origin` is the term used to instantiate the `match`-equation using E-matching.

src/Lean/Elab/Tactic/Grind.lean

@@ -214,27 +214,26 @@ def mkGrindOnly
   let mut foundFns : NameSet := {}
   for { origin, kind } in trace.thms.toList do
     if let .decl declName := origin then
-      unless Match.isMatchEqnTheorem (← getEnv) declName do
-        if let some declName ← isEqnThm? declName then
-          unless foundFns.contains declName do
-            foundFns := foundFns.insert declName
-            let decl : Ident := mkIdent (← unresolveNameGlobalAvoidingLocals declName)
-            let param ← `(Parser.Tactic.grindParam| $decl:ident)
-            params := params.push param
-        else
+      if let some declName ← isEqnThm? declName then
+        unless foundFns.contains declName do
+          foundFns := foundFns.insert declName
           let decl : Ident := mkIdent (← unresolveNameGlobalAvoidingLocals declName)
-          let param ← match kind with
-            | .eqLhs     => `(Parser.Tactic.grindParam| = $decl)
-            | .eqRhs     => `(Parser.Tactic.grindParam| =_ $decl)
-            | .eqBoth    => `(Parser.Tactic.grindParam| _=_ $decl)
-            | .eqBwd     => `(Parser.Tactic.grindParam| ←= $decl)
-            | .bwd       => `(Parser.Tactic.grindParam| ← $decl)
-            | .fwd       => `(Parser.Tactic.grindParam| → $decl)
-            | .leftRight => `(Parser.Tactic.grindParam| => $decl)
-            | .rightLeft => `(Parser.Tactic.grindParam| <= $decl)
-            | .user      => `(Parser.Tactic.grindParam| usr $decl)
-            | .default   => `(Parser.Tactic.grindParam| $decl:ident)
+          let param ← `(Parser.Tactic.grindParam| $decl:ident)
           params := params.push param
+      else
+        let decl : Ident := mkIdent (← unresolveNameGlobalAvoidingLocals declName)
+        let param ← match kind with
+          | .eqLhs     => `(Parser.Tactic.grindParam| = $decl)
+          | .eqRhs     => `(Parser.Tactic.grindParam| =_ $decl)
+          | .eqBoth    => `(Parser.Tactic.grindParam| _=_ $decl)
+          | .eqBwd     => `(Parser.Tactic.grindParam| ←= $decl)
+          | .bwd       => `(Parser.Tactic.grindParam| ← $decl)
+          | .fwd       => `(Parser.Tactic.grindParam| → $decl)
+          | .leftRight => `(Parser.Tactic.grindParam| => $decl)
+          | .rightLeft => `(Parser.Tactic.grindParam| <= $decl)
+          | .user      => `(Parser.Tactic.grindParam| usr $decl)
+          | .default   => `(Parser.Tactic.grindParam| $decl:ident)
+        params := params.push param
   for declName in trace.eagerCases.toList do
     unless Grind.isBuiltinEagerCases declName do
       let decl : Ident := mkIdent (← unresolveNameGlobalAvoidingLocals declName)

src/Lean/Meta/Match/MatchEqs.lean

@@ -833,7 +833,7 @@ def congrEqn1ThmSuffix := congrEqnThmSuffixBasePrefix ++ "1"
 example : congrEqn1ThmSuffix = "congr_eq_1" := rfl
 
 /-- Returns `true` if `s` is of the form `congr_eq_<idx>` -/
-def iscongrEqnReservedNameSuffix (s : String) : Bool :=
+def isCongrEqnReservedNameSuffix (s : String) : Bool :=
   congrEqnThmSuffixBasePrefix.isPrefixOf s && (s.drop congrEqnThmSuffixBasePrefix.length).isNat
 
 /- We generate the equations and splitter on demand, and do not save them on .olean files. -/
@@ -928,10 +928,10 @@ builtin_initialize registerTraceClass `Meta.Match.matchEqs
 
 private def isMatchEqName? (env : Environment) (n : Name) : Option (Name × Bool) := do
   let .str p s := n | failure
-  guard <| isEqnReservedNameSuffix s || s == "splitter" || iscongrEqnReservedNameSuffix s
+  guard <| isEqnReservedNameSuffix s || s == "splitter" || isCongrEqnReservedNameSuffix s
   let p ← privateToUserName? p
   guard <| isMatcherCore env p
-  return (p, iscongrEqnReservedNameSuffix s)
+  return (p, isCongrEqnReservedNameSuffix s)
 
 builtin_initialize registerReservedNamePredicate (isMatchEqName? · · |>.isSome)
 

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

@@ -17,9 +17,9 @@ namespace EMatch
 /-- We represent an `E-matching` problem as a list of constraints. -/
 inductive Cnstr where
   | /-- Matches pattern `pat` with term `e` -/
-    «match» (pat : Expr) (e : Expr)
+    «match» (gen? : Option GenPatternInfo) (pat : Expr) (e : Expr)
   | /-- Matches offset pattern `pat+k` with term `e` -/
-    offset (pat : Expr) (k : Nat) (e : Expr)
+    offset (gen? : Option GenPatternInfo) (pat : Expr) (k : Nat) (e : Expr)
   | /-- This constraint is used to encode multi-patterns. -/
     «continue» (pat : Expr)
   deriving Inhabited
@@ -30,6 +30,12 @@ This is a small hack to avoid one extra level of indirection by using `Option Ex
 -/
 private def unassigned : Expr := mkConst (Name.mkSimple "[grind_unassigned]")
 
+/--
+Internal "marker" for representing equality proofs for generalized patterns.
+They must be synthesized after we have the partial assignment.
+-/
+private def delayedEqProof : Expr := mkConst (Name.mkSimple "[grind_delayed_eq_proof]")
+
 private def assignmentToMessageData (assignment : Array Expr) : Array MessageData :=
   assignment.reverse.map fun e =>
     if isSameExpr e unassigned then m!"_" else m!"{e}"
@@ -89,6 +95,22 @@ private def assign? (c : Choice) (bidx : Nat) (e : Expr) : OptionT GoalM Choice
     -- `Choice` was not properly initialized
     unreachable!
 
+/--
+Assigns `bidx` with the marker for a delayed equality proof for generalized patterns.
+The proof is assigned after we have the complete assignment.
+-/
+private def assignDelayedEqProof? (c : Choice) (bidx : Nat) : OptionT GoalM Choice := do
+  if h : bidx < c.assignment.size then
+    let v := c.assignment[bidx]
+    if isSameExpr v unassigned then
+      return { c with assignment := c.assignment.set bidx delayedEqProof }
+    else
+      return c
+  else
+    -- `Choice` was not properly initialized
+    unreachable!
+
+
 private def unassign (c : Choice) (bidx : Nat) : Choice :=
   { c with assignment := c.assignment.set! bidx unassigned }
 
@@ -100,6 +122,11 @@ private def eqvFunctions (pFn eFn : Expr) : Bool :=
   (pFn.isFVar && pFn == eFn)
   || (pFn.isConst && eFn.isConstOf pFn.constName!)
 
+protected def _root_.Lean.Meta.Grind.GenPatternInfo.assign? (genInfo : GenPatternInfo) (c : Choice) (x : Expr) : OptionT GoalM Choice := do
+  let c ← assign? c genInfo.xIdx x
+  let c ← assignDelayedEqProof? c genInfo.hIdx
+  return c
+
 /-- Matches a pattern argument. See `matchArgs?`. -/
 private def matchArg? (c : Choice) (pArg : Expr) (eArg : Expr) : OptionT GoalM Choice := do
   if isPatternDontCare pArg then
@@ -128,9 +155,20 @@ private def matchArg? (c : Choice) (pArg : Expr) (eArg : Expr) : OptionT GoalM C
   else if let some (pArg, k) := isOffsetPattern? pArg then
     assert! Option.isNone <| isOffsetPattern? pArg
     assert! !isPatternDontCare pArg
-    return { c with cnstrs := .offset pArg k eArg :: c.cnstrs }
+    return { c with cnstrs := .offset none pArg k eArg :: c.cnstrs }
+  else if let some (genInfo, pArg) := isGenPattern? pArg then
+    if pArg.isBVar then
+      let c ← assign? c pArg.bvarIdx! eArg
+      genInfo.assign? c eArg
+    else if let some pArg := groundPattern? pArg then
+      guard (← isEqv pArg eArg <||> withReducibleAndInstances (isDefEq pArg eArg))
+      genInfo.assign? c eArg
+    else if let some (pArg, k) := isOffsetPattern? pArg then
+      return { c with cnstrs := .offset (some genInfo) pArg k eArg :: c.cnstrs }
+    else
+      return { c with cnstrs := .match (some genInfo) pArg eArg :: c.cnstrs }
   else
-    return { c with cnstrs := .match pArg eArg :: c.cnstrs }
+    return { c with cnstrs := .match none pArg eArg :: c.cnstrs }
 
 private def Choice.updateGen (c : Choice) (gen : Nat) : Choice :=
   { c with gen := Nat.max gen c.gen }
@@ -160,25 +198,35 @@ private partial def matchArgsPrefix? (c : Choice) (p : Expr) (e : Expr) : Option
   else
     matchArgs? c p (e.getAppPrefix pn)
 
+-- Helper function for `processMatch` and `processGenMatch`
+@[inline] private def isCandidate (n : ENode) (pFn : Expr) (pNumArgs : Nat) (maxGeneration : Nat) : Bool :=
+    -- Remark: we use `<` because the instance generation is the maximum term generation + 1
+    n.generation < maxGeneration
+    -- uses heterogeneous equality or is the root of its congruence class
+    && (n.heqProofs || n.isCongrRoot)
+    && eqvFunctions pFn n.self.getAppFn
+    && n.self.getAppNumArgs == pNumArgs
+
+private def assignGenInfo? (genInfo? : Option GenPatternInfo) (c : Choice) (x : Expr) : OptionT GoalM Choice := do
+  let some genInfo := genInfo? | return c
+  genInfo.assign? c x
+
 /--
 Matches pattern `p` with term `e` with respect to choice `c`.
 We traverse the equivalence class of `e` looking for applications compatible with `p`.
 For each candidate application, we match the arguments and may update `c`s assignments and constraints.
 We add the updated choices to the choice stack.
 -/
-private partial def processMatch (c : Choice) (p : Expr) (e : Expr) : M Unit := do
+private partial def processMatch (c : Choice) (genInfo? : Option GenPatternInfo) (p : Expr) (e : Expr) : M Unit := do
   let maxGeneration ← getMaxGeneration
   let pFn := p.getAppFn
   let numArgs := p.getAppNumArgs
   let mut curr := e
   repeat
     let n ← getENode curr
     -- Remark: we use `<` because the instance generation is the maximum term generation + 1
-    if n.generation < maxGeneration
-       -- uses heterogeneous equality or is the root of its congruence class
-       && (n.heqProofs || n.isCongrRoot)
-       && eqvFunctions pFn curr.getAppFn
-       && curr.getAppNumArgs == numArgs then
+    if isCandidate n pFn numArgs maxGeneration then
+      if let some c ← assignGenInfo? genInfo? c e |>.run then
       if let some c ← matchArgs? c p curr |>.run then
         pushChoice (c.updateGen n.generation)
     curr ← getNext curr
@@ -187,32 +235,34 @@ private partial def processMatch (c : Choice) (p : Expr) (e : Expr) : M Unit :=
 /--
 Matches offset pattern `pArg+k` with term `e` with respect to choice `c`.
 -/
-private partial def processOffset (c : Choice) (pArg : Expr) (k : Nat) (e : Expr) : M Unit := do
+private partial def processOffset (c : Choice) (genInfo? : Option GenPatternInfo) (pArg : Expr) (k : Nat) (e : Expr) : M Unit := do
   let maxGeneration ← getMaxGeneration
   let mut curr := e
   repeat
     let n ← getENode curr
     if n.generation < maxGeneration then
       if let some (eArg, k') ← isOffset? curr |>.run then
-        if k' < k then
-          let c := c.updateGen n.generation
-          pushChoice { c with cnstrs := .offset pArg (k - k') eArg :: c.cnstrs }
-        else if k' == k then
-          if let some c ← matchArg? c pArg eArg |>.run then
-            pushChoice (c.updateGen n.generation)
-        else if k' > k then
-          let eArg' := mkNatAdd eArg (mkNatLit (k' - k))
-          let eArg' ← shareCommon (← canon eArg')
-          internalize eArg' (n.generation+1)
-          if let some c ← matchArg? c pArg eArg' |>.run then
-            pushChoice (c.updateGen n.generation)
+        if let some c ← assignGenInfo? genInfo? c e |>.run then
+          if k' < k then
+            let c := c.updateGen n.generation
+            pushChoice { c with cnstrs := .offset none pArg (k - k') eArg :: c.cnstrs }
+          else if k' == k then
+            if let some c ← matchArg? c pArg eArg |>.run then
+              pushChoice (c.updateGen n.generation)
+          else if k' > k then
+            let eArg' := mkNatAdd eArg (mkNatLit (k' - k))
+            let eArg' ← shareCommon (← canon eArg')
+            internalize eArg' (n.generation+1)
+            if let some c ← matchArg? c pArg eArg' |>.run then
+              pushChoice (c.updateGen n.generation)
       else if let some k' ← evalNat curr |>.run then
-        if k' >= k then
-          let eArg' := mkNatLit (k' - k)
-          let eArg' ← shareCommon (← canon eArg')
-          internalize eArg' (n.generation+1)
-          if let some c ← matchArg? c pArg eArg' |>.run then
-            pushChoice (c.updateGen n.generation)
+        if let some c ← assignGenInfo? genInfo? c e |>.run then
+          if k' >= k then
+            let eArg' := mkNatLit (k' - k)
+            let eArg' ← shareCommon (← canon eArg')
+            internalize eArg' (n.generation+1)
+            if let some c ← matchArg? c pArg eArg' |>.run then
+              pushChoice (c.updateGen n.generation)
     curr ← getNext curr
     if isSameExpr curr e then break
 
@@ -280,7 +330,7 @@ private def addNewInstance (thm : EMatchTheorem) (proof : Expr) (generation : Na
     check proof
   let mut prop ← inferType proof
   let mut proof := proof
-  if Match.isMatchEqnTheorem (← getEnv) thm.origin.key then
+  if (← isMatchEqLikeDeclName thm.origin.key) then
     prop ← annotateMatchEqnType prop (← read).initApp
     -- We must add a hint here because `annotateMatchEqnType` introduces `simpMatchDiscrsOnly` and
     -- `Grind.PreMatchCond` which are not reducible.
@@ -302,13 +352,43 @@ private def synthesizeInsts (mvars : Array Expr) (bis : Array BinderInfo) : Opti
         reportIssue! "failed to synthesize instance when instantiating {← thm.origin.pp}{indentExpr type}"
         failure
 
+private def preprocessGeneralizedPatternRHS (lhs : Expr) (rhs : Expr) (origin : Origin) (expectedType : Expr) : OptionT (StateT Choice M) Expr := do
+  assert! (← alreadyInternalized lhs)
+  -- We use `dsimp` here to ensure terms such as `Nat.succ x` are normalized as `x+1`.
+  let rhs ← preprocessLight (← dsimpCore rhs)
+  internalize rhs (← getGeneration lhs)
+  processNewFacts
+  if (← isEqv lhs rhs) then
+    return rhs
+  else
+    reportIssue! "invalid generalized pattern at `{← origin.pp}`\nwhen processing argument with type{indentExpr expectedType}\nfailed to prove{indentExpr lhs}\nis equal to{indentExpr rhs}"
+    failure
+
+private def assignGeneralizedPatternProof (mvarId : MVarId) (eqProof : Expr) (origin : Origin) : OptionT (StateT Choice M) Unit := do
+  unless (← mvarId.checkedAssign eqProof) do
+    reportIssue! "invalid generalized pattern at `{← origin.pp}`\nfailed to assign {mkMVar mvarId}\nwith{indentExpr eqProof}"
+    failure
+
 private def applyAssignment (mvars : Array Expr) : OptionT (StateT Choice M) Unit := do
   let thm := (← read).thm
   let numParams := thm.numParams
   for h : i in [:mvars.size] do
     let bidx := numParams - i - 1
     let mut v := (← get).assignment[bidx]!
-    unless isSameExpr v unassigned do
+    if isSameExpr v delayedEqProof then
+      let mvarId := mvars[i].mvarId!
+      let mvarIdType ← instantiateMVars (← mvarId.getType)
+      match_expr mvarIdType with
+      | Eq α lhs rhs =>
+        let rhs ← preprocessGeneralizedPatternRHS lhs rhs thm.origin mvarIdType
+        assignGeneralizedPatternProof mvarId (← mkEqProof lhs rhs) thm.origin
+      | HEq α lhs β rhs =>
+        let rhs ← preprocessGeneralizedPatternRHS lhs rhs thm.origin mvarIdType
+        assignGeneralizedPatternProof mvarId (← mkHEqProof lhs rhs) thm.origin
+      | _ =>
+        reportIssue! "invalid generalized pattern at `{← thm.origin.pp}`\nequality type expected{indentExpr mvarIdType}"
+        failure
+    else if !isSameExpr v unassigned then
       let mvarId := mvars[i].mvarId!
       let mvarIdType ← mvarId.getType
       let vType ← inferType v
@@ -388,8 +468,8 @@ private def processChoices : M Unit := do
     if c.gen < maxGeneration then
       match c.cnstrs with
       | [] => instantiateTheorem c
-      | .match p e :: cnstrs => processMatch { c with cnstrs } p e
-      | .offset p k e :: cnstrs => processOffset { c with cnstrs } p k e
+      | .match genInfo? p e :: cnstrs => processMatch { c with cnstrs } genInfo? p e
+      | .offset genInfo? p k e :: cnstrs => processOffset { c with cnstrs } genInfo? p k e
       | .continue p :: cnstrs => processContinue { c with cnstrs } p
 
 private def main (p : Expr) (cnstrs : List Cnstr) : M Unit := do