95 lines
3.4 KiB
Text
95 lines
3.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
|
|
-/
|
|
import Lean.Compiler.LCNF.PassManager
|
|
|
|
namespace Lean.Compiler.LCNF
|
|
|
|
abbrev DeclExtState := PHashMap Name Decl
|
|
|
|
private abbrev declLt (a b : Decl) :=
|
|
Name.quickLt a.name b.name
|
|
|
|
private abbrev sortDecls (decls : Array Decl) : Array Decl :=
|
|
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
|
|
|
|
abbrev DeclExt := PersistentEnvExtension Decl Decl DeclExtState
|
|
|
|
def mkDeclExt (name : Name) : IO DeclExt := do
|
|
registerPersistentEnvExtension {
|
|
name := name
|
|
mkInitial := return {}
|
|
addImportedFn := fun _ => return {}
|
|
addEntryFn := fun decls decl => decls.insert decl.name decl
|
|
exportEntriesFn := fun s =>
|
|
let decls := s.foldl (init := #[]) fun decls _ decl => decls.push decl
|
|
sortDecls decls
|
|
}
|
|
|
|
builtin_initialize baseExt : PersistentEnvExtension Decl Decl DeclExtState ← mkDeclExt `compBaseDecls
|
|
builtin_initialize monoExt : PersistentEnvExtension Decl Decl DeclExtState ← mkDeclExt `compMonoDecls
|
|
|
|
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.addExtraName decl.name) decl
|
|
|
|
def saveMonoDeclCore (env : Environment) (decl : Decl) : Environment :=
|
|
monoExt.addEntry (env.addExtraName decl.name) decl
|
|
|
|
def Decl.saveBase (decl : Decl) : CoreM Unit :=
|
|
modifyEnv (saveBaseDeclCore · decl)
|
|
|
|
def Decl.saveMono (decl : Decl) : CoreM Unit :=
|
|
modifyEnv (saveMonoDeclCore · decl)
|
|
|
|
def getDeclAt? (declName : Name) (phase : Phase) : CoreM (Option Decl) :=
|
|
match phase with
|
|
| .base => getBaseDecl? declName
|
|
| .mono => getMonoDecl? declName
|
|
| _ => return none -- TODO
|
|
|
|
def getDecl? (declName : Name) : CompilerM (Option Decl) := do
|
|
getDeclAt? declName (← getPhase)
|
|
|
|
def getExt (phase : Phase) : DeclExt :=
|
|
match phase with
|
|
| .base => baseExt
|
|
| .mono => monoExt
|
|
| _ => unreachable!
|
|
|
|
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
|