fix: make new codegen async realization-compatible (#7316)

Follow-up to #7247

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

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
 

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

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)
   }
 
 /--

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

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 :=

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

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

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
   }
 
 /--

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 :=