feat: process congr lemmas at simp

src/Lean/Elab/Tactic/Simp.lean

@@ -12,11 +12,11 @@ import Lean.Meta.Tactic.Replace
 namespace Lean.Elab.Tactic
 open Meta
 
-def simpTarget (config : Meta.Simp.Config) (simpLemmas : SimpLemmas) : TacticM Unit := do
+def simpTarget (config : Meta.Simp.Config) (simpLemmas : SimpLemmas) (congrLemmas : CongrLemmas) : TacticM Unit := do
   let (g, gs) ← getMainGoal
   withMVarContext g do
     let target ← instantiateMVars (← getMVarDecl g).type
-    let r ← simp target config simpLemmas
+    let r ← simp target config simpLemmas congrLemmas
     match r.proof? with
     | some proof => setGoals ((← replaceTargetEq g r.expr proof) :: gs)
     | none => setGoals ((← replaceTargetDefEq g r.expr) :: gs)
@@ -25,28 +25,28 @@ def simpTarget (config : Meta.Simp.Config) (simpLemmas : SimpLemmas) : TacticM U
 -- TODO: issues: self simplification
 -- TODO: add new assertion with simplified result and clear old ones after simplifying all locals
 
-def simpLocalDeclFVarId (config : Meta.Simp.Config) (simpLemmas : SimpLemmas) (fvarId : FVarId) : TacticM Unit := do
+def simpLocalDeclFVarId (config : Meta.Simp.Config) (simpLemmas : SimpLemmas) (congrLemmas : CongrLemmas) (fvarId : FVarId) : TacticM Unit := do
   let (g, gs) ← getMainGoal
   withMVarContext g do
     let localDecl ← getLocalDecl fvarId
-    let r ← simp localDecl.type config simpLemmas
+    let r ← simp localDecl.type config simpLemmas congrLemmas
     match r.proof? with
     | some proof =>
       setGoals ((← replaceLocalDecl g fvarId r.expr proof).mvarId :: gs)
     | none => setGoals ((← changeLocalDecl g fvarId r.expr (checkDefEq := false)) :: gs)
 
-def simpLocalDecl (config : Meta.Simp.Config) (simpLemmas : SimpLemmas) (userName : Name) : TacticM Unit :=
+def simpLocalDecl (config : Meta.Simp.Config) (simpLemmas : SimpLemmas) (congrLemmas : CongrLemmas) (userName : Name) : TacticM Unit :=
   withMainMVarContext do
     let localDecl ← getLocalDeclFromUserName userName
-    simpLocalDeclFVarId config simpLemmas localDecl.fvarId
+    simpLocalDeclFVarId config simpLemmas congrLemmas localDecl.fvarId
 
-def simpAll (config : Meta.Simp.Config) (simpLemmas : SimpLemmas) : TacticM Unit := do
-  let worked ← «try» (simpTarget config simpLemmas)
+def simpAll (config : Meta.Simp.Config) (simpLemmas : SimpLemmas) (congrLemmas : CongrLemmas) : TacticM Unit := do
+  let worked ← «try» (simpTarget config simpLemmas congrLemmas)
   withMainMVarContext do
     let mut worked := worked
     -- We must traverse backwards because `replaceLocalDecl` uses the revert/intro idiom
     for fvarId in (← getLCtx).getFVarIds.reverse do
