177 lines
7.1 KiB
Text
177 lines
7.1 KiB
Text
/-
|
||
Copyright (c) 2020 Microsoft Corporation. All rights reserved.
|
||
Released under Apache 2.0 license as described in the file LICENSE.
|
||
Authors: Leonardo de Moura, Sebastian Ullrich
|
||
-/
|
||
prelude
|
||
import Lean.Compiler.InitAttr
|
||
import Lean.ScopedEnvExtension
|
||
import Lean.Compiler.IR.CompilerM
|
||
|
||
/-!
|
||
A builder for attributes that are applied to declarations of a common type and
|
||
group them by the given attribute argument (an arbitrary `Name`, currently).
|
||
Also creates a second "builtin" attribute used for bootstrapping, which saves
|
||
the applied declarations in an `IO.Ref` instead of an environment extension.
|
||
|
||
Used to register elaborators, macros, tactics, and delaborators.
|
||
-/
|
||
|
||
namespace Lean
|
||
namespace KeyedDeclsAttribute
|
||
|
||
-- could be a parameter as well, but right now it's all names
|
||
abbrev Key := Name
|
||
|
||
/--
|
||
`KeyedDeclsAttribute` definition.
|
||
|
||
Important: `mkConst valueTypeName` and `γ` must be definitionally equal. -/
|
||
structure Def (γ : Type) where
|
||
/-- Builtin attribute name, if any (e.g., `builtin_term_elab) -/
|
||
builtinName : Name := Name.anonymous
|
||
/-- Attribute name (e.g., `term_elab) -/
|
||
name : Name
|
||
/-- Attribute description -/
|
||
descr : String
|
||
valueTypeName : Name
|
||
/-- Convert `Syntax` into a `Key`, the default implementation expects an identifier. -/
|
||
evalKey (builtin : Bool) (stx : Syntax) : AttrM Key := do
|
||
let stx ← Attribute.Builtin.getIdent stx
|
||
let kind := stx.getId
|
||
if (← getEnv).contains kind && (← Elab.getInfoState).enabled then
|
||
Elab.addConstInfo stx kind none
|
||
pure kind
|
||
onAdded (builtin : Bool) (declName : Name) : AttrM Unit := pure ()
|
||
deriving Inhabited
|
||
|
||
structure OLeanEntry where
|
||
key : Key
|
||
/-- Name of a declaration stored in the environment which has type `mkConst Def.valueTypeName`. -/
|
||
declName : Name
|
||
deriving Inhabited
|
||
|
||
structure AttributeEntry (γ : Type) extends OLeanEntry where
|
||
/-- Recall that we cannot store `γ` into .olean files because it is a closure.
|
||
Given `OLeanEntry.declName`, we convert it into a `γ` by using the unsafe function `evalConstCheck`. -/
|
||
value : γ
|
||
|
||
abbrev Table (γ : Type) := SMap Key (List (AttributeEntry γ))
|
||
|
||
structure ExtensionState (γ : Type) where
|
||
newEntries : List OLeanEntry := []
|
||
table : Table γ := {}
|
||
declNames : PHashSet Name := {}
|
||
erased : PHashSet Name := {}
|
||
deriving Inhabited
|
||
|
||
abbrev Extension (γ : Type) := ScopedEnvExtension OLeanEntry (AttributeEntry γ) (ExtensionState γ)
|
||
|
||
end KeyedDeclsAttribute
|
||
|
||
structure KeyedDeclsAttribute (γ : Type) where
|
||
defn : KeyedDeclsAttribute.Def γ
|
||
-- imported/builtin instances
|
||
tableRef : IO.Ref (KeyedDeclsAttribute.Table γ)
|
||
-- instances from current module
|
||
ext : KeyedDeclsAttribute.Extension γ
|
||
deriving Nonempty
|
||
|
||
namespace KeyedDeclsAttribute
|
||
|
||
private def Table.insert (table : Table γ) (v : AttributeEntry γ) : Table γ :=
|
||
match table.find? v.key with
|
||
| some vs => SMap.insert table v.key (v::vs)
|
||
| none => SMap.insert table v.key [v]
|
||
|
||
def ExtensionState.insert (s : ExtensionState γ) (v : AttributeEntry γ) : ExtensionState γ := {
|
||
table := s.table.insert v
|
||
newEntries := v.toOLeanEntry :: s.newEntries
|
||
declNames := s.declNames.insert v.declName
|
||
erased := s.erased.erase v.declName
|
||
}
|
||
|
||
def addBuiltin (attr : KeyedDeclsAttribute γ) (key : Key) (declName : Name) (value : γ) : IO Unit :=
|
||
attr.tableRef.modify fun m => m.insert { key, declName, value }
|
||
|
||
def mkStateOfTable (table : Table γ) : ExtensionState γ := {
|
||
table
|
||
declNames := table.fold (init := {}) fun s _ es => es.foldl (init := s) fun s e => s.insert e.declName
|
||
}
|
||
|
||
def ExtensionState.erase (s : ExtensionState γ) (attrName : Name) (declName : Name) : CoreM (ExtensionState γ) := do
|
||
unless s.declNames.contains declName do
|
||
throwError "'{declName}' does not have [{attrName}] attribute"
|
||
return { s with erased := s.erased.insert declName, declNames := s.declNames.erase declName }
|
||
|
||
protected unsafe def init {γ} (df : Def γ) (attrDeclName : Name := by exact decl_name%) : IO (KeyedDeclsAttribute γ) := do
|
||
let tableRef ← IO.mkRef ({} : Table γ)
|
||
let ext : Extension γ ← registerScopedEnvExtension {
|
||
name := attrDeclName
|
||
mkInitial := return mkStateOfTable (← tableRef.get)
|
||
ofOLeanEntry := fun _ entry => do
|
||
let ctx ← read
|
||
match ctx.env.evalConstCheck γ ctx.opts df.valueTypeName entry.declName with
|
||
| Except.ok f => return { toOLeanEntry := entry, value := f }
|
||
| Except.error ex => throw (IO.userError ex)
|
||
addEntry := fun s e => s.insert e
|
||
toOLeanEntry := (·.toOLeanEntry)
|
||
}
|
||
unless df.builtinName.isAnonymous do
|
||
registerBuiltinAttribute {
|
||
ref := attrDeclName
|
||
name := df.builtinName
|
||
descr := "(builtin) " ++ df.descr
|
||
add := fun declName stx kind => do
|
||
unless kind == AttributeKind.global do throwError "invalid attribute '{df.builtinName}', must be global"
|
||
let key ← df.evalKey true stx
|
||
let decl ← getConstInfo declName
|
||
match decl.type with
|
||
| Expr.const c _ =>
|
||
if c != df.valueTypeName then throwError "unexpected type at '{declName}', '{df.valueTypeName}' expected"
|
||
else
|
||
/- builtin_initialize @addBuiltin $(mkConst valueTypeName) $(mkConst attrDeclName) $(key) $(declName) $(mkConst declName) -/
|
||
let val := mkAppN (mkConst ``addBuiltin) #[mkConst df.valueTypeName, mkConst attrDeclName, toExpr key, toExpr declName, mkConst declName]
|
||
declareBuiltin declName val
|
||
df.onAdded true declName
|
||
| _ => throwError "unexpected type at '{declName}', '{df.valueTypeName}' expected"
|
||
applicationTime := AttributeApplicationTime.afterCompilation
|
||
}
|
||
registerBuiltinAttribute {
|
||
ref := attrDeclName
|
||
name := df.name
|
||
descr := df.descr
|
||
erase := fun declName => do
|
||
let s := ext.getState (← getEnv)
|
||
let s ← s.erase df.name declName
|
||
modifyEnv fun env => ext.modifyState env fun _ => s
|
||
add := fun declName stx attrKind => do
|
||
let key ← df.evalKey false stx
|
||
match IR.getSorryDep (← getEnv) declName with
|
||
| none =>
|
||
let val ← evalConstCheck γ df.valueTypeName declName
|
||
ext.add { key := key, declName := declName, value := val } attrKind
|
||
df.onAdded false declName
|
||
| _ =>
|
||
-- If the declaration contains `sorry`, we skip `evalConstCheck` to avoid unnecessary bizarre error message
|
||
pure ()
|
||
applicationTime := AttributeApplicationTime.afterCompilation
|
||
}
|
||
pure { defn := df, tableRef := tableRef, ext := ext }
|
||
|
||
/-- Retrieve entries tagged with `[attr key]` or `[builtinAttr key]`. -/
|
||
def getEntries {γ} (attr : KeyedDeclsAttribute γ) (env : Environment) (key : Name) : List (AttributeEntry γ) :=
|
||
let s := attr.ext.getState env
|
||
let attrs := s.table.findD key []
|
||
if s.erased.isEmpty then
|
||
attrs
|
||
else
|
||
attrs.filter fun attr => !s.erased.contains attr.declName
|
||
|
||
/-- Retrieve values tagged with `[attr key]` or `[builtinAttr key]`. -/
|
||
def getValues {γ} (attr : KeyedDeclsAttribute γ) (env : Environment) (key : Name) : List γ :=
|
||
(getEntries attr env key).map AttributeEntry.value
|
||
|
||
end KeyedDeclsAttribute
|
||
|
||
end Lean
|