373 lines
17 KiB
Text
373 lines
17 KiB
Text
/-
|
||
Copyright (c) 2019 Microsoft Corporation. All rights reserved.
|
||
Released under Apache 2.0 license as described in the file LICENSE.
|
||
Authors: Leonardo de Moura
|
||
-/
|
||
import Lean.Syntax
|
||
import Lean.CoreM
|
||
import Lean.ResolveName
|
||
|
||
namespace Lean
|
||
|
||
inductive AttributeApplicationTime :=
|
||
| afterTypeChecking | afterCompilation | beforeElaboration
|
||
|
||
def AttributeApplicationTime.beq : AttributeApplicationTime → AttributeApplicationTime → Bool
|
||
| AttributeApplicationTime.afterTypeChecking, AttributeApplicationTime.afterTypeChecking => true
|
||
| AttributeApplicationTime.afterCompilation, AttributeApplicationTime.afterCompilation => true
|
||
| AttributeApplicationTime.beforeElaboration, AttributeApplicationTime.beforeElaboration => true
|
||
| _, _ => false
|
||
|
||
instance : BEq AttributeApplicationTime := ⟨AttributeApplicationTime.beq⟩
|
||
|
||
structure Attr.Context :=
|
||
(currNamespace : Name)
|
||
(openDecls : List OpenDecl)
|
||
|
||
abbrev AttrM := ReaderT Attr.Context CoreM
|
||
|
||
instance : MonadLift ImportM AttrM := {
|
||
monadLift := fun x => do liftM (m := IO) (x { env := (← getEnv), opts := (← getOptions) })
|
||
}
|
||
|
||
instance : MonadResolveName AttrM := {
|
||
getCurrNamespace := do pure (← read).currNamespace,
|
||
getOpenDecls := do pure (← read).openDecls
|
||
}
|
||
|
||
structure AttributeImplCore :=
|
||
(name : Name)
|
||
(descr : String)
|
||
(applicationTime := AttributeApplicationTime.afterTypeChecking)
|
||
|
||
structure AttributeImpl extends AttributeImplCore :=
|
||
(add (decl : Name) (args : Syntax) (persistent : Bool) : AttrM Unit)
|
||
|
||
instance : Inhabited AttributeImpl :=
|
||
⟨{ name := arbitrary _, descr := arbitrary _, add := fun env _ _ _ => pure () }⟩
|
||
|
||
open Std (PersistentHashMap)
|
||
|
||
builtin_initialize attributeMapRef : IO.Ref (PersistentHashMap Name AttributeImpl) ← IO.mkRef {}
|
||
|
||
/- Low level attribute registration function. -/
|
||
def registerBuiltinAttribute (attr : AttributeImpl) : IO Unit := do
|
||
let m ← attributeMapRef.get
|
||
if m.contains attr.name then throw (IO.userError ("invalid builtin attribute declaration, '" ++ toString attr.name ++ "' has already been used"))
|
||
unless (← IO.initializing) do throw (IO.userError "failed to register attribute, attributes can only be registered during initialization")
|
||
attributeMapRef.modify fun m => m.insert attr.name attr
|
||
|
||
abbrev AttributeImplBuilder := List DataValue → Except String AttributeImpl
|
||
abbrev AttributeImplBuilderTable := Std.HashMap Name AttributeImplBuilder
|
||
|
||
builtin_initialize attributeImplBuilderTableRef : IO.Ref AttributeImplBuilderTable ← IO.mkRef {}
|
||
|
||
def registerAttributeImplBuilder (builderId : Name) (builder : AttributeImplBuilder) : IO Unit := do
|
||
let table ← attributeImplBuilderTableRef.get
|
||
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
|
||
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
|
||
|
||
inductive AttributeExtensionOLeanEntry :=
|
||
| decl (declName : Name) -- `declName` has type `AttributeImpl`
|
||
| builder (builderId : Name) (args : List DataValue)
|
||
|
||
structure AttributeExtensionState :=
|
||
(newEntries : List AttributeExtensionOLeanEntry := [])
|
||
(map : PersistentHashMap Name AttributeImpl)
|
||
|
||
abbrev AttributeExtension := PersistentEnvExtension AttributeExtensionOLeanEntry (AttributeExtensionOLeanEntry × AttributeImpl) AttributeExtensionState
|
||
|
||
instance : Inhabited AttributeExtensionState := ⟨{ map := {} }⟩
|
||
|
||
private def AttributeExtension.mkInitial : IO AttributeExtensionState := do
|
||
let map ← attributeMapRef.get
|
||
pure { map := map }
|
||
|
||
unsafe def mkAttributeImplOfConstantUnsafe (env : Environment) (opts : Options) (declName : Name) : Except String AttributeImpl :=
|
||
match env.find? declName with
|
||
| none => throw ("unknow constant '" ++ toString declName ++ "'")
|
||
| some info =>
|
||
match info.type with
|
||
| Expr.const `Lean.AttributeImpl _ _ => env.evalConst AttributeImpl opts declName
|
||
| _ => throw ("unexpected attribute implementation type at '" ++ toString declName ++ "' (`AttributeImpl` expected")
|
||
|
||
@[implementedBy mkAttributeImplOfConstantUnsafe]
|
||
constant mkAttributeImplOfConstant (env : Environment) (opts : Options) (declName : Name) : Except String AttributeImpl
|
||
|
||
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
|
||
|
||
private def AttributeExtension.addImported (es : Array (Array AttributeExtensionOLeanEntry)) : ImportM AttributeExtensionState := do
|
||
let ctx ← read
|
||
let map ← attributeMapRef.get
|
||
let map ← es.foldlM
|
||
(fun map entries =>
|
||
entries.foldlM
|
||
(fun (map : PersistentHashMap Name AttributeImpl) entry => do
|
||
let attrImpl ← liftM $ mkAttributeImplOfEntry ctx.env ctx.opts entry
|
||
pure $ map.insert attrImpl.name attrImpl)
|
||
map)
|
||
map
|
||
pure { map := map }
|
||
|
||
private def addAttrEntry (s : AttributeExtensionState) (e : AttributeExtensionOLeanEntry × AttributeImpl) : AttributeExtensionState :=
|
||
{ s with map := s.map.insert e.2.name e.2, newEntries := e.1 :: s.newEntries }
|
||
|
||
builtin_initialize attributeExtension : AttributeExtension ←
|
||
registerPersistentEnvExtension {
|
||
name := `attrExt,
|
||
mkInitial := AttributeExtension.mkInitial,
|
||
addImportedFn := AttributeExtension.addImported,
|
||
addEntryFn := addAttrEntry,
|
||
exportEntriesFn := fun s => s.newEntries.reverse.toArray,
|
||
statsFn := fun s => format "number of local entries: " ++ format s.newEntries.length
|
||
}
|
||
|
||
/- Return true iff `n` is the name of a registered attribute. -/
|
||
@[export lean_is_attribute]
|
||
def isBuiltinAttribute (n : Name) : IO Bool := do
|
||
let m ← attributeMapRef.get; pure (m.contains n)
|
||
|
||
/- Return the name of all registered attributes. -/
|
||
def getBuiltinAttributeNames : IO (List Name) := do
|
||
let m ← attributeMapRef.get; pure $ m.foldl (fun r n _ => n::r) []
|
||
|
||
def getBuiltinAttributeImpl (attrName : Name) : IO AttributeImpl := do
|
||
let m ← attributeMapRef.get
|
||
match m.find? attrName with
|
||
| some attr => pure attr
|
||
| none => throw (IO.userError ("unknown attribute '" ++ toString attrName ++ "'"))
|
||
|
||
@[export lean_attribute_application_time]
|
||
def getBuiltinAttributeApplicationTime (n : Name) : IO AttributeApplicationTime := do
|
||
let attr ← getBuiltinAttributeImpl n
|
||
pure attr.applicationTime
|
||
|
||
def isAttribute (env : Environment) (attrName : Name) : Bool :=
|
||
(attributeExtension.getState env).map.contains attrName
|
||
|
||
def getAttributeNames (env : Environment) : List Name :=
|
||
let m := (attributeExtension.getState env).map
|
||
m.foldl (fun r n _ => n::r) []
|
||
|
||
def getAttributeImpl (env : Environment) (attrName : Name) : Except String AttributeImpl :=
|
||
let m := (attributeExtension.getState env).map
|
||
match m.find? attrName with
|
||
| some attr => pure attr
|
||
| none => throw ("unknown attribute '" ++ toString attrName ++ "'")
|
||
|
||
def registerAttributeOfDecl (env : Environment) (opts : Options) (attrDeclName : Name) : Except String Environment := do
|
||
let attrImpl ← mkAttributeImplOfConstant env opts attrDeclName
|
||
if isAttribute env attrImpl.name then
|
||
throw ("invalid builtin attribute declaration, '" ++ toString attrImpl.name ++ "' has already been used")
|
||
else
|
||
pure $ attributeExtension.addEntry env (AttributeExtensionOLeanEntry.decl attrDeclName, attrImpl)
|
||
|
||
def registerAttributeOfBuilder (env : Environment) (builderId : Name) (args : List DataValue) : IO Environment := do
|
||
let attrImpl ← mkAttributeImplOfBuilder builderId args
|
||
if isAttribute env attrImpl.name then
|
||
throw (IO.userError ("invalid builtin attribute declaration, '" ++ toString attrImpl.name ++ "' has already been used"))
|
||
else
|
||
pure $ attributeExtension.addEntry env (AttributeExtensionOLeanEntry.builder builderId args, attrImpl)
|
||
|
||
def addAttribute (decl : Name) (attrName : Name) (args : Syntax) (persistent : Bool := true) : AttrM Unit := do
|
||
let env ← getEnv
|
||
let attr ← ofExcept $ getAttributeImpl env attrName
|
||
attr.add decl args persistent
|
||
|
||
/--
|
||
Tag attributes are simple and efficient. They are useful for marking declarations in the modules where
|
||
they were defined.
|
||
|
||
The startup cost for this kind of attribute is very small since `addImportedFn` is a constant function.
|
||
|
||
They provide the predicate `tagAttr.hasTag env decl` which returns true iff declaration `decl`
|
||
is tagged in the environment `env`. -/
|
||
structure TagAttribute :=
|
||
(attr : AttributeImpl)
|
||
(ext : PersistentEnvExtension Name Name NameSet)
|
||
|
||
def registerTagAttribute (name : Name) (descr : String) (validate : Name → AttrM Unit := fun _ => pure ()) : IO TagAttribute := do
|
||
let ext : PersistentEnvExtension Name Name NameSet ← registerPersistentEnvExtension {
|
||
name := name,
|
||
mkInitial := pure {},
|
||
addImportedFn := fun _ _ => pure {},
|
||
addEntryFn := fun (s : NameSet) n => s.insert n,
|
||
exportEntriesFn := fun es =>
|
||
let r : Array Name := es.fold (fun a e => a.push e) #[]
|
||
r.qsort Name.quickLt,
|
||
statsFn := fun s => "tag attribute" ++ Format.line ++ "number of local entries: " ++ format s.size
|
||
}
|
||
let attrImpl : AttributeImpl := {
|
||
name := name,
|
||
descr := descr,
|
||
add := fun decl args persistent => do
|
||
if args.hasArgs then throwError! "invalid attribute '{name}', unexpected argument"
|
||
unless persistent do throwError! "invalid attribute '{name}', must be persistent"
|
||
let env ← getEnv
|
||
unless (env.getModuleIdxFor? decl).isNone do
|
||
throwError! "invalid attribute '{name}', declaration is in an imported module"
|
||
validate decl
|
||
let env ← getEnv
|
||
setEnv $ ext.addEntry env decl
|
||
}
|
||
registerBuiltinAttribute attrImpl
|
||
pure { attr := attrImpl, ext := ext }
|
||
|
||
namespace TagAttribute
|
||
|
||
instance : Inhabited TagAttribute := ⟨{attr := arbitrary _, ext := arbitrary _}⟩
|
||
|
||
def hasTag (attr : TagAttribute) (env : Environment) (decl : Name) : Bool :=
|
||
match env.getModuleIdxFor? decl with
|
||
| some modIdx => (attr.ext.getModuleEntries env modIdx).binSearchContains decl Name.quickLt
|
||
| none => (attr.ext.getState env).contains decl
|
||
|
||
end TagAttribute
|
||
|
||
/--
|
||
A `TagAttribute` variant where we can attach parameters to attributes.
|
||
It is slightly more expensive and consumes a little bit more memory than `TagAttribute`.
|
||
|
||
They provide the function `pAttr.getParam env decl` which returns `some p` iff declaration `decl`
|
||
contains the attribute `pAttr` with parameter `p`. -/
|
||
structure ParametricAttribute (α : Type) :=
|
||
(attr : AttributeImpl)
|
||
(ext : PersistentEnvExtension (Name × α) (Name × α) (NameMap α))
|
||
|
||
structure ParametricAttributeImpl (α : Type) extends AttributeImplCore :=
|
||
(getParam : Name → Syntax → AttrM α)
|
||
(afterSet : Name → α → AttrM Unit := fun env _ _ => pure ())
|
||
(afterImport : Array (Array (Name × α)) → ImportM Unit := fun _ => pure ())
|
||
|
||
def registerParametricAttribute {α : Type} [Inhabited α] (impl : ParametricAttributeImpl α) : IO (ParametricAttribute α) := do
|
||
let ext : PersistentEnvExtension (Name × α) (Name × α) (NameMap α) ← registerPersistentEnvExtension {
|
||
name := impl.name,
|
||
mkInitial := pure {},
|
||
addImportedFn := fun s => impl.afterImport s *> pure {},
|
||
addEntryFn := fun (s : NameMap α) (p : Name × α) => s.insert p.1 p.2,
|
||
exportEntriesFn := fun m =>
|
||
let r : Array (Name × α) := m.fold (fun a n p => a.push (n, p)) #[]
|
||
r.qsort (fun a b => Name.quickLt a.1 b.1),
|
||
statsFn := fun s => "parametric attribute" ++ Format.line ++ "number of local entries: " ++ format s.size
|
||
}
|
||
let attrImpl : AttributeImpl := { impl with
|
||
add := fun decl args persistent => do
|
||
unless persistent do throwError! "invalid attribute '{impl.name}', must be persistent"
|
||
let env ← getEnv
|
||
unless (env.getModuleIdxFor? decl).isNone do
|
||
throwError! "invalid attribute '{impl.name}', declaration is in an imported module"
|
||
let val ← impl.getParam decl args
|
||
let env' := ext.addEntry env (decl, val)
|
||
setEnv env'
|
||
try impl.afterSet decl val catch _ => setEnv env
|
||
}
|
||
registerBuiltinAttribute attrImpl
|
||
pure { attr := attrImpl, ext := ext }
|
||
|
||
namespace ParametricAttribute
|
||
|
||
instance {α : Type} : Inhabited (ParametricAttribute α) := ⟨{attr := arbitrary _, ext := arbitrary _}⟩
|
||
|
||
def getParam {α : Type} [Inhabited α] (attr : ParametricAttribute α) (env : Environment) (decl : Name) : Option α :=
|
||
match env.getModuleIdxFor? decl with
|
||
| some modIdx =>
|
||
match (attr.ext.getModuleEntries env modIdx).binSearch (decl, arbitrary _) (fun a b => Name.quickLt a.1 b.1) with
|
||
| some (_, val) => some val
|
||
| none => none
|
||
| none => (attr.ext.getState env).find? decl
|
||
|
||
def setParam {α : Type} (attr : ParametricAttribute α) (env : Environment) (decl : Name) (param : α) : Except String Environment :=
|
||
if (env.getModuleIdxFor? decl).isSome then
|
||
Except.error ("invalid '" ++ toString attr.attr.name ++ "'.setParam, declaration is in an imported module")
|
||
else if ((attr.ext.getState env).find? decl).isSome then
|
||
Except.error ("invalid '" ++ toString attr.attr.name ++ "'.setParam, attribute has already been set")
|
||
else
|
||
Except.ok (attr.ext.addEntry env (decl, param))
|
||
|
||
end ParametricAttribute
|
||
|
||
/-
|
||
Given a list `[a₁, ..., a_n]` of elements of type `α`, `EnumAttributes` provides an attribute `Attr_i` for
|
||
associating a value `a_i` with an declaration. `α` is usually an enumeration type.
|
||
Note that whenever we register an `EnumAttributes`, we create `n` attributes, but only one environment extension. -/
|
||
structure EnumAttributes (α : Type) :=
|
||
(attrs : List AttributeImpl)
|
||
(ext : PersistentEnvExtension (Name × α) (Name × α) (NameMap α))
|
||
|
||
def registerEnumAttributes {α : Type} [Inhabited α] (extName : Name) (attrDescrs : List (Name × String × α))
|
||
(validate : Name → α → AttrM Unit := fun _ _ => pure ())
|
||
(applicationTime := AttributeApplicationTime.afterTypeChecking) : IO (EnumAttributes α) := do
|
||
let ext : PersistentEnvExtension (Name × α) (Name × α) (NameMap α) ← registerPersistentEnvExtension {
|
||
name := extName,
|
||
mkInitial := pure {},
|
||
addImportedFn := fun _ _ => pure {},
|
||
addEntryFn := fun (s : NameMap α) (p : Name × α) => s.insert p.1 p.2,
|
||
exportEntriesFn := fun m =>
|
||
let r : Array (Name × α) := m.fold (fun a n p => a.push (n, p)) #[]
|
||
r.qsort (fun a b => Name.quickLt a.1 b.1),
|
||
statsFn := fun s => "enumeration attribute extension" ++ Format.line ++ "number of local entries: " ++ format s.size
|
||
}
|
||
let attrs := attrDescrs.map fun (name, descr, val) => {
|
||
name := name,
|
||
descr := descr,
|
||
add := fun decl args persistent => do
|
||
unless persistent do throwError! "invalid attribute '{name}', must be persistent"
|
||
let env ← getEnv
|
||
unless (env.getModuleIdxFor? decl).isNone do
|
||
throwError! "invalid attribute '{name}', declaration is in an imported module"
|
||
validate decl val
|
||
setEnv $ ext.addEntry env (decl, val),
|
||
applicationTime := applicationTime
|
||
: AttributeImpl
|
||
}
|
||
attrs.forM registerBuiltinAttribute
|
||
pure { ext := ext, attrs := attrs }
|
||
|
||
namespace EnumAttributes
|
||
|
||
instance {α : Type} : Inhabited (EnumAttributes α) := ⟨{attrs := [], ext := arbitrary _}⟩
|
||
|
||
def getValue {α : Type} [Inhabited α] (attr : EnumAttributes α) (env : Environment) (decl : Name) : Option α :=
|
||
match env.getModuleIdxFor? decl with
|
||
| some modIdx =>
|
||
match (attr.ext.getModuleEntries env modIdx).binSearch (decl, arbitrary _) (fun a b => Name.quickLt a.1 b.1) with
|
||
| some (_, val) => some val
|
||
| none => none
|
||
| none => (attr.ext.getState env).find? decl
|
||
|
||
def setValue {α : Type} (attrs : EnumAttributes α) (env : Environment) (decl : Name) (val : α) : Except String Environment :=
|
||
if (env.getModuleIdxFor? decl).isSome then
|
||
Except.error ("invalid '" ++ toString attrs.ext.name ++ "'.setValue, declaration is in an imported module")
|
||
else if ((attrs.ext.getState env).find? decl).isSome then
|
||
Except.error ("invalid '" ++ toString attrs.ext.name ++ "'.setValue, attribute has already been set")
|
||
else
|
||
Except.ok (attrs.ext.addEntry env (decl, val))
|
||
|
||
end EnumAttributes
|
||
|
||
/--
|
||
Helper function for converting a Syntax object representing attribute parameters into an identifier.
|
||
It returns `none` if the parameter is not a simple identifier.
|
||
|
||
Remark: in the future, attributes should define their own parsers, and we should use `match_syntax` to
|
||
decode the Syntax object. -/
|
||
def attrParamSyntaxToIdentifier (s : Syntax) : Option Name :=
|
||
match s with
|
||
| Syntax.node k args =>
|
||
if k == nullKind && args.size == 1 then match args.get! 0 with
|
||
| Syntax.ident _ _ id _ => some id
|
||
| _ => none
|
||
else
|
||
none
|
||
| _ => none
|
||
|
||
end Lean
|