feat: add @[congr] attribute

src/Lean/Meta/Tactic/Simp.lean

@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 import Lean.Meta.Tactic.Simp.SimpLemmas
+import Lean.Meta.Tactic.Simp.CongrLemmas
 import Lean.Meta.Tactic.Simp.Types
 import Lean.Meta.Tactic.Simp.Main
 import Lean.Meta.Tactic.Simp.Rewrite

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

@@ -0,0 +1,118 @@
+/-
+Copyright (c) 2021 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+import Lean.ScopedEnvExtension
+import Lean.Util.Recognizers
+import Lean.Meta.Basic
+
+namespace Lean.Meta
+
+structure CongrLemma where
+  theoremName   : Name
+  funName       : Name
+  hypothesesPos : Array Nat
+  priority      : Nat
+deriving Inhabited, Repr
+
+structure CongrLemmas where
+  lemmas : SMap Name (List CongrLemma) := {}
+  deriving Inhabited
+
+def addCongrLemmaEntry (d : CongrLemmas) (e : CongrLemma) : CongrLemmas :=
+  { d with lemmas :=
+      match d.lemmas.find? e.funName with
+      | none    => d.lemmas.insert e.funName [e]
+      | some es => d.lemmas.insert e.funName <| insert es }
+where
+  insert : List CongrLemma → List CongrLemma
+    | []     => [e]
+    | e'::es => if e.priority ≥ e'.priority then e::e'::es else e' :: insert es
+
+builtin_initialize congrExtension : SimpleScopedEnvExtension CongrLemma CongrLemmas ←
+  registerSimpleScopedEnvExtension {
+    name           := `congrExt
+    initial        := {}
+    addEntry       := addCongrLemmaEntry
+    finalizeImport := fun s => { s with lemmas := s.lemmas.switch }
+  }
+
+def mkCongrLemma (declName : Name) (prio : Nat) : MetaM CongrLemma := withReducible do
+  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) =>
+    lhs.withApp fun lhsFn lhsArgs => rhs.withApp fun rhsFn rhsArgs => do
+      unless lhsFn.isConst && rhsFn.isConst && lhsFn.constName! == rhsFn.constName! && lhsArgs.size == rhsArgs.size do
+        throwError! "invalid 'congr' lemma, equality left/right-hand sides must be applications of the same function{indentExpr type}"
+      let mut foundMVars : NameSet := {}
+      for lhsArg in lhsArgs do
+        unless lhsArg.isSort do
+          unless lhsArg.isMVar do
+            throwError! "invalid 'congr' lemma, arguments in the left-hand-side must be variables or sorts{indentExpr lhs}"
+          foundMVars := foundMVars.insert lhsArg.mvarId!
+      let mut i := 0
+      let mut hypothesesPos := #[]
+      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) =>
+              let mut j := 0
+              for y in ys do
+                let yType ← inferType y
+                unless onlyMVarsAt yType foundMVars do
+                  throwError! "invalid 'congr' lemma, argument #{j+1} of parameter #{i+1} contains unresolved parameter{indentExpr yType}"
+                j := j + 1
+              unless onlyMVarsAt xLhs foundMVars do
+                throwError! "invalid 'congr' lemma, parameter #{i+1} is not a valid hypothesis, the left-hand-side contains unresolved parameters{indentExpr xLhs}"
+              let xRhsFn := xRhs.getAppFn
+              unless xRhsFn.isMVar do
+                throwError! "invalid 'congr' lemma, parameter #{i+1} is not a valid hypothesis, the right-hand-side head is not a metavariable{indentExpr xRhs}"
+              unless !foundMVars.contains xRhsFn.mvarId! do
+                throwError! "invalid 'congr' lemma, parameter #{i+1} is not a valid hypothesis, the right-hand-side head was already resolved{indentExpr xRhs}"
+              for arg in xRhs.getAppArgs do
+                unless arg.isFVar do
+                  throwError! "invalid 'congr' lemma, parameter #{i+1} is not a valid hypothesis, the right-hand-side argument is not local variable{indentExpr xRhs}"
+              pure (some xRhsFn)
+          match rhsFn? with
+          | none       => pure ()
+          | some rhsFn =>
+            foundMVars    := foundMVars.insert x.mvarId! |>.insert rhsFn.mvarId!
+            hypothesesPos := hypothesesPos.push i
+        i := i + 1
+      trace[Meta.debug]! "c: {c} : {type}"
+      return {
+        theoremName   := declName
+        funName       := lhsFn.constName!
+        hypothesesPos := hypothesesPos
+        priority      := prio
+      }
+where
+  /-- Return `true` if `t` contains a metavariable that is not in `mvarSet` -/
+  onlyMVarsAt (t : Expr) (mvarSet : NameSet) : Bool :=
+    Option.isNone <| t.find? fun e => e.isMVar && !mvarSet.contains e.mvarId!
+
+def addCongrLemma (declName : Name) (attrKind : AttributeKind) (prio : Nat) : MetaM Unit := do
+  let lemma ← mkCongrLemma declName prio
+  congrExtension.add lemma attrKind
+
+builtin_initialize
+  registerBuiltinAttribute {
+    name  := `congr
+    descr := "congruence lemma"
+    add   := fun declName stx attrKind => do
+      let prio ← getAttrParamOptPrio stx[1]
+      discard <| addCongrLemma declName attrKind prio |>.run {} {}
+  }
+
+def getCongrLemmas : MetaM CongrLemmas :=
+  return congrExtension.getState (← getEnv)
+
+end Lean.Meta

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

