lean4-htt/src/Lean/Compiler/MetaAttr.lean
2025-12-16 16:28:39 +00:00

114 lines
4.3 KiB
Text

/-
Copyright (c) 2025 Lean FRO. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Sebastian Ullrich
-/
module
prelude
public import Lean.EnvExtension
public section
namespace Lean
/-- Environment extension collecting `meta` annotations. -/
private builtin_initialize metaExt : TagDeclarationExtension ←
-- set by `addPreDefinitions`; if we ever make `def` elaboration async, it should be moved to
-- remain on the main environment branch
mkTagDeclarationExtension (asyncMode := .async .mainEnv)
/--
Environment extension collecting declarations *could* have been marked as `meta` by the user but
were not, so should not allow access to `meta` declarations to surface phase distinction errors as
soon as possible.
-/
private builtin_initialize notMetaExt : EnvExtension NameSet ←
registerEnvExtension
(mkInitial := pure {})
(asyncMode := .async .mainEnv)
(replay? := some fun _ _ newEntries s => newEntries.foldl (·.insert) s)
/-- Marks in the environment extension that the given declaration has been declared by the user as `meta`. -/
def markMeta (env : Environment) (declName : Name) : Environment :=
metaExt.tag env declName
/--
Marks the given declaration as not being annotated with `meta` even if it could have been by the
user.
-/
def markNotMeta (env : Environment) (declName : Name) : Environment :=
if declName.isAnonymous then -- avoid panic from `modifyState` on partial input
env
else
notMetaExt.modifyState (asyncDecl := declName) env (·.insert declName)
/-- Returns true iff the user has declared the given declaration as `meta`. -/
def isMarkedMeta (env : Environment) (declName : Name) : Bool :=
metaExt.isTagged env declName
/--
Set of IR decls that should be made available to any importer. This is a superset of `metaExt`
(except for non-defs such as `meta structure`), which is managed by the elaborator and has a
different async mode. More precisely, it contains the closure of `metaExt` as well as further
derived decls such as `_boxed` versions. We store this set primarily to filter exports in
`declMapExt`; we persist it in `.olean.private` for the benefit of `shake`.
-/
private builtin_initialize declMetaExt : SimplePersistentEnvExtension Name NameSet ←
registerSimplePersistentEnvExtension {
addImportedFn := fun _ => {}
addEntryFn := fun s n => s.insert n
asyncMode := .sync
replay? := some <| SimplePersistentEnvExtension.replayOfFilter (!·.contains ·) (·.insert ·)
exportEntriesFnEx? := some fun env s entries => fun
| .private =>
let decls := entries.foldl (init := #[]) fun decls decl => decls.push decl
decls.qsort Name.quickLt
| _ => #[]
}
/-- Whether a declaration should be exported for interpretation. -/
def isDeclMeta (env : Environment) (declName : Name) : Bool :=
if !env.header.isModule then
true
else
-- The interpreter may call the boxed variant even if the IR does not directly reference it, so
-- use same visibility as base decl.
-- Note that boxed decls are created after the `inferVisibility` pass.
let inferFor := match declName with
| .str n "_boxed" => n
| n => n
match env.getModuleIdxFor? declName with
| some idx => declMetaExt.getModuleEntries env idx |>.binSearchContains inferFor Name.quickLt
| none => declMetaExt.getState env |>.contains inferFor
/-- Marks a declaration to be exported for interpretation. -/
def setDeclMeta (env : Environment) (declName : Name) : Environment :=
if isDeclMeta env declName then
env
else
declMetaExt.addEntry env declName
/--
Returns the IR phases of the given declaration that should be considered accessible. Does not take
additional IR loaded for language server purposes into account.
-/
def getIRPhases (env : Environment) (declName : Name) : IRPhases := Id.run do
if !env.header.isModule then
return .all
match env.getModuleIdxFor? declName with
| some idx =>
if isMarkedMeta env declName then
.comptime
else
env.header.modules[idx]?.map (·.irPhases) |>.get!
| none =>
if isMarkedMeta env declName then
.comptime
else if notMetaExt.getState env |>.contains declName then
.runtime
else
-- Allow `meta`->non-`meta` references in the same module for auxiliary declarations the user
-- could not have marked as `meta` themselves.
.all
end Lean