feat: reuse specialized functions between different compilation units

leanprover · Oct 7, 2022 · 9eb641e · 9eb641e
1 parent f11e449
commit 9eb641e
Showing 1 changed file with 15 additions and 4 deletions.
diff --git a/src/Lean/Compiler/LCNF/Specialize.lean b/src/Lean/Compiler/LCNF/Specialize.lean
@@ -16,13 +16,24 @@ import Lean.Compiler.LCNF.Closure
 namespace Lean.Compiler.LCNF
 namespace Specialize
 
-abbrev Cache := PHashMap Expr Name
+abbrev Cache := SMap Expr Name
 
-builtin_initialize specCacheExt : EnvExtension Cache ←
-  registerEnvExtension (pure {})
+structure CacheEntry where
+  key : Expr
+  declName : Name
+  deriving Inhabited
+
+def addEntry (cache : Cache) (e : CacheEntry) : Cache :=
+  cache.insert e.key e.declName
+
+builtin_initialize specCacheExt : SimplePersistentEnvExtension CacheEntry Cache ←
+  registerSimplePersistentEnvExtension {
+    addEntryFn    := addEntry
+    addImportedFn := fun es => (mkStateFromImportedEntries addEntry {} es).switch
+  }
 
 def cacheSpec (key : Expr) (declName : Name) : CoreM Unit :=
-  modifyEnv fun env => specCacheExt.modifyState env (·.insert key declName)
+  modifyEnv fun env => specCacheExt.addEntry env { key, declName }
 
 def findSpecCache? (key : Expr) : CoreM (Option Name) :=
   return specCacheExt.getState (← getEnv) |>.find? key