feat: add @[attribute] for registering custom attributes

src/Lean/Attributes.lean

@@ -87,7 +87,7 @@ def registerBuiltinAttribute (attr : AttributeImpl) : IO Unit := do
 def Attribute.Builtin.ensureNoArgs (stx : Syntax) : AttrM Unit := do
   if stx.getKind == `Lean.Parser.Attr.simple && stx[1].isNone && stx[2].isNone then
     return ()
-  else if stx.getKind == `Lean.Parser.Attr.«class» then
+  else if stx.getKind == `Lean.Parser.Attr.«class» || stx.getKind == `Lean.Parser.Attr.«attribute» then
     return ()
   else match stx with
     | Syntax.missing => return () -- In the elaborator, we use `Syntax.missing` when creating attribute views for simple attributes such as `class and `inline
@@ -457,6 +457,45 @@ builtin_initialize attributeExtension : AttributeExtension ←
     statsFn         := fun s => format "number of local entries: " ++ format s.newEntries.length
   }
 
+/-!
+  `attributeExt`: a separate environment extension for attributes registered via
+  the `@[attribute]` attribute (i.e., as Lean constants of type `AttributeImpl`).
+  Kept distinct from `attributeExtension` (which serves builtin attributes and
+  parser-category builder entries) so that the two registration mechanisms remain
+  cleanly separated.
+-/
+
+structure CustomAttributeExtensionState where
+  newEntries : Array Name := #[]
+  map : Std.HashMap Name AttributeImpl := {}
+  deriving Inhabited
+
+abbrev CustomAttributeExtension :=
+  PersistentEnvExtension Name (Name × AttributeImpl) CustomAttributeExtensionState
+
+private def CustomAttributeExtension.addImported
+    (es : Array (Array Name)) : ImportM CustomAttributeExtensionState := do
+  let ctx ← read
+  let mut map : Std.HashMap Name AttributeImpl := {}
+  for entries in es do
+    for declName in entries do
+      let attrImpl ← IO.ofExcept <| mkAttributeImplOfConstant ctx.env ctx.opts declName
+      map := map.insert attrImpl.name attrImpl
+  pure { map := map }
+
+private def addCustomAttrEntry
+    (s : CustomAttributeExtensionState) (e : Name × AttributeImpl) : CustomAttributeExtensionState :=
+  { s with map := s.map.insert e.2.name e.2, newEntries := s.newEntries.push e.1 }
+
+builtin_initialize attributeExt : CustomAttributeExtension ←
+  registerPersistentEnvExtension {
+    mkInitial       := pure {}
+    addImportedFn   := CustomAttributeExtension.addImported
+    addEntryFn      := addCustomAttrEntry
+    exportEntriesFn := fun s => s.newEntries
+    statsFn         := fun s => format "number of custom attributes: " ++ format s.newEntries.size
+  }
+
 /-- Return true iff `n` is the name of a registered attribute. -/
 @[export lean_is_attribute]
 def isBuiltinAttribute (n : Name) : IO Bool := do
@@ -478,17 +517,21 @@ def getBuiltinAttributeApplicationTime (n : Name) : IO AttributeApplicationTime
   pure attr.applicationTime
 
 def isAttribute (env : Environment) (attrName : Name) : Bool :=
-  (attributeExtension.getState env).map.contains attrName
+  (attributeExtension.getState env).map.contains attrName ||
+  (attributeExt.getState env).map.contains attrName
 
 def getAttributeNames (env : Environment) : List Name :=
-  let m := (attributeExtension.getState env).map
-  m.fold (fun r n _ => n::r) []
+  let m1 := (attributeExtension.getState env).map
+  let m2 := (attributeExt.getState env).map
+  m2.fold (init := m1.fold (init := []) fun r n _ => n::r) (fun r n _ => n::r)
 
 def getAttributeImpl (env : Environment) (attrName : Name) : Except String AttributeImpl :=
-  let m := (attributeExtension.getState env).map
-  match m[attrName]? with
+  match (attributeExtension.getState env).map[attrName]? with
   | some attr => pure attr
-  | none      => throw s!"Unknown attribute `{attrName}`"
+  | none =>
+    match (attributeExt.getState env).map[attrName]? with
+    | some attr => pure attr
+    | none      => throw s!"Unknown attribute `{attrName}`"
 
 def registerAttributeOfBuilder (env : Environment) (builderId ref : Name) (args : List DataValue) : IO Environment := do
   let entry := {builderId, ref, args}
