diff --git a/src/Lean/Compiler/ClosedTermCache.lean b/src/Lean/Compiler/ClosedTermCache.lean index a620e09663..b4d4fa2a5e 100644 --- a/src/Lean/Compiler/ClosedTermCache.lean +++ b/src/Lean/Compiler/ClosedTermCache.lean @@ -11,10 +11,17 @@ namespace Lean structure ClosedTermCache where map : PHashMap Expr Name := {} constNames : NameSet := {} + -- used for `replay?` only + revExprs : List Expr := [] deriving Inhabited builtin_initialize closedTermCacheExt : EnvExtension ClosedTermCache ← registerEnvExtension (pure {}) (asyncMode := .sync) -- compilation is non-parallel anyway + (replay? := some fun oldState newState _ s => + let newExprs := newState.revExprs.take (newState.revExprs.length - oldState.revExprs.length) + newExprs.foldl (init := s) fun s e => + let c := newState.map.find! e + { s with map := s.map.insert e c, constNames := s.constNames.insert c, revExprs := e :: s.revExprs }) @[export lean_cache_closed_term_name] def cacheClosedTermName (env : Environment) (e : Expr) (n : Name) : Environment := diff --git a/src/Lean/Compiler/IR/CompilerM.lean b/src/Lean/Compiler/IR/CompilerM.lean index eeab571774..6e8a36ba9e 100644 --- a/src/Lean/Compiler/IR/CompilerM.lean +++ b/src/Lean/Compiler/IR/CompilerM.lean @@ -94,6 +94,7 @@ builtin_initialize declMapExt : SimplePersistentEnvExtension Decl DeclMap ← -- share a name prefix with the top-level Lean declaration being compiled, e.g. from -- specialization. asyncMode := .sync + replay? := some <| SimplePersistentEnvExtension.replayOfFilter (!·.contains ·.name) (fun s d => s.insert d.name d) } @[export lean_ir_find_env_decl] diff --git a/src/Lean/Compiler/IR/ElimDeadBranches.lean b/src/Lean/Compiler/IR/ElimDeadBranches.lean index bc77895958..e74c3eb05c 100644 --- a/src/Lean/Compiler/IR/ElimDeadBranches.lean +++ b/src/Lean/Compiler/IR/ElimDeadBranches.lean @@ -143,6 +143,7 @@ builtin_initialize functionSummariesExt : SimplePersistentEnvExtension (FunId × addEntryFn := fun s ⟨e, n⟩ => s.insert e n toArrayFn := fun s => sortEntries s.toArray asyncMode := .sync -- compilation is non-parallel anyway + replay? := some <| SimplePersistentEnvExtension.replayOfFilter (!·.contains ·.1) (fun s ⟨e, n⟩ => s.insert e n) } def addFunctionSummary (env : Environment) (fid : FunId) (v : Value) : Environment := diff --git a/src/Lean/Compiler/Specialize.lean b/src/Lean/Compiler/Specialize.lean index 2b841e5f1e..3723b0edaf 100644 --- a/src/Lean/Compiler/Specialize.lean +++ b/src/Lean/Compiler/Specialize.lean @@ -111,6 +111,9 @@ builtin_initialize specExtension : SimplePersistentEnvExtension SpecEntry SpecSt addEntryFn := SpecState.addEntry, addImportedFn := fun es => (mkStateFromImportedEntries SpecState.addEntry {} es).switch asyncMode := .sync -- compilation is non-parallel anyway + replay? := some <| SimplePersistentEnvExtension.replayOfFilter (fun + | s, .info n _ => !s.specInfo.contains n + | s, .cache key _ => !s.cache.contains key) SpecState.addEntry } @[export lean_add_specialization_info] diff --git a/src/Lean/Environment.lean b/src/Lean/Environment.lean index d17e40a649..9d834ff897 100644 --- a/src/Lean/Environment.lean +++ b/src/Lean/Environment.lean @@ -1351,6 +1351,18 @@ structure SimplePersistentEnvExtensionDescr (α σ : Type) where addImportedFn : Array (Array α) → σ toArrayFn : List α → Array α := fun es => es.toArray asyncMode : EnvExtension.AsyncMode := .mainOnly + replay? : Option ((newEntries : List α) → (newState : σ) → σ → List α × σ) := none + +/-- +Returns a function suitable for `SimplePersistentEnvExtensionDescr.replay?` that replays all new +entries onto the state using `addEntryFn`. `p` should filter out entries that have already been +added to the state by a prior replay of the same realizable constant. +-/ +def SimplePersistentEnvExtension.replayOfFilter (p : σ → α → Bool) + (addEntryFn : σ → α → σ) : List α → σ → σ → List α × σ := + fun newEntries _ s => + let newEntries := newEntries.filter (p s) + (newEntries, newEntries.foldl (init := s) addEntryFn) def registerSimplePersistentEnvExtension {α σ : Type} [Inhabited σ] (descr : SimplePersistentEnvExtensionDescr α σ) : IO (SimplePersistentEnvExtension α σ) := registerPersistentEnvExtension { @@ -1362,9 +1374,10 @@ def registerSimplePersistentEnvExtension {α σ : Type} [Inhabited σ] (descr : exportEntriesFn := fun s => descr.toArrayFn s.1.reverse, statsFn := fun s => format "number of local entries: " ++ format s.1.length asyncMode := descr.asyncMode - replay? := some fun oldState newState _ (entries, s) => - let newEntries := newState.1.drop oldState.1.length - (newEntries ++ entries, newEntries.foldl descr.addEntryFn s) + replay? := descr.replay?.map fun replay oldState newState _ (entries, s) => + let newEntries := newState.1.take (newState.1.length - oldState.1.length) + let (newEntries, s) := replay newEntries newState.2 s + (entries ++ newEntries, s) } namespace SimplePersistentEnvExtension diff --git a/src/Lean/Meta/Basic.lean b/src/Lean/Meta/Basic.lean index 7b1311140d..689c71c8ae 100644 --- a/src/Lean/Meta/Basic.lean +++ b/src/Lean/Meta/Basic.lean @@ -2255,8 +2255,12 @@ def realizeConst (forConst : Name) (constName : Name) (realize : MetaM Unit) : return withTraceNode `Meta.realizeConst (fun _ => return constName) do let coreCtx ← readThe Core.Context - -- these fields should be invariant throughout the file - let coreCtx := { fileName := coreCtx.fileName, fileMap := coreCtx.fileMap } + let coreCtx := { + -- these fields should be invariant throughout the file + fileName := coreCtx.fileName, fileMap := coreCtx.fileMap + -- heartbeat limits inside `realizeAndReport` should be measured from this point on + initHeartbeats := (← IO.getNumHeartbeats) + } let (env, dyn) ← env.realizeConst forConst constName (realizeAndReport coreCtx) if let some res := dyn.get? RealizeConstantResult then let mut snap := res.snap