130 lines
5 KiB
Text
130 lines
5 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
|
|
-/
|
|
prelude
|
|
import Lean.Meta.Basic
|
|
import Lean.Attributes
|
|
|
|
namespace Lean.Compiler
|
|
|
|
inductive SpecializeAttributeKind where
|
|
| specialize | nospecialize
|
|
deriving Inhabited, BEq
|
|
|
|
builtin_initialize nospecializeAttr : TagAttribute ←
|
|
registerTagAttribute `nospecialize "mark definition to never be specialized"
|
|
|
|
private def elabSpecArgs (declName : Name) (args : Array Syntax) : MetaM (Array Nat) := do
|
|
if args.isEmpty then return #[]
|
|
let info ← getConstInfo declName
|
|
Meta.forallTelescopeReducing info.type fun xs _ => do
|
|
let argNames ← xs.mapM fun x => x.fvarId!.getUserName
|
|
let mut result := #[]
|
|
for arg in args do
|
|
if let some idx := arg.isNatLit? then
|
|
if idx == 0 then throwErrorAt arg "invalid specialization argument index, index must be greater than 0"
|
|
let idx := idx - 1
|
|
if idx >= argNames.size then
|
|
throwErrorAt arg "invalid argument index, `{declName}` has #{argNames.size} arguments"
|
|
if result.contains idx then throwErrorAt arg "invalid specialization argument index, `{argNames[idx]!}` has already been specified as a specialization candidate"
|
|
result := result.push idx
|
|
else
|
|
let argName := arg.getId
|
|
if let some idx := argNames.getIdx? argName then
|
|
if result.contains idx then throwErrorAt arg "invalid specialization argument name `{argName}`, it has already been specified as a specialization candidate"
|
|
result := result.push idx
|
|
else
|
|
throwErrorAt arg "invalid specialization argument name `{argName}`, `{declName}` does have an argument with this name"
|
|
return result.qsort (·<·)
|
|
|
|
builtin_initialize specializeAttr : ParametricAttribute (Array Nat) ←
|
|
registerParametricAttribute {
|
|
name := `specialize
|
|
descr := "mark definition to always be specialized"
|
|
getParam := fun declName stx => do
|
|
let args := stx[1].getArgs
|
|
elabSpecArgs declName args |>.run'
|
|
}
|
|
|
|
def getSpecializationArgs? (env : Environment) (declName : Name) : Option (Array Nat) :=
|
|
specializeAttr.getParam? env declName
|
|
|
|
def hasSpecializeAttribute (env : Environment) (declName : Name) : Bool :=
|
|
getSpecializationArgs? env declName |>.isSome
|
|
|
|
def hasNospecializeAttribute (env : Environment) (declName : Name) : Bool :=
|
|
nospecializeAttr.hasTag env declName
|
|
|
|
/- TODO: the rest of the file is for the old / current code generator. We should remove it as soon as we move to the new one. -/
|
|
|
|
@[export lean_has_specialize_attribute]
|
|
partial def hasSpecializeAttributeOld (env : Environment) (n : Name) : Bool :=
|
|
match specializeAttr.getParam? env n with
|
|
| some _ => true
|
|
| none => if n.isInternal then hasSpecializeAttributeOld env n.getPrefix else false -- TODO: remove recursion after we move to new compiler
|
|
|
|
@[export lean_has_nospecialize_attribute]
|
|
partial def hasNospecializeAttributeOld (env : Environment) (n : Name) : Bool :=
|
|
nospecializeAttr.hasTag env n ||
|
|
(n.isInternal && hasNospecializeAttributeOld env n.getPrefix) -- TODO: remove recursion after we move to new compiler
|
|
|
|
inductive SpecArgKind where
|
|
| fixed
|
|
| fixedNeutral -- computationally neutral
|
|
| fixedHO -- higher order
|
|
| fixedInst -- type class instance
|
|
| other
|
|
deriving Inhabited
|
|
|
|
structure SpecInfo where
|
|
mutualDecls : List Name
|
|
argKinds : List SpecArgKind
|
|
deriving Inhabited
|
|
|
|
structure SpecState where
|
|
specInfo : SMap Name SpecInfo := {}
|
|
cache : SMap Expr Name := {}
|
|
deriving Inhabited
|
|
|
|
inductive SpecEntry where
|
|
| info (name : Name) (info : SpecInfo)
|
|
| cache (key : Expr) (fn : Name)
|
|
deriving Inhabited
|
|
|
|
namespace SpecState
|
|
|
|
def addEntry (s : SpecState) (e : SpecEntry) : SpecState :=
|
|
match e with
|
|
| SpecEntry.info name info => { s with specInfo := s.specInfo.insert name info }
|
|
| SpecEntry.cache key fn => { s with cache := s.cache.insert key fn }
|
|
|
|
def switch : SpecState → SpecState
|
|
| ⟨m₁, m₂⟩ => ⟨m₁.switch, m₂.switch⟩
|
|
|
|
end SpecState
|
|
|
|
builtin_initialize specExtension : SimplePersistentEnvExtension SpecEntry SpecState ←
|
|
registerSimplePersistentEnvExtension {
|
|
addEntryFn := SpecState.addEntry,
|
|
addImportedFn := fun es => (mkStateFromImportedEntries SpecState.addEntry {} es).switch
|
|
}
|
|
|
|
@[export lean_add_specialization_info]
|
|
def addSpecializationInfo (env : Environment) (fn : Name) (info : SpecInfo) : Environment :=
|
|
specExtension.addEntry env (SpecEntry.info fn info)
|
|
|
|
@[export lean_get_specialization_info]
|
|
def getSpecializationInfo (env : Environment) (fn : Name) : Option SpecInfo :=
|
|
(specExtension.getState env).specInfo.find? fn
|
|
|
|
@[export lean_cache_specialization]
|
|
def cacheSpecialization (env : Environment) (e : Expr) (fn : Name) : Environment :=
|
|
specExtension.addEntry env (SpecEntry.cache e fn)
|
|
|
|
@[export lean_get_cached_specialization]
|
|
def getCachedSpecialization (env : Environment) (e : Expr) : Option Name :=
|
|
(specExtension.getState env).cache.find? e
|
|
|
|
end Lean.Compiler
|