201 lines
7.4 KiB
Text
201 lines
7.4 KiB
Text
/-
|
||
Copyright (c) 2022 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.Compiler.LCNF.PassManager
|
||
|
||
public section
|
||
|
||
namespace Lean.Compiler.LCNF
|
||
|
||
/-- Creates a replayable local environment extension holding a name set. -/
|
||
private def mkDeclSetExt : IO (EnvExtension (List Name × NameSet)) :=
|
||
registerEnvExtension
|
||
(mkInitial := pure ([], {}))
|
||
(asyncMode := .sync)
|
||
(replay? := some <| fun oldState newState _ s =>
|
||
let newEntries := newState.1.take (newState.1.length - oldState.1.length)
|
||
newEntries.foldl (init := s) fun s n =>
|
||
if s.1.contains n then
|
||
s
|
||
else
|
||
(n :: s.1, if newState.2.contains n then s.2.insert n else s.2))
|
||
|
||
/--
|
||
Set of declarations to be exported to other modules; visibility shared by base/mono/IR phases.
|
||
-/
|
||
private builtin_initialize publicDeclsExt : EnvExtension (List Name × NameSet) ← mkDeclSetExt
|
||
|
||
-- The following extensions are per-phase as e.g. a declaration may be inlineable in the mono phase
|
||
-- because it directly unfolds to a `_redArg` call, but making its body in the base phase available
|
||
-- as well may lead to opportunistic inlining there, exposing references to declarations that we did
|
||
-- not export, leading to compilation errors. Ensuring those declarations are exported as well would
|
||
-- be possible, but would mean more information than intended is exported, potentially leading to
|
||
-- surprising rebuilds.
|
||
|
||
/--
|
||
Set of public declarations whose base bodies should be exported to other modules
|
||
-/
|
||
private builtin_initialize baseTransparentDeclsExt : EnvExtension (List Name × NameSet) ← mkDeclSetExt
|
||
/--
|
||
Set of public declarations whose mono bodies should be exported to other modules
|
||
-/
|
||
private builtin_initialize monoTransparentDeclsExt : EnvExtension (List Name × NameSet) ← mkDeclSetExt
|
||
|
||
private def getTransparencyExt : Phase → EnvExtension (List Name × NameSet)
|
||
| .base => baseTransparentDeclsExt
|
||
| .mono => monoTransparentDeclsExt
|
||
|
||
def isDeclPublic (env : Environment) (declName : Name) : Bool := Id.run do
|
||
if !env.header.isModule then
|
||
return true
|
||
-- The IR compiler may call the boxed variant it introduces after we do visibility inference, so
|
||
-- use same visibility as base decl.
|
||
let inferFor := match declName with
|
||
| .str n "_boxed" => n
|
||
| n => n
|
||
let (_, map) := publicDeclsExt.getState env
|
||
map.contains inferFor
|
||
|
||
def setDeclPublic (env : Environment) (declName : Name) : Environment :=
|
||
if isDeclPublic env declName then
|
||
env
|
||
else
|
||
publicDeclsExt.modifyState env fun s =>
|
||
(declName :: s.1, s.2.insert declName)
|
||
|
||
def isDeclTransparent (env : Environment) (phase : Phase) (declName : Name) : Bool := Id.run do
|
||
if !env.header.isModule then
|
||
return true
|
||
let (_, map) := getTransparencyExt phase |>.getState env
|
||
map.contains declName
|
||
|
||
def setDeclTransparent (env : Environment) (phase : Phase) (declName : Name) : Environment :=
|
||
if isDeclTransparent env phase declName then
|
||
env
|
||
else
|
||
getTransparencyExt phase |>.modifyState env fun s =>
|
||
(declName :: s.1, s.2.insert declName)
|
||
|
||
abbrev DeclExtState := PHashMap Name Decl
|
||
|
||
private abbrev declLt (a b : Decl) :=
|
||
Name.quickLt a.name b.name
|
||
|
||
private def sortedDecls (s : DeclExtState) : Array Decl :=
|
||
let decls := s.foldl (init := #[]) fun ps _ v => ps.push v
|
||
decls.qsort declLt
|
||
|
||
private abbrev findAtSorted? (decls : Array Decl) (declName : Name) : Option Decl :=
|
||
let tmpDecl : Decl := default
|
||
let tmpDecl := { tmpDecl with name := declName }
|
||
decls.binSearch tmpDecl declLt
|
||
|
||
@[expose] def DeclExt := PersistentEnvExtension Decl Decl DeclExtState
|
||
|
||
instance : Inhabited DeclExt :=
|
||
inferInstanceAs (Inhabited (PersistentEnvExtension Decl Decl DeclExtState))
|
||
|
||
def mkDeclExt (phase : Phase) (name : Name := by exact decl_name%) : IO DeclExt :=
|
||
registerPersistentEnvExtension {
|
||
name,
|
||
mkInitial := pure {},
|
||
addImportedFn := fun _ => pure {},
|
||
addEntryFn := fun s decl => s.insert decl.name decl
|
||
exportEntriesFnEx env s level := Id.run do
|
||
let mut entries := sortedDecls s
|
||
if level != .private then
|
||
entries := entries.filterMap fun decl => do
|
||
guard <| isDeclPublic env decl.name
|
||
if isDeclTransparent env phase decl.name then
|
||
some decl
|
||
else
|
||
some { decl with value := .extern { arity? := decl.getArity, entries := [.opaque decl.name] } }
|
||
return entries
|
||
statsFn := fun s =>
|
||
let numEntries := s.foldl (init := 0) (fun count _ _ => count + 1)
|
||
format "number of local entries: " ++ format numEntries
|
||
asyncMode := .sync,
|
||
replay? := some <| fun oldState newState _ otherState =>
|
||
newState.foldl (init := otherState) fun otherState k v =>
|
||
if oldState.contains k then
|
||
otherState
|
||
else
|
||
otherState.insert k v
|
||
}
|
||
|
||
builtin_initialize baseExt : DeclExt ← mkDeclExt .base
|
||
builtin_initialize monoExt : DeclExt ← mkDeclExt .mono
|
||
|
||
def getDeclCore? (env : Environment) (ext : DeclExt) (declName : Name) : Option Decl :=
|
||
match env.getModuleIdxFor? declName with
|
||
| some modIdx => findAtSorted? (ext.getModuleEntries env modIdx) declName
|
||
| none => ext.getState env |>.find? declName
|
||
|
||
def getBaseDecl? (declName : Name) : CoreM (Option Decl) := do
|
||
return getDeclCore? (← getEnv) baseExt declName
|
||
|
||
def getMonoDecl? (declName : Name) : CoreM (Option Decl) := do
|
||
return getDeclCore? (← getEnv) monoExt declName
|
||
|
||
def saveBaseDeclCore (env : Environment) (decl : Decl) : Environment :=
|
||
baseExt.addEntry env decl
|
||
|
||
def saveMonoDeclCore (env : Environment) (decl : Decl) : Environment :=
|
||
monoExt.addEntry env decl
|
||
|
||
def Decl.saveBase (decl : Decl) : CoreM Unit :=
|
||
modifyEnv (saveBaseDeclCore · decl)
|
||
|
||
def Decl.saveMono (decl : Decl) : CoreM Unit :=
|
||
modifyEnv (saveMonoDeclCore · decl)
|
||
|
||
def Decl.save (decl : Decl) : CompilerM Unit := do
|
||
match (← getPhase) with
|
||
| .base => decl.saveBase
|
||
| .mono => decl.saveMono
|
||
|
||
def getDeclAt? (declName : Name) (phase : Phase) : CoreM (Option Decl) :=
|
||
match phase with
|
||
| .base => getBaseDecl? declName
|
||
| .mono => getMonoDecl? declName
|
||
|
||
def getDecl? (declName : Name) : CompilerM (Option Decl) := do
|
||
getDeclAt? declName (← getPhase)
|
||
|
||
def getLocalDeclAt? (declName : Name) (phase : Phase) : CompilerM (Option Decl) := do
|
||
match phase with
|
||
| .base => return baseExt.getState (← getEnv) |>.find? declName
|
||
| .mono => return monoExt.getState (← getEnv) |>.find? declName
|
||
|
||
def getLocalDecl? (declName : Name) : CompilerM (Option Decl) := do
|
||
getLocalDeclAt? declName (← getPhase)
|
||
|
||
def getExt (phase : Phase) : DeclExt :=
|
||
match phase with
|
||
| .base => baseExt
|
||
| .mono => monoExt
|
||
|
||
def forEachDecl (f : Decl → CoreM Unit) (phase := Phase.base) : CoreM Unit := do
|
||
let ext := getExt phase
|
||
let env ← getEnv
|
||
for modIdx in *...env.allImportedModuleNames.size do
|
||
for decl in ext.getModuleEntries env modIdx do
|
||
f decl
|
||
ext.getState env |>.forM fun _ decl => f decl
|
||
|
||
def forEachModuleDecl (moduleName : Name) (f : Decl → CoreM Unit) (phase := Phase.base) : CoreM Unit := do
|
||
let ext := getExt phase
|
||
let env ← getEnv
|
||
let some modIdx := env.getModuleIdx? moduleName | throwError "module `{moduleName}` not found"
|
||
for decl in ext.getModuleEntries env modIdx do
|
||
f decl
|
||
|
||
def forEachMainModuleDecl (f : Decl → CoreM Unit) (phase := Phase.base) : CoreM Unit := do
|
||
(getExt phase).getState (← getEnv) |>.forM fun _ decl => f decl
|
||
|
||
end Lean.Compiler.LCNF
|