3bfeb0bc1f
This PR uses `isRecursiveDefinition` when validating `macro_inline`, instead of rummaging in the internals of the definition.
106 lines
4.4 KiB
Text
106 lines
4.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.Attributes
|
|
import Lean.Meta.RecExt
|
|
|
|
public section
|
|
|
|
|
|
namespace Lean.Compiler
|
|
|
|
inductive InlineAttributeKind where
|
|
| inline | noinline | macroInline | inlineIfReduce | alwaysInline
|
|
deriving Inhabited, BEq, Hashable
|
|
|
|
private def InlineAttributeKind.toAttrString : InlineAttributeKind → String
|
|
| .inline => "inline"
|
|
| .noinline => "noinline"
|
|
| .macroInline => "macro_inline"
|
|
| .inlineIfReduce => "inline_if_reduce"
|
|
| .alwaysInline => "always_inline"
|
|
|
|
/--
|
|
This is an approximate test for testing whether `declName` can be annotated with the `[macro_inline]` attribute or not.
|
|
-/
|
|
private def isValidMacroInline (declName : Name) : CoreM Bool := do
|
|
let .defnInfo info ← getConstInfo declName
|
|
| return false
|
|
unless info.all.length = 1 do
|
|
-- We do not allow `[macro_inline]` attributes at mutual recursive definitions
|
|
return false
|
|
if (← Meta.isRecursiveDefinition declName) then
|
|
-- It is recursive
|
|
return false
|
|
return true
|
|
|
|
/--
|
|
Changes the inlining behavior. This attribute comes in several variants:
|
|
- `@[inline]`: marks the definition to be inlined when it is appropriate.
|
|
- `@[inline_if_reduce]`: marks the definition to be inlined if an application of it after inlining
|
|
and applying reduction isn't a `match` expression. This attribute can be used for inlining
|
|
structurally recursive functions.
|
|
- `@[noinline]`: marks the definition to never be inlined.
|
|
- `@[always_inline]`: marks the definition to always be inlined.
|
|
- `@[macro_inline]`: marks the definition to always be inlined at the beginning of compilation.
|
|
This makes it possible to define functions that evaluate some of their parameters lazily.
|
|
Example:
|
|
```
|
|
@[macro_inline]
|
|
def test (x y : Nat) : Nat :=
|
|
if x = 42 then x else y
|
|
|
|
#eval test 42 (2^1000000000000) -- doesn't compute 2^1000000000000
|
|
```
|
|
Only non-recursive functions may be marked `@[macro_inline]`.
|
|
-/
|
|
@[builtin_doc]
|
|
builtin_initialize inlineAttrs : EnumAttributes InlineAttributeKind ←
|
|
registerEnumAttributes
|
|
[(`inline, "mark definition to be inlined", .inline),
|
|
(`inline_if_reduce, "mark definition to be inlined when resultant term after reduction is not a `cases_on` application", .inlineIfReduce),
|
|
(`noinline, "mark definition to never be inlined", .noinline),
|
|
(`macro_inline, "mark definition to always be inlined before ANF conversion", .macroInline),
|
|
(`always_inline, "mark definition to be always inlined", .alwaysInline)]
|
|
fun declName kind => do
|
|
if kind matches .macroInline then
|
|
if !(checkIsDefinition (← getEnv) declName |>.isOk) then
|
|
throwError "invalid `[macro_inline]` attribute, `{.ofConstName declName}` must be an exposed definition"
|
|
unless (← isValidMacroInline declName) do
|
|
throwError "Cannot add `[macro_inline]` attribute to `{.ofConstName declName}`: This attribute does not support this kind of declaration; only non-recursive definitions are supported"
|
|
else
|
|
ofExcept <| (checkIsDefinition (← withoutExporting <| getEnv) declName).mapError fun e =>
|
|
s!"Cannot add attribute `[{kind.toAttrString}]`: {e}"
|
|
|
|
def setInlineAttribute (env : Environment) (declName : Name) (kind : InlineAttributeKind) : Except String Environment :=
|
|
inlineAttrs.setValue env declName kind
|
|
|
|
def getInlineAttribute? (env : Environment) (declName : Name) : Option InlineAttributeKind :=
|
|
inlineAttrs.getValue env declName
|
|
|
|
private def hasInlineAttrCore (env : Environment) (kind : InlineAttributeKind) (declName : Name) : Bool :=
|
|
match inlineAttrs.getValue env declName with
|
|
| some k => kind == k
|
|
| _ => false
|
|
|
|
@[inline] def hasInlineAttribute (env : Environment) (declName : Name) : Bool :=
|
|
hasInlineAttrCore env .inline declName
|
|
|
|
def hasInlineIfReduceAttribute (env : Environment) (declName : Name) : Bool :=
|
|
hasInlineAttrCore env .inlineIfReduce declName
|
|
|
|
def hasNoInlineAttribute (env : Environment) (declName : Name) : Bool :=
|
|
hasInlineAttrCore env .noinline declName
|
|
|
|
def hasMacroInlineAttribute (env : Environment) (declName : Name) : Bool :=
|
|
hasInlineAttrCore env .macroInline declName
|
|
|
|
@[inline] def hasAlwaysInlineAttribute (env : Environment) (declName : Name) : Bool :=
|
|
hasInlineAttrCore env .alwaysInline declName
|
|
|
|
end Lean.Compiler
|