feat: invoke simproc attribute handlers from simp attribute handlers

leanprover · Feb 26, 2024 · 5171aed · 5171aed
1 parent dde5e2d
commit 5171aed
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 27 deletions.
diff --git a/src/Lean/Meta/Tactic/Simp/Attr.lean b/src/Lean/Meta/Tactic/Simp/Attr.lean
@@ -4,10 +4,12 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 prelude
-import Lean.Meta.Tactic.Simp.SimpTheorems
 import Lean.Meta.Tactic.Simp.Types
+import Lean.Meta.Tactic.Simp.SimpTheorems
+import Lean.Meta.Tactic.Simp.Simproc
 
 namespace Lean.Meta
+open Simp
 
 def mkSimpAttr (attrName : Name) (attrDescr : String) (ext : SimpExtension)
     (ref : Name := by exact decl_name%) : IO Unit :=
@@ -16,29 +18,37 @@ def mkSimpAttr (attrName : Name) (attrDescr : String) (ext : SimpExtension)
     name  := attrName
     descr := attrDescr
     applicationTime := AttributeApplicationTime.afterCompilation
-    add   := fun declName stx attrKind =>
-      let go : MetaM Unit := do
-        let info ← getConstInfo declName
-        let post := if stx[1].isNone then true else stx[1][0].getKind == ``Lean.Parser.Tactic.simpPost
-        let prio ← getAttrParamOptPrio stx[2]
-        if (← isProp info.type) then
-          addSimpTheorem ext declName post (inv := false) attrKind prio
-        else if info.hasValue then
-          if let some eqns ← getEqnsFor? declName then
-            for eqn in eqns do
-              addSimpTheorem ext eqn post (inv := false) attrKind prio
-            ext.add (SimpEntry.toUnfoldThms declName eqns) attrKind
-            if hasSmartUnfoldingDecl (← getEnv) declName then
+    add   := fun declName stx attrKind => do
+      if (← isSimproc declName <||> isBuiltinSimproc declName) then
+        let simprocAttrName := simpAttrNameToSimprocAttrName attrName
+        Attribute.add declName simprocAttrName stx attrKind
+      else
+        let go : MetaM Unit := do
+          let info ← getConstInfo declName
+          let post := if stx[1].isNone then true else stx[1][0].getKind == ``Lean.Parser.Tactic.simpPost
+          let prio ← getAttrParamOptPrio stx[2]
+          if (← isProp info.type) then
+            addSimpTheorem ext declName post (inv := false) attrKind prio
+          else if info.hasValue then
+            if let some eqns ← getEqnsFor? declName then
+              for eqn in eqns do
+                addSimpTheorem ext eqn post (inv := false) attrKind prio
+              ext.add (SimpEntry.toUnfoldThms declName eqns) attrKind
+              if hasSmartUnfoldingDecl (← getEnv) declName then
+                ext.add (SimpEntry.toUnfold declName) attrKind
+            else
               ext.add (SimpEntry.toUnfold declName) attrKind
           else
-            ext.add (SimpEntry.toUnfold declName) attrKind
-        else
-          throwError "invalid 'simp', it is not a proposition nor a definition (to unfold)"
-      discard <| go.run {} {}
+            throwError "invalid 'simp', it is not a proposition nor a definition (to unfold)"
+        discard <| go.run {} {}
     erase := fun declName => do
-      let s := ext.getState (← getEnv)
-      let s ← s.erase (.decl declName)
-      modifyEnv fun env => ext.modifyState env fun _ => s
+      if (← isSimproc declName <||> isBuiltinSimproc declName) then
+        let simprocAttrName := simpAttrNameToSimprocAttrName attrName
+        Attribute.erase declName simprocAttrName
+      else
+        let s := ext.getState (← getEnv)
+        let s ← s.erase (.decl declName)
+        modifyEnv fun env => ext.modifyState env fun _ => s
   }
 
 def registerSimpAttr (attrName : Name) (attrDescr : String)

diff --git a/src/Lean/Meta/Tactic/Simp/RegisterCommand.lean b/src/Lean/Meta/Tactic/Simp/RegisterCommand.lean
@@ -6,6 +6,7 @@ Authors: Leonardo de Moura
 prelude
 import Lean.Meta.Tactic.Simp.SimpTheorems
 import Lean.Meta.Tactic.Simp.Simproc
+import Lean.Meta.Tactic.Simp.Attr
 
 namespace Lean.Meta.Simp
 

diff --git a/src/Lean/Meta/Tactic/Simp/Simproc.lean b/src/Lean/Meta/Tactic/Simp/Simproc.lean
@@ -130,7 +130,7 @@ def eraseSimprocAttr (ext : SimprocExtension) (declName : Name) : AttrM Unit :=
     throwError "'{declName}' does not have a simproc attribute"
   modifyEnv fun env => ext.modifyState env fun s => s.erase declName
 
-def addSimprocAttr (ext : SimprocExtension) (declName : Name) (kind : AttributeKind) (post : Bool) : CoreM Unit := do
+def addSimprocAttrCore (ext : SimprocExtension) (declName : Name) (kind : AttributeKind) (post : Bool) : CoreM Unit := do
   let proc ← getSimprocFromDecl declName
   let some keys ← getSimprocDeclKeys? declName |
     throwError "invalid [simproc] attribute, '{declName}' is not a simproc"
@@ -280,17 +280,19 @@ def mkSimprocExt (name : Name := by exact decl_name%) (ref? : Option (IO.Ref Sim
     addEntry      := fun s e => s.addCore e.keys e.declName e.post e.proc
   }
 
+def addSimprocAttr (ext : SimprocExtension) (declName : Name) (stx : Syntax) (attrKind : AttributeKind) : AttrM Unit := do
+  let go : MetaM Unit := do
+    let post := if stx[1].isNone then true else stx[1][0].getKind == ``Lean.Parser.Tactic.simpPost
+    addSimprocAttrCore ext declName attrKind post
+  discard <| go.run {} {}
+
 def mkSimprocAttr (attrName : Name) (attrDescr : String) (ext : SimprocExtension) (name : Name) : IO Unit := do
   registerBuiltinAttribute {
     ref   := name
     name  := attrName
     descr := attrDescr
     applicationTime := AttributeApplicationTime.afterCompilation
-    add   := fun declName stx attrKind =>
-      let go : MetaM Unit := do
-        let post := if stx[1].isNone then true else stx[1][0].getKind == ``Lean.Parser.Tactic.simpPost
-        addSimprocAttr ext declName attrKind post
-      discard <| go.run {} {}
+    add   := addSimprocAttr ext
     erase := eraseSimprocAttr ext
   }