@@ -498,6 +541,19 @@ def registerAttributeOfBuilder (env : Environment) (builderId ref : Name) (args
   else
     return attributeExtension.addEntry env (entry, attrImpl)
 
+/--
+  Register a declaration of type `Lean.AttributeImpl` as a custom attribute.
+  Used by the `@[attribute]` attribute to allow attribute registration outside
+  of `initialize` blocks.
+-/
+def registerAttributeOfDecl (declName : Name) : AttrM Unit := do
+  let env ← getEnv
+  let opts ← getOptions
+  let attrImpl ← ofExcept <| mkAttributeImplOfConstant env opts declName
+  if isAttribute env attrImpl.name then
+    throwError m!"Invalid attribute declaration: `{attrImpl.name}` has already been registered"
+  setEnv <| attributeExt.addEntry env (declName, attrImpl)
+
 def Attribute.add (declName : Name) (attrName : Name) (stx : Syntax) (kind := AttributeKind.global) : AttrM Unit := do
   let attr ← ofExcept <| getAttributeImpl (← getEnv) attrName
   attr.add declName stx kind
@@ -522,4 +578,15 @@ def updateEnvAttributesImpl (env : Environment) : IO Environment := do
 @[export lean_get_num_attributes] def getNumBuiltinAttributesImpl : IO Nat :=
   return (← attributeMapRef.get).size
 
+builtin_initialize
+  registerBuiltinAttribute {
+    name  := `attribute
+    descr := "register the tagged declaration (which must have type `Lean.AttributeImpl`) as a custom attribute"
+    applicationTime := .afterCompilation
+    add   := fun declName stx kind => do
+      Attribute.Builtin.ensureNoArgs stx
+      unless kind == .global do throwAttrMustBeGlobal `attribute kind
+      registerAttributeOfDecl declName
+  }
+
 end Lean

src/Lean/Parser/Attr.lean

@@ -37,9 +37,10 @@ end Priority
 namespace Attr
 
 @[builtin_attr_parser] def simple     := leading_parser ident >> optional (ppSpace >> (priorityParser <|> ident))
-/- Remark, We can't use `simple` for `class`, `instance`, `export`, and `macro` because they are keywords. -/
+/- Remark, We can't use `simple` for `class`, `instance`, `export`, `macro`, and `attribute` because they are keywords. -/
 @[builtin_attr_parser] def «macro»    := leading_parser "macro " >> ident
 @[builtin_attr_parser] def «export»   := leading_parser "export " >> ident
+@[builtin_attr_parser] def «attribute» := leading_parser "attribute"
 
 /- We don't use `simple` for recursor because the argument is not a priority -/
 @[builtin_attr_parser] def recursor         := leading_parser nonReservedSymbol "recursor " >> numLit

tests/elab/attributeAttr.lean

@@ -0,0 +1,74 @@
+import Lean
+
+/-!
+Test the `@[attribute]` attribute, which registers a `Lean.AttributeImpl`
+constant as a custom attribute without requiring an `initialize` block.
+-/
+
+open Lean
+
+/-- A simple custom attribute that just logs the declaration it's applied to. -/
+@[attribute] def myCustomAttr : AttributeImpl where
+  name := `myCustomAttr
+  descr := "a test custom attribute"
+  applicationTime := .afterTypeChecking
+  add := fun decl _stx _kind => do
+    logInfo m!"myCustomAttr applied to `{decl}`"
+
+/-- info: myCustomAttr applied to `foo` -/
+#guard_msgs in
+@[myCustomAttr] def foo := 42
+
+/-- info: myCustomAttr applied to `bar` -/
+#guard_msgs in
+@[myCustomAttr] def bar := "hello"
+
+/-!
+The custom attribute should be visible via `getAttributeImpl` / `isAttribute`.
+-/
+
+run_meta do
+  unless isAttribute (← getEnv) `myCustomAttr do
+    throwError "myCustomAttr is not a registered attribute"
+
+/-!
+Re-registering the same attribute name should be rejected.
+-/
+
+def myCustomAttr' : AttributeImpl where
+  name := `myCustomAttr  -- collides with the one above
+  descr := "duplicate"
+  add := fun _ _ _ => pure ()
+
+/--
+error: Invalid attribute declaration: `myCustomAttr` has already been registered
+-/
+#guard_msgs in
+attribute [attribute] myCustomAttr'
+
+/-!
+Tagging a declaration that isn't an `AttributeImpl` should be rejected.
+-/
+
+def notAnAttribute : Nat := 0
+
+/--
+error: Unexpected attribute implementation type: `{.ofConstName declName}` is not of type `Lean.AttributeImpl`
+-/
+#guard_msgs in
+attribute [attribute] notAnAttribute
+
+/-!
+Local/scoped registration should be rejected.
+-/
+
+def anotherAttr : AttributeImpl where
+  name := `anotherAttr
+  descr := "another"
+  add := fun _ _ _ => pure ()
+
+/--
+error: Invalid attribute scope: Attribute `[attribute]` must be global, not `local`
+-/
+#guard_msgs in
+attribute [local attribute] anotherAttr

tests/pkg/attribute_attr/AttrTest.lean

@@ -0,0 +1 @@
+import AttrTest.Apply

tests/pkg/attribute_attr/AttrTest/Apply.lean

@@ -0,0 +1,19 @@
+import AttrTest.Ext
+
+/-! The attribute is defined here, alongside its applications. -/
+
+open Lean
+
+@[attribute] def markAttr : AttributeImpl where
+  name := `mark
+  descr := "mark a declaration into markedExt"
+  add := fun decl _stx _kind => do
+    modifyEnv fun env => markedExt.addEntry env decl
+
+@[mark] def foo := 1
+@[mark] def bar := 2
+@[mark] def baz := 3
+
+/-- info: [foo, bar, baz] -/
+#guard_msgs in
+run_meta Lean.logInfo m!"{getMarked (← Lean.getEnv)}"

tests/pkg/attribute_attr/AttrTest/Ext.lean

@@ -0,0 +1,14 @@
+import Lean
+
+/-! Defines the env extension only.  No attribute here. -/
+
+open Lean
+
+initialize markedExt : SimplePersistentEnvExtension Name (Array Name) ←
+  registerSimplePersistentEnvExtension {
+    addEntryFn    := Array.push
+    addImportedFn := fun ess => ess.foldl (init := #[]) Array.append
+  }
+
+def getMarked (env : Environment) : Array Name :=
+  markedExt.getState env

tests/pkg/attribute_attr/lakefile.lean

@@ -0,0 +1,5 @@
+import Lake
+open System Lake DSL
+
+package attribute_attr
+@[default_target] lean_lib AttrTest

tests/pkg/attribute_attr/lean-toolchain

@@ -0,0 +1 @@
+leanprover/lean4:nightly

tests/pkg/attribute_attr/run_test.sh

@@ -0,0 +1,2 @@
+rm -rf .lake/build
+lake build