-      worked := worked || (← «try» <| simpLocalDeclFVarId config simpLemmas fvarId)
+      worked := worked || (← «try» <| simpLocalDeclFVarId config simpLemmas congrLemmas fvarId)
     unless worked do
       let (mvarId, _) ← getMainGoal
       throwTacticEx `simp mvarId "failed to simplify"
@@ -74,13 +74,14 @@ def elabSimpConfig (optConfig : Syntax) : TermElabM Meta.Simp.Config := do
       evalSimpConfig (← instantiateMVars c)
 
 @[builtinTactic Lean.Parser.Tactic.simp] def evalSimp : Tactic := fun stx => do
-  let lemmas ← mkSimpLemmas stx[1]
+  let simpLemmas ← mkSimpLemmas stx[1]
+  let congrLemmas ← getCongrLemmas
   let config ← elabSimpConfig stx[2]
   let loc := expandOptLocation stx[3]
   match loc with
-  | Location.target => simpTarget config lemmas
-  | Location.localDecls userNames => userNames.forM (simpLocalDecl config lemmas)
-  | Location.wildcard => simpAll config lemmas
+  | Location.target => simpTarget config simpLemmas congrLemmas
+  | Location.localDecls userNames => userNames.forM (simpLocalDecl config simpLemmas congrLemmas)
+  | Location.wildcard => simpAll config simpLemmas congrLemmas
   tryExactTrivial
 where
   mkSimpLemmas (stx : Syntax) := do

src/Lean/Meta/Tactic/Simp/CongrLemmas.lean

@@ -20,6 +20,11 @@ structure CongrLemmas where
   lemmas : SMap Name (List CongrLemma) := {}
   deriving Inhabited, Repr
 
+def CongrLemmas.get (d : CongrLemmas) (declName : Name) : List CongrLemma :=
+  match d.lemmas.find? declName with
+  | none    => []
+  | some cs => cs
+
 def addCongrLemmaEntry (d : CongrLemmas) (e : CongrLemma) : CongrLemmas :=
   { d with lemmas :=
       match d.lemmas.find? e.funName with
@@ -42,7 +47,6 @@ def mkCongrLemma (declName : Name) (prio : Nat) : MetaM CongrLemma := withReduci
   let info ← getConstInfo declName
   let c := mkConst declName (info.levelParams.map mkLevelParam)
   let (xs, bis, type) ← forallMetaTelescopeReducing (← inferType c)
-  let type ← whnf type
   match type.eq? with
   | none => throwError! "invalid 'congr' lemma, equality expected{indentExpr type}"
   | some (_, lhs, rhs) =>
@@ -60,7 +64,6 @@ def mkCongrLemma (declName : Name) (prio : Nat) : MetaM CongrLemma := withReduci
       for x in xs, bi in bis do
         if bi.isExplicit && !foundMVars.contains x.mvarId! then
           let rhsFn? ← forallTelescopeReducing (← inferType x) fun ys xType => do
-            let xType ← whnf xType
             match xType.eq? with
             | none => pure none -- skip
             | some (_, xLhs, xRhs) =>

src/Lean/Meta/Tactic/Simp/Main.lean

@@ -9,6 +9,12 @@ import Lean.Meta.Tactic.Simp.Rewrite
 namespace Lean.Meta
 namespace Simp
 
+builtin_initialize congrHypothesisExceptionId : InternalExceptionId ←
+  registerInternalExceptionId `congrHypothesisFailed
+
+def throwCongrHypothesisFailed {α} : MetaM α :=
+  throw <| Exception.internal congrHypothesisExceptionId
+
 def Result.getProof (r : Result) : MetaM Expr := do
   match r.proof? with
   | some p => return p
@@ -127,25 +133,82 @@ where
     | Expr.mvar ..     => pure { expr := (← instantiateMVars e) }
     | Expr.fvar ..     => pure { expr := (← reduceFVar (← getConfig) e) }
 
+  congrDefault (e : Expr) : M Result :=
+    withParent e <| e.withApp fun f args => do
+      let infos := (← getFunInfoNArgs f args.size).paramInfo
+      let mut r ← simp f
+      let mut i := 0
+      for arg in args do
+        trace[Meta.Tactic.simp]! "app [{i}] {infos.size} {arg} hasFwdDeps: {infos[i].hasFwdDeps}"
+        if i < infos.size && !infos[i].hasFwdDeps then
+          r ← mkCongr r (← simp arg)
+        else
+          r ← mkCongrFun r (← dsimp arg)
+        i := i + 1
+      return r
+
+  /- Return true iff processing the given congruence lemma hypothesis produced a non-refl proof. -/
+  processCongrHypothesis (h : Expr) : M Bool := do
+    forallTelescopeReducing (← inferType h) fun xs hType => withNewLemmas xs do
+      let lhs ← instantiateMVars hType.appFn!.appArg!
+      let r ← simp lhs
+      let rhs := hType.appArg!
+      rhs.withApp fun m zs => do
+        let val ← mkLambdaFVars zs r.expr
+        unless (← isDefEq m val) do
+          throwCongrHypothesisFailed
+        unless (← isDefEq h (← mkLambdaFVars xs (← r.getProof))) do
+          throwCongrHypothesisFailed
+        return r.proof?.isSome
+
+  /- Try to rewrite `e` children using the given congruence lemma -/
+  tryCongrLemma? (c : CongrLemma) (e : Expr) : M (Option Result) := withNewMCtxDepth do
+    let info ← getConstInfo c.theoremName
+    let lemma := mkConst c.theoremName (← info.levelParams.mapM fun _ => mkFreshLevelMVar)
+    let (xs, bis, type) ← forallMetaTelescopeReducing (← inferType lemma)
+    if c.hypothesesPos.any (· ≥ xs.size) then
+      return none
+    let lhs := type.appFn!.appArg!
+    let rhs := type.appArg!
+    if (← isDefEq lhs e) then
+      let mut modified := false
+      for i in c.hypothesesPos do
+        let x := xs[i]
+        try
+          if (← processCongrHypothesis x) then
+            modified := true
+        catch _ =>
+          return none
+      unless modified do
+        return none
+      unless (← synthesizeArgs c.theoremName xs bis (← read).discharge?) do
+        return none
+      let eNew ← instantiateMVars rhs
+      return some { expr := eNew, proof? := mkAppN lemma xs }
+    else
+      return none
+
+  congr (e : Expr) : M Result := do
+    let f := e.getAppFn
+    if f.isConst then
+      let congrLemmas ← getCongrLemmas
+      let cs := congrLemmas.get f.constName!
+      for c in cs do
+        match (← tryCongrLemma? c e) with
+        | none   => pure ()
+        | some r => return r
+      congrDefault e
+    else
+      congrDefault e
+
   simpApp (e : Expr) : M Result := do
     let e ← reduce (← getConfig) e
     if !e.isApp then
       simp e
     else
