feat: add go-to-def for simple attributes

src/Lean/Attributes.lean

@@ -19,6 +19,8 @@ instance : MonadLift ImportM AttrM where
   monadLift x := do liftM (m := IO) (x { env := (← getEnv), opts := (← getOptions) })
 
 structure AttributeImplCore where
+  /-- This is used as the target for go-to-definition queries for simple attributes -/
+  ref : Name := by exact decl_name%
   name : Name
   descr : String
   applicationTime := AttributeApplicationTime.afterTypeChecking
@@ -136,7 +138,8 @@ structure TagAttribute where
   ext  : PersistentEnvExtension Name Name NameSet
   deriving Inhabited
 
-def registerTagAttribute (name : Name) (descr : String) (validate : Name → AttrM Unit := fun _ => pure ()) : IO TagAttribute := do
+def registerTagAttribute (name : Name) (descr : String)
+    (validate : Name → AttrM Unit := fun _ => pure ()) (ref : Name := by exact decl_name%) : IO TagAttribute := do
   let ext : PersistentEnvExtension Name Name NameSet ← registerPersistentEnvExtension {
     name            := name
     mkInitial       := pure {}
@@ -148,6 +151,7 @@ def registerTagAttribute (name : Name) (descr : String) (validate : Name → Att
     statsFn         := fun s => "tag attribute" ++ Format.line ++ "number of local entries: " ++ format s.size
   }
   let attrImpl : AttributeImpl := {
+    ref   := ref
     name  := name
     descr := descr
     add   := fun decl stx kind => do
@@ -184,6 +188,7 @@ structure ParametricAttribute (α : Type) where
   deriving Inhabited
 
 structure ParametricAttributeImpl (α : Type) extends AttributeImplCore where
+  /-- This is used as the target for go-to-definition queries for simple attributes -/
   getParam : Name → Syntax → AttrM α
   afterSet : Name → α → AttrM Unit := fun _ _ _ => pure ()
   afterImport : Array (Array (Name × α)) → ImportM Unit := fun _ => pure ()
@@ -200,6 +205,7 @@ def registerParametricAttribute {α : Type} [Inhabited α] (impl : ParametricAtt
     statsFn         := fun s => "parametric attribute" ++ Format.line ++ "number of local entries: " ++ format s.size
   }
   let attrImpl : AttributeImpl := {
+    ref   := impl.ref
     name  := impl.name
     descr := impl.descr
     add   := fun decl stx kind => do
@@ -246,7 +252,8 @@ structure EnumAttributes (α : Type) where
 
 def registerEnumAttributes {α : Type} [Inhabited α] (extName : Name) (attrDescrs : List (Name × String × α))
     (validate : Name → α → AttrM Unit := fun _ _ => pure ())
-    (applicationTime := AttributeApplicationTime.afterTypeChecking) : IO (EnumAttributes α) := do
+    (applicationTime := AttributeApplicationTime.afterTypeChecking)
+    (ref : Name := by exact decl_name%) : IO (EnumAttributes α) := do
   let ext : PersistentEnvExtension (Name × α) (Name × α) (NameMap α) ← registerPersistentEnvExtension {
     name            := extName
     mkInitial       := pure {}
@@ -258,6 +265,7 @@ def registerEnumAttributes {α : Type} [Inhabited α] (extName : Name) (attrDesc
     statsFn         := fun s => "enumeration attribute extension" ++ Format.line ++ "number of local entries: " ++ format s.size
   }
   let attrs := attrDescrs.map fun (name, descr, val) => {
+    ref             := ref
     name            := name
     descr           := descr
     add             := fun decl stx kind => do
@@ -298,7 +306,7 @@ end EnumAttributes
   Attribute extension and builders. We use builders to implement attribute factories for parser categories.
 -/
 
-abbrev AttributeImplBuilder := List DataValue → Except String AttributeImpl
+abbrev AttributeImplBuilder := Name → List DataValue → Except String AttributeImpl
 abbrev AttributeImplBuilderTable := Std.HashMap Name AttributeImplBuilder
 
 builtin_initialize attributeImplBuilderTableRef : IO.Ref AttributeImplBuilderTable ← IO.mkRef {}
@@ -308,15 +316,15 @@ def registerAttributeImplBuilder (builderId : Name) (builder : AttributeImplBuil
   if table.contains builderId then throw (IO.userError ("attribute implementation builder '" ++ toString builderId ++ "' has already been declared"))
   attributeImplBuilderTableRef.modify fun table => table.insert builderId builder
 
-def mkAttributeImplOfBuilder (builderId : Name) (args : List DataValue) : IO AttributeImpl := do
+def mkAttributeImplOfBuilder (builderId ref : Name) (args : List DataValue) : IO AttributeImpl := do
   let table ← attributeImplBuilderTableRef.get
   match table.find? builderId with
   | none         => throw (IO.userError ("unknown attribute implementation builder '" ++ toString builderId ++ "'"))
-  | some builder => IO.ofExcept <| builder args
+  | some builder => IO.ofExcept <| builder ref args
 
 inductive AttributeExtensionOLeanEntry where
   | decl (declName : Name) -- `declName` has type `AttributeImpl`
-  | builder (builderId : Name) (args : List DataValue)
+  | builder (builderId ref : Name) (args : List DataValue)
 
 structure AttributeExtensionState where
   newEntries : List AttributeExtensionOLeanEntry := []
@@ -342,8 +350,8 @@ opaque mkAttributeImplOfConstant (env : Environment) (opts : Options) (declName
 
 def mkAttributeImplOfEntry (env : Environment) (opts : Options) (e : AttributeExtensionOLeanEntry) : IO AttributeImpl :=
   match e with
-  | AttributeExtensionOLeanEntry.decl declName          => IO.ofExcept <| mkAttributeImplOfConstant env opts declName
-  | AttributeExtensionOLeanEntry.builder builderId args => mkAttributeImplOfBuilder builderId args
+  | .decl declName              => IO.ofExcept <| mkAttributeImplOfConstant env opts declName
+  | .builder builderId ref args => mkAttributeImplOfBuilder builderId ref args
 
 private def AttributeExtension.addImported (es : Array (Array AttributeExtensionOLeanEntry)) : ImportM AttributeExtensionState := do
   let ctx ← read
@@ -409,14 +417,14 @@ def registerAttributeOfDecl (env : Environment) (opts : Options) (attrDeclName :
   if isAttribute env attrImpl.name then
     throw ("invalid builtin attribute declaration, '" ++ toString attrImpl.name ++ "' has already been used")
   else
-    return attributeExtension.addEntry env (AttributeExtensionOLeanEntry.decl attrDeclName, attrImpl)
+    return attributeExtension.addEntry env (.decl attrDeclName, attrImpl)
 
-def registerAttributeOfBuilder (env : Environment) (builderId : Name) (args : List DataValue) : IO Environment := do
-  let attrImpl ← mkAttributeImplOfBuilder builderId args
+def registerAttributeOfBuilder (env : Environment) (builderId ref : Name) (args : List DataValue) : IO Environment := do
+  let attrImpl ← mkAttributeImplOfBuilder builderId ref args
   if isAttribute env attrImpl.name then
     throw (IO.userError ("invalid builtin attribute declaration, '" ++ toString attrImpl.name ++ "' has already been used"))
   else
-    return attributeExtension.addEntry env (AttributeExtensionOLeanEntry.builder builderId args, attrImpl)
+    return attributeExtension.addEntry env (.builder builderId ref args, attrImpl)
 
 def Attribute.add (declName : Name) (attrName : Name) (stx : Syntax) (kind := AttributeKind.global) : AttrM Unit := do
   let attr ← ofExcept <| getAttributeImpl (← getEnv) attrName

src/Lean/Compiler/InitAttr.lean

@@ -32,10 +32,11 @@ unsafe opaque runModInit (mod : Name) : IO Bool
 @[extern "lean_run_init"]
 unsafe opaque runInit (env : @& Environment) (opts : @& Options) (decl initDecl : @& Name) : IO Unit
 
-unsafe def registerInitAttrUnsafe (attrName : Name) (runAfterImport : Bool) : IO (ParametricAttribute Name) :=
+unsafe def registerInitAttrUnsafe (attrName : Name) (runAfterImport : Bool) (ref : Name) : IO (ParametricAttribute Name) :=
   registerParametricAttribute {
-    name := attrName,
-    descr := "initialization procedure for global references",
+    ref := ref
+    name := attrName
+    descr := "initialization procedure for global references"
     getParam := fun declName stx => do
       let decl ← getConstInfo declName
       match (← Attribute.Builtin.getIdent? stx) with
@@ -70,7 +71,11 @@ unsafe def registerInitAttrUnsafe (attrName : Name) (runAfterImport : Bool) : IO
   }
 
 @[implementedBy registerInitAttrUnsafe]
-opaque registerInitAttr (attrName : Name) (runAfterImport : Bool) : IO (ParametricAttribute Name)
+private opaque registerInitAttrInner (attrName : Name) (runAfterImport : Bool) (ref : Name) : IO (ParametricAttribute Name)
+
+@[inline]
+def registerInitAttr (attrName : Name) (runAfterImport : Bool) (ref : Name := by exact decl_name%) : IO (ParametricAttribute Name) :=
+  registerInitAttrInner attrName runAfterImport ref
 
 builtin_initialize regularInitAttr : ParametricAttribute Name ← registerInitAttr `init true
 builtin_initialize builtinInitAttr : ParametricAttribute Name ← registerInitAttr `builtinInit false

src/Lean/Compiler/Specialize.lean

@@ -25,9 +25,7 @@ builtin_initialize specializeAttrs : EnumAttributes SpecializeAttributeKind ←
        In the new equation compiler we should pass all attributes and allow it to apply them to auxiliary definitions.
        In the current implementation, we workaround this issue by using functions such as `hasSpecializeAttrAux`.
      -/
-    (fun _ _ => pure ())
-    AttributeApplicationTime.beforeElaboration
-
+    (applicationTime := .beforeElaboration)
 private partial def hasSpecializeAttrAux (env : Environment) (kind : SpecializeAttributeKind) (n : Name) : Bool :=
   match specializeAttrs.getValue env n with
   | some k => kind == k

src/Lean/Elab/App.lean

@@ -660,16 +660,15 @@ end
 end ElabAppArgs
 
 builtin_initialize elabAsElim : TagAttribute ←
-  registerTagAttribute
-    `elabAsElim
+  registerTagAttribute `elabAsElim
     "instructs elaborator that the arguments of the function application should be elaborated as were an eliminator"
-   fun declName => do
-     let go : MetaM Unit := do
-       discard <| getElimInfo declName
-       let info ← getConstInfo declName
-       if (← hasOptAutoParams info.type) then
-         throwError "[elabAsElim] attribute cannot be used in declarations containing optional and auto parameters"
-     go.run' {} {}
+    fun declName => do
+      let go : MetaM Unit := do
+        discard <| getElimInfo declName
+        let info ← getConstInfo declName
+        if (← hasOptAutoParams info.type) then
+          throwError "[elabAsElim] attribute cannot be used in declarations containing optional and auto parameters"
+      go.run' {} {}
 
 /-! # Eliminator-like function application elaborator -/
 namespace ElabElim

src/Lean/Elab/Attributes.lean

@@ -51,20 +51,20 @@ def elabAttr [Monad m] [MonadEnv m] [MonadResolveName m] [MonadError m] [MonadMa
   else match attr.getKind with
     | .str _ s => pure <| Name.mkSimple s
     | _ => throwErrorAt attr  "unknown attribute"
-  unless isAttribute (← getEnv) attrName do
-    throwError "unknown attribute [{attrName}]"
+  let .ok impl := getAttributeImpl (← getEnv) attrName
+    | throwError "unknown attribute [{attrName}]"
   let attrSyntaxNodeKind := attrInstance[1].getKind
-  unless attrSyntaxNodeKind == ``Lean.Parser.Attr.simple do
-    -- `Lean.Parser.Attr.simple` is a generic `attribute` parser used in simple attributes.
-    -- We don't want to create an info tree node from a simple attribute to the generic parser.
-    if (← getEnv).contains attrSyntaxNodeKind && (← getInfoState).enabled then
-      pushInfoLeaf <| Info.ofTermInfo {
-        elaborator    := .anonymous
-        lctx          := {}
-        expr          := mkConst attrSyntaxNodeKind
-        stx           := attrInstance[1][0] -- We want to associate the information to the first atom only
-        expectedType? := none
-      }
+  -- `Lean.Parser.Attr.simple` is a generic `attribute` parser used in simple attributes.
+  -- We don't want to create an info tree node from a simple attribute to the generic parser.
+  let declTarget := if attrSyntaxNodeKind == ``Lean.Parser.Attr.simple then impl.ref else attrSyntaxNodeKind
+  if (← getEnv).contains declTarget && (← getInfoState).enabled then
+    pushInfoLeaf <| Info.ofTermInfo {
+      elaborator    := .anonymous
+      lctx          := {}
+      expr          := mkConst declTarget
+      stx           := attrInstance[1][0] -- We want to associate the information to the first atom only
+      expectedType? := none
+    }
   /- The `AttrM` does not have sufficient information for expanding macros in `args`.
      So, we expand them before here before we invoke the attributer handlers implemented using `AttrM`. -/
   return { kind := attrKind, name := attrName, stx := attr }

src/Lean/Elab/Syntax.lean

@@ -272,9 +272,9 @@ private def declareSyntaxCatQuotParser (catName : Name) : CommandElabM Unit := d
     else
       Parser.LeadingIdentBehavior.symbol
   let attrName := catName.appendAfter "Parser"
-  setEnv (← Parser.registerParserCategory (← getEnv) attrName catName catBehavior)
-  let catDeclName := `_root_ ++ ``Lean.Parser.Category ++ catName
-  let cmd ← `($[$docString?]? def $(mkIdentFrom stx[2] catDeclName) : Lean.Parser.Category := {})
+  let catDeclName := ``Lean.Parser.Category ++ catName
+  setEnv (← Parser.registerParserCategory (← getEnv) attrName catName catBehavior catDeclName)
+  let cmd ← `($[$docString?]? def $(mkIdentFrom stx[2] (`_root_ ++ catDeclName)) : Lean.Parser.Category := {})
   declareSyntaxCatQuotParser catName
   elabCommand cmd
 

src/Lean/Elab/Term.lean

@@ -870,7 +870,7 @@ Extensions for monads.
 
   ```
 
-3. If there is a monad lif from `m` to `n` and a coercion from `α` to `β`, we use
+3. If there is a monad lift from `m` to `n` and a coercion from `α` to `β`, we use
   ```
   liftCoeM {m : Type u → Type v} {n : Type u → Type w} {α β : Type u} [MonadLiftT m n] [∀ a, CoeT α a β] [Monad n] (x : m α) : n β
   ```

src/Lean/KeyedDeclsAttribute.lean

@@ -117,6 +117,7 @@ protected unsafe def init {γ} (df : Def γ) (attrDeclName : Name) : IO (KeyedDe
   }
   unless df.builtinName.isAnonymous do
     registerBuiltinAttribute {
+      ref   := attrDeclName
       name  := df.builtinName
       descr := "(builtin) " ++ df.descr
       add   := fun declName stx kind => do
@@ -135,6 +136,7 @@ protected unsafe def init {γ} (df : Def γ) (attrDeclName : Name) : IO (KeyedDe
       applicationTime := AttributeApplicationTime.afterCompilation
     }
   registerBuiltinAttribute {
+    ref             := attrDeclName
     name            := df.name
     descr           := df.descr
     erase           := fun declName => do

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

@@ -274,8 +274,10 @@ def addSimpTheorem (ext : SimpExtension) (declName : Name) (post : Bool) (inv :
   for simpThm in simpThms do
     ext.add (SimpEntry.thm simpThm) attrKind
 
-def mkSimpAttr (attrName : Name) (attrDescr : String) (ext : SimpExtension) : IO Unit :=
+def mkSimpAttr (attrName : Name) (attrDescr : String) (ext : SimpExtension)
+    (ref : Name := by exact decl_name%) : IO Unit :=
   registerBuiltinAttribute {
+    ref   := ref
     name  := attrName
     descr := attrDescr
     applicationTime := AttributeApplicationTime.afterCompilation
@@ -319,9 +321,11 @@ abbrev SimpExtensionMap := Std.HashMap Name SimpExtension
 
 builtin_initialize simpExtensionMapRef : IO.Ref SimpExtensionMap ← IO.mkRef {}
 
-def registerSimpAttr (attrName : Name) (attrDescr : String) (extName : Name := attrName.appendAfter "Ext") : IO SimpExtension := do
+def registerSimpAttr (attrName : Name) (attrDescr : String)
+    (ref : Name := by exact decl_name%)
+    (extName : Name := attrName.appendAfter "Ext") : IO SimpExtension := do
   let ext ← mkSimpExt extName
-  mkSimpAttr attrName attrDescr ext -- Remark: it will fail if it is not performed during initialization
+  mkSimpAttr attrName attrDescr ext ref -- Remark: it will fail if it is not performed during initialization
   simpExtensionMapRef.modify fun map => map.insert attrName ext
   return ext
 
@@ -414,7 +418,7 @@ def SimpTheoremsArray.isDeclToUnfold (thmsArray : SimpTheoremsArray) (declName :
 macro doc?:(docComment)? "register_simp_attr" id:ident descr:str : command => do
   let str := id.getId.toString
   let idParser := mkIdentFrom id (`Parser.Attr ++ id.getId)
-  `($[$doc?]? initialize ext : SimpExtension ← registerSimpAttr $(quote id.getId) $descr
+  `($[$doc?]? initialize ext : SimpExtension ← registerSimpAttr $(quote id.getId) $descr $(quote id.getId)
     $[$doc?]? syntax (name := $idParser:ident) $(quote str):str (Parser.Tactic.simpPre <|> Parser.Tactic.simpPost)? (prio)? : attr)
 
 end Meta

src/Lean/Parser/Extension.lean

@@ -308,8 +308,8 @@ def runParserAttributeHooks (catName : Name) (declName : Name) (builtin : Bool)
 
 builtin_initialize
   registerBuiltinAttribute {
-    name  := `runBuiltinParserAttributeHooks,
-    descr := "explicitly run hooks normally activated by builtin parser attributes",
+    name  := `runBuiltinParserAttributeHooks
+    descr := "explicitly run hooks normally activated by builtin parser attributes"
     add   := fun decl stx _ => do
       Attribute.Builtin.ensureNoArgs stx
       runParserAttributeHooks Name.anonymous decl (builtin := true)
@@ -317,8 +317,8 @@ builtin_initialize
 
 builtin_initialize
   registerBuiltinAttribute {
-    name  := `runParserAttributeHooks,
-    descr := "explicitly run hooks normally activated by parser attributes",
+    name  := `runParserAttributeHooks
+    descr := "explicitly run hooks normally activated by parser attributes"
     add   := fun decl stx _ => do
       Attribute.Builtin.ensureNoArgs stx
       runParserAttributeHooks Name.anonymous decl (builtin := false)
@@ -502,13 +502,15 @@ private def BuiltinParserAttribute.add (attrName : Name) (catName : Name)
 /--
 The parsing tables for builtin parsers are "stored" in the extracted source code.
 -/
-def registerBuiltinParserAttribute (attrName : Name) (catName : Name) (behavior := LeadingIdentBehavior.default) : IO Unit := do
+def registerBuiltinParserAttribute (attrName : Name) (catName : Name)
+    (behavior := LeadingIdentBehavior.default) (ref : Name := by exact decl_name%) : IO Unit := do
   addBuiltinParserCategory catName behavior
   registerBuiltinAttribute {
-   name            := attrName,
-   descr           := "Builtin parser",
-   add             := fun declName stx kind => liftM $ BuiltinParserAttribute.add attrName catName declName stx kind,
-   applicationTime := AttributeApplicationTime.afterCompilation
+    ref             := ref
+    name            := attrName
+    descr           := "Builtin parser"
+    add             := fun declName stx kind => liftM $ BuiltinParserAttribute.add attrName catName declName stx kind
+    applicationTime := AttributeApplicationTime.afterCompilation
   }
 
 private def ParserAttribute.add (_attrName : Name) (catName : Name) (declName : Name) (stx : Syntax) (attrKind : AttributeKind) : AttrM Unit := do
@@ -533,25 +535,27 @@ private def ParserAttribute.add (_attrName : Name) (catName : Name) (declName :
   | Except.ok _     => parserExtension.add entry attrKind
   runParserAttributeHooks catName declName (builtin := false)
 
-def mkParserAttributeImpl (attrName : Name) (catName : Name) : AttributeImpl where
+def mkParserAttributeImpl (attrName catName : Name) (ref : Name := by exact decl_name%) : AttributeImpl where
+  ref                       := ref
   name                      := attrName
   descr                     := "parser"
   add declName stx attrKind := ParserAttribute.add attrName catName declName stx attrKind
   applicationTime           := AttributeApplicationTime.afterCompilation
 
 /-- A builtin parser attribute that can be extended by users. -/
-def registerBuiltinDynamicParserAttribute (attrName : Name) (catName : Name) : IO Unit := do
-  registerBuiltinAttribute (mkParserAttributeImpl attrName catName)
+def registerBuiltinDynamicParserAttribute (attrName catName : Name) (ref : Name := by exact decl_name%) : IO Unit := do
+  registerBuiltinAttribute (mkParserAttributeImpl attrName catName ref)
 
 builtin_initialize
-  registerAttributeImplBuilder `parserAttr fun args =>
+  registerAttributeImplBuilder `parserAttr fun ref args =>
     match args with
-    | [DataValue.ofName attrName, DataValue.ofName catName] => pure $ mkParserAttributeImpl attrName catName
+    | [DataValue.ofName attrName, DataValue.ofName catName] => pure $ mkParserAttributeImpl attrName catName ref
     | _ => throw "invalid parser attribute implementation builder arguments"
 
-def registerParserCategory (env : Environment) (attrName : Name) (catName : Name) (behavior := LeadingIdentBehavior.default) : IO Environment := do
+def registerParserCategory (env : Environment) (attrName catName : Name)
+    (behavior := LeadingIdentBehavior.default) (ref : Name := by exact decl_name%) : IO Environment := do
   let env ← IO.ofExcept $ addParserCategory env catName behavior
-  registerAttributeOfBuilder env `parserAttr [DataValue.ofName attrName, DataValue.ofName catName]
+  registerAttributeOfBuilder env `parserAttr ref [DataValue.ofName attrName, DataValue.ofName catName]
 
 -- declare `termParser` here since it is used everywhere via antiquotations