lean4-htt/src/Lean/LabelAttribute.lean
2025-07-25 12:02:51 +00:00

100 lines
3.6 KiB
Text

/-
Copyright (c) 2023 Kim Morrison. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Kim Morrison
-/
module
prelude
public import Lean.ScopedEnvExtension
public import Lean.DocString
public meta import Init.Data.String.Extra
public section
/-!
# "Label" attributes
Allow creating attributes using `register_label_attr`,
and retrieving the array of `Name`s of declarations which have been tagged with such an attribute.
These differ slightly from the built-in "tag attributes" which can be initialized with the syntax:
```
initialize someName : TagAttribute ← registerTagAttribute `tagName "description"
```
in that a "tag attribute" can only be put on a declaration at the moment it is declared,
and can not be modified by scoping commands.
The "label attributes" constructed here support adding (or locally removing) the attribute
either at the moment of declaration, or later.
-/
namespace Lean
/-- An environment extension that just tracks an array of names. -/
abbrev LabelExtension := SimpleScopedEnvExtension Name (Array Name)
/-- The collection of all current `LabelExtension`s, indexed by name. -/
abbrev LabelExtensionMap := Std.HashMap Name LabelExtension
/-- Store the current `LabelExtension`s. -/
builtin_initialize labelExtensionMapRef : IO.Ref LabelExtensionMap ← IO.mkRef {}
/-- Helper function for `registerLabelAttr`. -/
def mkLabelExt (name : Name := by exact decl_name%) : IO LabelExtension :=
registerSimpleScopedEnvExtension {
name := name
initial := #[]
addEntry := fun d e => if d.contains e then d else d.push e
}
/-- Helper function for `registerLabelAttr`. -/
def mkLabelAttr (attrName : Name) (attrDescr : String) (ext : LabelExtension)
(ref : Name := by exact decl_name%) : IO Unit :=
registerBuiltinAttribute {
ref := ref
name := attrName
descr := attrDescr
applicationTime := AttributeApplicationTime.afterCompilation
add := fun declName _ kind =>
ext.add declName kind
erase := fun declName => do
let s := ext.getState (← getEnv)
modifyEnv fun env => ext.modifyState env fun _ => s.erase declName
}
/--
Construct a new "label attribute",
which does nothing except keep track of the names of the declarations with that attribute.
Users will generally use the `register_label_attr` macro defined below.
-/
def registerLabelAttr (attrName : Name) (attrDescr : String)
(ref : Name := by exact decl_name%) : IO LabelExtension := do
let ext ← mkLabelExt ref
mkLabelAttr attrName attrDescr ext ref
labelExtensionMapRef.modify fun map => map.insert attrName ext
return ext
/--
Initialize a new "label" attribute.
Declarations tagged with the attribute can be retrieved using `Lean.labelled`.
-/
macro (name := _root_.Lean.Parser.Command.registerLabelAttr)
doc:(docComment)? "register_label_attr " id:ident : command => do
let str := id.getId.toString
let idParser := mkIdentFrom id (`Parser.Attr ++ id.getId)
let descr := quote ((doc.map (·.getDocString) |>.getD ("labelled declarations for " ++ id.getId.toString)).removeLeadingSpaces)
`($[$doc:docComment]? initialize ext : Lean.LabelExtension ←
registerLabelAttr $(quote id.getId) $descr $(quote id.getId)
$[$doc:docComment]? syntax (name := $idParser:ident) $(quote str):str : attr)
/-- When `attrName` is an attribute created using `register_labelled_attr`,
return the names of all declarations labelled using that attribute. -/
def labelled (attrName : Name) : CoreM (Array Name) := do
match (← labelExtensionMapRef.get)[attrName]? with
| none => throwError "No extension named {attrName}"
| some ext => pure <| ext.getState (← getEnv)
end Lean