From 944beb69fda300da12a4611f35ceeea23d614b2d Mon Sep 17 00:00:00 2001 From: Sebastian Ullrich Date: Sun, 3 May 2026 21:44:16 +0000 Subject: [PATCH] feat: add `@[attribute]` for registering custom attributes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR adds a new `@[attribute]` attribute that registers a constant of type `Lean.AttributeImpl` as a custom attribute, lifting the previous restriction that all attributes had to be registered via `registerBuiltinAttribute` from inside an `initialize` block. A new `attributeExt` environment extension stores the registered attributes, and `getAttributeImpl` now consults both it and the existing `attributeExtension` (which keeps serving builtin attributes and parser category builders). Because `@[attribute]` registers eagerly during elaboration (via `applicationTime := .afterCompilation`) rather than at module load, custom attributes can be applied in modules that import the defining module — including in the same file as the orphan applications they target — without the artificial split that `initialize`-based registration forces. The accompanying parser change registers `attribute` as a syntax kind in the `attr` category alongside the other keyword-named attribute forms (`class`, `instance`, `macro`, `export`). Co-Authored-By: Claude Opus 4.7 (1M context) --- src/Lean/Attributes.lean | 81 ++++++++++++++++++-- src/Lean/Parser/Attr.lean | 3 +- tests/elab/attributeAttr.lean | 74 ++++++++++++++++++ tests/pkg/attribute_attr/AttrTest.lean | 1 + tests/pkg/attribute_attr/AttrTest/Apply.lean | 19 +++++ tests/pkg/attribute_attr/AttrTest/Ext.lean | 14 ++++ tests/pkg/attribute_attr/lakefile.lean | 5 ++ tests/pkg/attribute_attr/lean-toolchain | 1 + tests/pkg/attribute_attr/run_test.sh | 2 + 9 files changed, 192 insertions(+), 8 deletions(-) create mode 100644 tests/elab/attributeAttr.lean create mode 100644 tests/pkg/attribute_attr/AttrTest.lean create mode 100644 tests/pkg/attribute_attr/AttrTest/Apply.lean create mode 100644 tests/pkg/attribute_attr/AttrTest/Ext.lean create mode 100644 tests/pkg/attribute_attr/lakefile.lean create mode 100644 tests/pkg/attribute_attr/lean-toolchain create mode 100644 tests/pkg/attribute_attr/run_test.sh diff --git a/src/Lean/Attributes.lean b/src/Lean/Attributes.lean index dc523132fa35..dc59891a9f6e 100644 --- a/src/Lean/Attributes.lean +++ b/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 diff --git a/src/Lean/Parser/Attr.lean b/src/Lean/Parser/Attr.lean index 5b170c3d4d01..6da74e5fb820 100644 --- a/src/Lean/Parser/Attr.lean +++ b/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 diff --git a/tests/elab/attributeAttr.lean b/tests/elab/attributeAttr.lean new file mode 100644 index 000000000000..e23af8690753 --- /dev/null +++ b/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 diff --git a/tests/pkg/attribute_attr/AttrTest.lean b/tests/pkg/attribute_attr/AttrTest.lean new file mode 100644 index 000000000000..e7c168d177b1 --- /dev/null +++ b/tests/pkg/attribute_attr/AttrTest.lean @@ -0,0 +1 @@ +import AttrTest.Apply diff --git a/tests/pkg/attribute_attr/AttrTest/Apply.lean b/tests/pkg/attribute_attr/AttrTest/Apply.lean new file mode 100644 index 000000000000..f2b890eda336 --- /dev/null +++ b/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)}" diff --git a/tests/pkg/attribute_attr/AttrTest/Ext.lean b/tests/pkg/attribute_attr/AttrTest/Ext.lean new file mode 100644 index 000000000000..c019d46f8630 --- /dev/null +++ b/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 diff --git a/tests/pkg/attribute_attr/lakefile.lean b/tests/pkg/attribute_attr/lakefile.lean new file mode 100644 index 000000000000..54fcf1fbd659 --- /dev/null +++ b/tests/pkg/attribute_attr/lakefile.lean @@ -0,0 +1,5 @@ +import Lake +open System Lake DSL + +package attribute_attr +@[default_target] lean_lib AttrTest diff --git a/tests/pkg/attribute_attr/lean-toolchain b/tests/pkg/attribute_attr/lean-toolchain new file mode 100644 index 000000000000..6a9c6db38965 --- /dev/null +++ b/tests/pkg/attribute_attr/lean-toolchain @@ -0,0 +1 @@ +leanprover/lean4:nightly diff --git a/tests/pkg/attribute_attr/run_test.sh b/tests/pkg/attribute_attr/run_test.sh new file mode 100644 index 000000000000..b911bf4de8da --- /dev/null +++ b/tests/pkg/attribute_attr/run_test.sh @@ -0,0 +1,2 @@ +rm -rf .lake/build +lake build