84 lines
3.4 KiB
Text
84 lines
3.4 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
|
|
-/
|
|
module
|
|
|
|
prelude
|
|
public import Lean.Meta.Basic
|
|
|
|
public section
|
|
|
|
namespace Lean.Compiler
|
|
|
|
inductive SpecializeAttributeKind where
|
|
| specialize | nospecialize
|
|
deriving Inhabited, BEq
|
|
|
|
/--
|
|
Marks a definition to never be specialized during code generation.
|
|
-/
|
|
@[builtin_doc]
|
|
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 `0`: Index must be greater than 0"
|
|
let idx := idx - 1
|
|
if h : idx >= argNames.size then
|
|
throwErrorAt arg "Invalid argument index `{idx}`: `{.ofConstName declName}` has only {argNames.size} arguments"
|
|
else
|
|
if result.contains idx then throwErrorAt arg "Invalid specialization argument index `{idx + 1}`: \
|
|
The argument at this 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.idxOf? 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}`: `{.ofConstName declName}` does not have an argument with this name"
|
|
return result.qsort (·<·)
|
|
|
|
/--
|
|
Marks a definition to always be specialized during code generation.
|
|
|
|
Specialization is an optimization in the code generator for generating variants of a function that
|
|
are specialized to specific parameter values. This is in particular useful for functions that take
|
|
other functions as parameters: Usually when passing functions as parameters, a closure needs to be
|
|
allocated that will then be called. Using `@[specialize]` prevents both of these operations by
|
|
using the provided function directly in the specialization of the inner function.
|
|
|
|
`@[specialize]` can take additional arguments for the parameter names or indices (starting at 1) of
|
|
the parameters that should be specialized. By default, instance and function parameters are
|
|
specialized.
|
|
-/
|
|
@[builtin_doc]
|
|
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
|
|
|
|
end Lean.Compiler
|