@@ -94,7 +94,6 @@ def addSimpLemma (declName : Name) (post : Bool) (attrKind : AttributeKind) (pri
      See `SimpLemma.getValue` -/
   let lemma ← mkSimpLemmaCore (mkConst declName (cinfo.levelParams.map mkLevelParam)) (mkConst declName) post prio declName
   simpExtension.add lemma attrKind
-  pure ()
 
 builtin_initialize
   registerBuiltinAttribute {

src/Lean/ScopedEnvExtension.lean

@@ -35,6 +35,7 @@ structure Descr (α : Type) (β : Type) (σ : Type) where
   ofOLeanEntry   : σ → α → ImportM β
   toOLeanEntry   : β → α
   addEntry       : σ → β → σ
+  finalizeImport : σ → σ := id
 
 instance [Inhabited α] : Inhabited (Descr α β σ) where
   default := {
@@ -65,6 +66,7 @@ def addImportedFn (descr : Descr α β σ) (as : Array (Array (Entry α))) : Imp
       | Entry.scoped ns a =>
         let b ← descr.ofOLeanEntry s a
         scopedEntries := scopedEntries.insert ns b
+  s := descr.finalizeImport s
   return { stateStack := [ { state := s } ], scopedEntries := scopedEntries }
 
 def addEntryFn (descr : Descr α β σ) (s : StateStack α β σ) (e : Entry β) : StateStack α β σ :=
@@ -192,14 +194,16 @@ structure SimpleScopedEnvExtension.Descr (α : Type) (σ : Type) where
   name           : Name
   addEntry       : σ → α → σ
   initial        : σ
+  finalizeImport : σ → σ := id
 
 def registerSimpleScopedEnvExtension (descr : SimpleScopedEnvExtension.Descr α σ) : IO (SimpleScopedEnvExtension α σ) := do
   registerScopedEnvExtension {
-    name          := descr.name
-    mkInitial     := return descr.initial
-    addEntry      := descr.addEntry
-    toOLeanEntry  := id
-    ofOLeanEntry  := fun s a => return a
+    name           := descr.name
+    mkInitial      := return descr.initial
+    addEntry       := descr.addEntry
+    toOLeanEntry   := id
+    ofOLeanEntry   := fun s a => return a
+    finalizeImport := descr.finalizeImport
   }
 
 end Lean