diff --git a/src/Lean/Compiler/LCNF/AuxDeclCache.lean b/src/Lean/Compiler/LCNF/AuxDeclCache.lean index 8552de6805..1e1347138b 100644 --- a/src/Lean/Compiler/LCNF/AuxDeclCache.lean +++ b/src/Lean/Compiler/LCNF/AuxDeclCache.lean @@ -10,10 +10,7 @@ import Lean.Compiler.LCNF.Internalize namespace Lean.Compiler.LCNF -abbrev AuxDeclCache := PHashMap Decl Name - -builtin_initialize auxDeclCacheExt : EnvExtension AuxDeclCache ← - registerEnvExtension (pure {}) (asyncMode := .sync) -- compilation is non-parallel anyway +builtin_initialize auxDeclCacheExt : CacheExtension Decl Name ← CacheExtension.register inductive CacheAuxDeclResult where | new @@ -22,11 +19,11 @@ inductive CacheAuxDeclResult where def cacheAuxDecl (decl : Decl) : CompilerM CacheAuxDeclResult := do let key := { decl with name := .anonymous } let key ← normalizeFVarIds key - match auxDeclCacheExt.getState (← getEnv) |>.find? key with + match (← auxDeclCacheExt.find? key) with | some declName => return .alreadyCached declName | none => - modifyEnv fun env => auxDeclCacheExt.modifyState env fun s => s.insert key decl.name + auxDeclCacheExt.insert key decl.name return .new end Lean.Compiler.LCNF diff --git a/src/Lean/Compiler/LCNF/BaseTypes.lean b/src/Lean/Compiler/LCNF/BaseTypes.lean index 70a6a1266a..a6d96e2aef 100644 --- a/src/Lean/Compiler/LCNF/BaseTypes.lean +++ b/src/Lean/Compiler/LCNF/BaseTypes.lean @@ -14,21 +14,15 @@ State for the environment extension used to save the LCNF base phase type for de that do not have code associated with them. Example: constructors, inductive types, foreign functions. -/ -structure BaseTypeExtState where - /-- The LCNF type for the `base` phase. -/ - base : PHashMap Name Expr := {} - deriving Inhabited - -builtin_initialize baseTypeExt : EnvExtension BaseTypeExtState ← - registerEnvExtension (pure {}) (asyncMode := .sync) -- compilation is non-parallel anyway +builtin_initialize baseTypeExt : CacheExtension Name Expr ← CacheExtension.register def getOtherDeclBaseType (declName : Name) (us : List Level) : CoreM Expr := do let info ← getConstInfo declName - let type ← match baseTypeExt.getState (← getEnv) |>.base.find? declName with + let type ← match (← baseTypeExt.find? declName) with | some type => pure type | none => let type ← Meta.MetaM.run' <| toLCNFType info.type - modifyEnv fun env => baseTypeExt.modifyState env fun s => { s with base := s.base.insert declName type } + baseTypeExt.insert declName type pure type return type.instantiateLevelParamsNoCache info.levelParams us diff --git a/src/Lean/Compiler/LCNF/CompilerM.lean b/src/Lean/Compiler/LCNF/CompilerM.lean index 4853db635b..db9527894d 100644 --- a/src/Lean/Compiler/LCNF/CompilerM.lean +++ b/src/Lean/Compiler/LCNF/CompilerM.lean @@ -483,4 +483,26 @@ def getConfig : CompilerM ConfigOptions := def CompilerM.run (x : CompilerM α) (s : State := {}) (phase : Phase := .base) : CoreM α := do x { phase, config := toConfigOptions (← getOptions) } |>.run' s +/-- Environment extension for local caching of key-value pairs, not persisted in .olean files. -/ +structure CacheExtension (α β : Type) [BEq α] [Hashable α] extends EnvExtension (List α × PHashMap α β) +deriving Inhabited + +namespace CacheExtension + +def register [BEq α] [Hashable α] [Inhabited β] : + IO (CacheExtension α β) := + CacheExtension.mk <$> registerEnvExtension (pure ([], {})) (asyncMode := .sync) -- compilation is non-parallel anyway + (replay? := some fun oldState newState _ s => + let newEntries := newState.1.take (newState.1.length - oldState.1.length) + newEntries.foldl (init := s) fun s e => + (e :: s.1, s.2.insert e (newState.2.find! e))) + +def insert [BEq α] [Hashable α] [Inhabited β] (ext : CacheExtension α β) (a : α) (b : β) : CoreM Unit := do + modifyEnv (ext.modifyState · fun ⟨as, m⟩ => (a :: as, m.insert a b)) + +def find? [BEq α] [Hashable α] [Inhabited β] (ext : CacheExtension α β) (a : α) : CoreM (Option β) := do + return ext.toEnvExtension.getState (← getEnv) |>.2.find? a + +end CacheExtension + end Lean.Compiler.LCNF diff --git a/src/Lean/Compiler/LCNF/ElimDeadBranches.lean b/src/Lean/Compiler/LCNF/ElimDeadBranches.lean index fb64f6cc98..75051f7946 100644 --- a/src/Lean/Compiler/LCNF/ElimDeadBranches.lean +++ b/src/Lean/Compiler/LCNF/ElimDeadBranches.lean @@ -249,6 +249,7 @@ builtin_initialize functionSummariesExt : SimplePersistentEnvExtension (Name × addEntryFn := fun s ⟨e, n⟩ => s.insert e n toArrayFn := fun s => s.toArray.qsort decLt asyncMode := .sync -- compilation is non-parallel anyway + replay? := some <| SimplePersistentEnvExtension.replayOfFilter (!·.contains ·.1) (fun s ⟨e, n⟩ => s.insert e n) } /-- diff --git a/src/Lean/Compiler/LCNF/MonoTypes.lean b/src/Lean/Compiler/LCNF/MonoTypes.lean index cd38edcf88..8983ca43e1 100644 --- a/src/Lean/Compiler/LCNF/MonoTypes.lean +++ b/src/Lean/Compiler/LCNF/MonoTypes.lean @@ -111,20 +111,14 @@ State for the environment extension used to save the LCNF mono phase type for de that do not have code associated with them. Example: constructors, inductive types, foreign functions. -/ -structure MonoTypeExtState where - /-- The LCNF type for the `mono` phase. -/ - mono : PHashMap Name Expr := {} - deriving Inhabited - -builtin_initialize monoTypeExt : EnvExtension MonoTypeExtState ← - registerEnvExtension (pure {}) (asyncMode := .sync) -- compilation is non-parallel anyway +builtin_initialize monoTypeExt : CacheExtension Name Expr ← CacheExtension.register def getOtherDeclMonoType (declName : Name) : CoreM Expr := do - match monoTypeExt.getState (← getEnv) |>.mono.find? declName with + match (← monoTypeExt.find? declName) with | some type => return type | none => let type ← toMonoType (← getOtherDeclBaseType declName []) - modifyEnv fun env => monoTypeExt.modifyState env fun s => { s with mono := s.mono.insert declName type } + monoTypeExt.insert declName type return type end Lean.Compiler.LCNF diff --git a/src/Lean/Compiler/LCNF/Passes.lean b/src/Lean/Compiler/LCNF/Passes.lean index 4579073511..1356a0a4b4 100644 --- a/src/Lean/Compiler/LCNF/Passes.lean +++ b/src/Lean/Compiler/LCNF/Passes.lean @@ -94,7 +94,6 @@ builtin_initialize passManagerExt : PersistentEnvExtension Name (Name × PassMan addImportedFn := fun ns => return ([], ← ImportM.runCoreM <| runImportedDecls ns) addEntryFn := fun (installerDeclNames, _) (installerDeclName, managerNew) => (installerDeclName :: installerDeclNames, managerNew) exportEntriesFn := fun s => s.1.reverse.toArray - asyncMode := .sync } def getPassManager : CoreM PassManager := diff --git a/src/Lean/Compiler/LCNF/PhaseExt.lean b/src/Lean/Compiler/LCNF/PhaseExt.lean index d41fe344aa..9da5dfd007 100644 --- a/src/Lean/Compiler/LCNF/PhaseExt.lean +++ b/src/Lean/Compiler/LCNF/PhaseExt.lean @@ -21,22 +21,21 @@ private abbrev findAtSorted? (decls : Array Decl) (declName : Name) : Option Dec let tmpDecl := { tmpDecl with name := declName } decls.binSearch tmpDecl declLt -abbrev DeclExt := PersistentEnvExtension Decl Decl DeclExtState +abbrev DeclExt := SimplePersistentEnvExtension Decl DeclExtState def mkDeclExt (name : Name := by exact decl_name%) : IO DeclExt := do - registerPersistentEnvExtension { + registerSimplePersistentEnvExtension { name := name - mkInitial := return {} - addImportedFn := fun _ => return {} + addImportedFn := fun _ => {} 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 + toArrayFn := (sortDecls ·.toArray) asyncMode := .sync -- compilation is non-parallel anyway + replay? := some <| SimplePersistentEnvExtension.replayOfFilter + (fun s d => !s.contains d.name) (fun decls decl => decls.insert decl.name decl) } -builtin_initialize baseExt : PersistentEnvExtension Decl Decl DeclExtState ← mkDeclExt -builtin_initialize monoExt : PersistentEnvExtension Decl Decl DeclExtState ← mkDeclExt +builtin_initialize baseExt : DeclExt ← mkDeclExt +builtin_initialize monoExt : DeclExt ← mkDeclExt def getDeclCore? (env : Environment) (ext : DeclExt) (declName : Name) : Option Decl := match env.getModuleIdxFor? declName with diff --git a/src/Lean/Compiler/LCNF/Simp/ConstantFold.lean b/src/Lean/Compiler/LCNF/Simp/ConstantFold.lean index 5551fb028c..ecdc9cbafb 100644 --- a/src/Lean/Compiler/LCNF/Simp/ConstantFold.lean +++ b/src/Lean/Compiler/LCNF/Simp/ConstantFold.lean @@ -397,7 +397,7 @@ structure FolderOleanEntry where structure FolderEntry extends FolderOleanEntry where folder : Folder -builtin_initialize folderExt : PersistentEnvExtension FolderOleanEntry FolderEntry (List FolderOleanEntry × SMap Name Folder) ← +builtin_initialize folderExt : PersistentEnvExtension FolderOleanEntry FolderEntry (List FolderEntry × SMap Name Folder) ← registerPersistentEnvExtension { mkInitial := return ([], builtinFolders) addImportedFn := fun entriesArray => do @@ -408,9 +408,12 @@ builtin_initialize folderExt : PersistentEnvExtension FolderOleanEntry FolderEnt let folder ← IO.ofExcept <| getFolderCore ctx.env ctx.opts folderDeclName folders := folders.insert declName folder return ([], folders.switch) - addEntryFn := fun (entries, map) entry => (entry.toFolderOleanEntry :: entries, map.insert entry.declName entry.folder) - exportEntriesFn := fun (entries, _) => entries.reverse.toArray + addEntryFn := fun (entries, map) entry => (entry :: entries, map.insert entry.declName entry.folder) + exportEntriesFn := fun (entries, _) => entries.reverse.toArray.map (·.toFolderOleanEntry) asyncMode := .sync + replay? := some fun oldState newState _ s => + let newEntries := newState.1.take (newState.1.length - oldState.1.length) + (newEntries ++ s.1, newEntries.foldl (init := s.2) fun s e => s.insert e.declName (newState.2.find! e.declName)) } def registerFolder (declName : Name) (folderDeclName : Name) : CoreM Unit := do diff --git a/src/Lean/Compiler/LCNF/SpecInfo.lean b/src/Lean/Compiler/LCNF/SpecInfo.lean index 3571d7fda3..8c8403d57e 100644 --- a/src/Lean/Compiler/LCNF/SpecInfo.lean +++ b/src/Lean/Compiler/LCNF/SpecInfo.lean @@ -86,6 +86,8 @@ builtin_initialize specExtension : SimplePersistentEnvExtension SpecEntry SpecSt addImportedFn := fun _ => {} toArrayFn := fun s => sortEntries s.toArray asyncMode := .sync + replay? := some <| SimplePersistentEnvExtension.replayOfFilter + (!·.specInfo.contains ·.declName) SpecState.addEntry } /-- diff --git a/src/Lean/Compiler/LCNF/Specialize.lean b/src/Lean/Compiler/LCNF/Specialize.lean index ec626d3b5f..00404b92c5 100644 --- a/src/Lean/Compiler/LCNF/Specialize.lean +++ b/src/Lean/Compiler/LCNF/Specialize.lean @@ -33,6 +33,8 @@ builtin_initialize specCacheExt : SimplePersistentEnvExtension CacheEntry Cache addEntryFn := addEntry addImportedFn := fun es => (mkStateFromImportedEntries addEntry {} es).switch asyncMode := .sync + replay? := some <| SimplePersistentEnvExtension.replayOfFilter + (!·.contains ·.key) addEntry } def cacheSpec (key : Expr) (declName : Name) : CoreM Unit :=