-      withParent e <| e.withApp fun f args => do
-        let infos := (← getFunInfoNArgs f args.size).paramInfo
-        let mut r ← simp f
-        let mut i := 0
-        for arg in args do
-          trace[Meta.Tactic.simp]! "app [{i}] {infos.size} {arg} hasFwdDeps: {infos[i].hasFwdDeps}"
-          if i < infos.size && !infos[i].hasFwdDeps then
-            r ← mkCongr r (← simp arg)
-          else
-            r ← mkCongrFun r (← dsimp arg)
-          i := i + 1
-        return r
+      congr e
 
-  withNewLemmas (xs : Array Expr) (f : M Result) : M Result := do
+  withNewLemmas {α} (xs : Array Expr) (f : M α) : M α := do
     if (← getConfig).contextual then
       let mut s ← getSimpLemmas
       let mut updated := false
@@ -221,16 +284,16 @@ where
       modify fun s => { s with cache := s.cache.insert e r }
     return r
 
-def main (e : Expr) (config : Config := {}) (methods : Methods := {}) (simpLemmas : SimpLemmas := {}) : MetaM Result := do
+def main (e : Expr) (config : Config := {}) (methods : Methods := {}) (simpLemmas : SimpLemmas := {}) (congrLemmas : CongrLemmas := {}) : MetaM Result := do
   withReducible do
-    simp e methods { config := config, simpLemmas := simpLemmas } |>.run' {}
+    simp e methods { config := config, simpLemmas := simpLemmas, congrLemmas := congrLemmas } |>.run' {}
 
 end Simp
 
-def simp (e : Expr) (config : Simp.Config := {}) (simpLemmas : SimpLemmas := {}) : MetaM Simp.Result := do
+def simp (e : Expr) (config : Simp.Config := {}) (simpLemmas : SimpLemmas := {}) (congrLemmas : CongrLemmas := {}) : MetaM Simp.Result := do
   let discharge? (e : Expr) : SimpM (Option Expr) := return none -- TODO: use simp, and add config option
   let pre  := (Simp.preDefault · discharge?)
   let post := (Simp.postDefault · discharge?)
-  Simp.main e (config := config) (methods := { pre := pre, post := post, discharge? := discharge? }) (simpLemmas := simpLemmas)
+  Simp.main e (config := config) (methods := { pre := pre, post := post, discharge? := discharge? }) (simpLemmas := simpLemmas) (congrLemmas := congrLemmas)
 
 end Lean.Meta

src/Lean/Meta/Tactic/Simp/SimpLemmas.lean

@@ -24,13 +24,15 @@ structure SimpLemma where
   kind     : SimpLemmaKind
   deriving Inhabited
 
+def SimpLemma.getName (s : SimpLemma) : Name :=
+  match s.name? with
+  | some n => n
+  | none   => "<unknown>"
+
 instance : ToFormat SimpLemma where
   format s :=
     let perm := if s.perm then ":perm" else ""
-    let name :=
-      match s.name? with
-      | some n => fmt n
-      | none   => "<unknown>"
+    let name := fmt s.getName
     let prio := f!":{s.priority}"
     name ++ prio ++ perm
 

src/Lean/Meta/Tactic/Simp/Types.lean

@@ -5,6 +5,7 @@ Authors: Leonardo de Moura
 -/
 import Lean.Meta.AppBuilder
 import Lean.Meta.Tactic.Simp.SimpLemmas
+import Lean.Meta.Tactic.Simp.CongrLemmas
 
 namespace Lean.Meta
 namespace Simp
@@ -17,9 +18,10 @@ structure Result where
 abbrev Cache := ExprMap Result
 
 structure Context where
-  config     : Config
-  parent?    : Option Expr := none
-  simpLemmas : SimpLemmas
+  config      : Config
+  parent?     : Option Expr := none
+  simpLemmas  : SimpLemmas
+  congrLemmas : CongrLemmas
 
 structure State where
   cache    : Cache := {}
@@ -61,6 +63,9 @@ def getConfig : M Config :=
 def getSimpLemmas : M SimpLemmas :=
   return (← readThe Context).simpLemmas
 
+def getCongrLemmas : M CongrLemmas :=
+  return (← readThe Context).congrLemmas
+
 @[inline] def withSimpLemmas (s : SimpLemmas) (x : M α) : M α := do
   let cacheSaved := (← get).cache
   modify fun s => { s with cache := {} }