Caches belong to runtime, not IR compiler (#8178)

engine/runtime-instrument-common/src/main/java/org/enso/interpreter/instrument/job/SerializeModuleJob.java

@@ -1,9 +1,8 @@
 package org.enso.interpreter.instrument.job;
 
 import java.util.logging.Level;
-import org.enso.compiler.SerializationManager;
 import org.enso.interpreter.instrument.execution.RuntimeContext;
-import org.enso.interpreter.runtime.EnsoContext;
+import org.enso.interpreter.runtime.SerializationManager;
 import org.enso.pkg.QualifiedName;
 import org.enso.polyglot.CompilationStage;
 
@@ -21,8 +20,9 @@ public final class SerializeModuleJob extends BackgroundJob<Void> {
 
   @Override
   public Void run(RuntimeContext ctx) {
-    EnsoContext ensoContext = ctx.executionService().getContext();
+    var ensoContext = ctx.executionService().getContext();
-    SerializationManager serializationManager = ensoContext.getCompiler().getSerializationManager();
+    var compiler = ensoContext.getCompiler();
+    SerializationManager serializationManager = SerializationManager.apply(compiler.context());
     boolean useGlobalCacheLocations = ensoContext.isUseGlobalCache();
     var writeLockTimestamp = ctx.locking().acquireWriteCompilationLock();
     try {
@@ -42,7 +42,7 @@ public final class SerializeModuleJob extends BackgroundJob<Void> {
                 }
 
                 serializationManager.serializeModule(
-                    module.asCompilerModule(), useGlobalCacheLocations, false);
+                    compiler, module.asCompilerModule(), useGlobalCacheLocations, false);
               });
     } finally {
       ctx.locking().releaseWriteCompilationLock();

engine/runtime-instrument-common/src/main/scala/org/enso/interpreter/instrument/job/DeserializeLibrarySuggestionsJob.scala

@@ -1,6 +1,7 @@
 package org.enso.interpreter.instrument.job
 
 import org.enso.editions.LibraryName
+import org.enso.interpreter.runtime.SerializationManager
 import org.enso.interpreter.instrument.execution.RuntimeContext
 import org.enso.polyglot.runtime.Runtime.Api
 
@@ -31,8 +32,9 @@ final class DeserializeLibrarySuggestionsJob(
       Level.FINE,
       s"Deserializing suggestions for library [$libraryName]."
     )
-    val serializationManager =
+    val serializationManager = SerializationManager(
-      ctx.executionService.getContext.getCompiler.getSerializationManager
+      ctx.executionService.getContext.getCompiler.context
+    )
     serializationManager
       .deserializeSuggestions(libraryName)
       .foreach { cachedSuggestions =>

engine/runtime/src/main/java/org/enso/compiler/context/CompilerContext.java

@@ -5,17 +5,15 @@ import com.oracle.truffle.api.source.Source;
 import java.io.IOException;
 import java.io.PrintStream;
 import java.util.List;
-import java.util.Optional;
+import java.util.concurrent.Future;
 import java.util.function.Consumer;
 import java.util.logging.Level;
-import org.enso.compiler.Cache;
 import org.enso.compiler.Compiler;
-import org.enso.compiler.ModuleCache;
 import org.enso.compiler.PackageRepository;
 import org.enso.compiler.Passes;
-import org.enso.compiler.SerializationManager;
 import org.enso.compiler.data.BindingsMap;
 import org.enso.compiler.data.CompilerConfig;
+import org.enso.editions.LibraryName;
 import org.enso.pkg.Package;
 import org.enso.pkg.QualifiedName;
 import org.enso.polyglot.CompilationStage;
@@ -66,10 +64,7 @@ public interface CompilerContext {
   boolean typeContainsValues(String name);
 
   void initializeBuiltinsIr(
-      boolean irCachingEnabled,
+      Compiler compiler, boolean irCachingEnabled, FreshNameSupply freshNameSupply, Passes passes);
-      SerializationManager serializationManager,
-      FreshNameSupply freshNameSupply,
-      Passes passes);
 
   QualifiedName getModuleName(Module module);
 
@@ -89,12 +84,18 @@ public interface CompilerContext {
 
   CompilationStage getCompilationStage(Module module);
 
-  <T> Optional<T> loadCache(Cache<T, ?> cache);
-
-  <T> Optional<TruffleFile> saveCache(Cache<T, ?> cache, T entry, boolean useGlobalCacheLocations);
-
   TypeGraph getTypeHierarchy();
 
+  Future<Boolean> serializeLibrary(
+      Compiler compiler, LibraryName libraryName, boolean useGlobalCacheLocations);
+
+  Future<Boolean> serializeModule(
+      Compiler compiler, Module module, boolean useGlobalCacheLocations);
+
+  boolean deserializeModule(Compiler compiler, Module module);
+
+  void shutdown(boolean waitForPendingJobCompletion);
+
   public static interface Updater {
     void bindingsMap(BindingsMap map);
 
@@ -128,8 +129,6 @@ public interface CompilerContext {
 
     public abstract List<QualifiedName> getDirectModulesRefs();
 
-    public abstract ModuleCache getCache();
-
     public abstract CompilationStage getCompilationStage();
 
     public abstract boolean isSynthetic();

engine/runtime/src/main/java/org/enso/interpreter/caches/Cache.java

@@ -1,4 +1,4 @@
-package org.enso.compiler;
+package org.enso.interpreter.caches;
 
 import com.oracle.truffle.api.TruffleFile;
 import com.oracle.truffle.api.TruffleLogger;

engine/runtime/src/main/java/org/enso/interpreter/caches/ImportExportCache.java

@@ -1,4 +1,4 @@
-package org.enso.compiler;
+package org.enso.interpreter.caches;
 
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
@@ -114,40 +114,19 @@ public final class ImportExportCache extends Cache<ImportExportCache.CachedBindi
       return byteStream.toByteArray();
     }
 
-    static class MapToBindings implements Serializable {
+    public static final class MapToBindings implements Serializable {
         private final Map<QualifiedName, BindingsMap> _entries;
 
         public MapToBindings(Map<QualifiedName, BindingsMap> entries) {
             this._entries = entries;
         }
 
-        Map<QualifiedName, BindingsMap> entries() {
+        public Map<QualifiedName, BindingsMap> entries() {
             return _entries;
         }
     }
 
-    // CachedBindings is not a record **on purpose**. There appears to be a Frgaal bug leading to invalid compilation error.
+    public static record CachedBindings(LibraryName libraryName, MapToBindings bindings, Optional<List<SourceFile<TruffleFile>>> sources) {
-    public static final class CachedBindings {
-        private final LibraryName _libraryName;
-        private final MapToBindings _bindings;
-        private final Optional<List<SourceFile<TruffleFile>>> _sources;
-
-        CachedBindings(LibraryName libraryName, MapToBindings bindings, Optional<List<SourceFile<TruffleFile>>> sources) {
-            this._libraryName = libraryName;
-            this._bindings = bindings;
-            this._sources = sources;
-        }
-
-        LibraryName libraryName() {
-            return _libraryName;
-        }
-        MapToBindings bindings() {
-            return _bindings;
-        }
-        Optional<List<SourceFile<TruffleFile>>> sources() {
-            return _sources;
-        }
-
     }
 
     public record Metadata(

engine/runtime/src/main/java/org/enso/interpreter/caches/ModuleCache.java

@@ -1,4 +1,4 @@
-package org.enso.compiler;
+package org.enso.interpreter.caches;
 
 import buildinfo.Info;
 import com.fasterxml.jackson.annotation.JsonProperty;
@@ -175,27 +175,7 @@ public final class ModuleCache extends Cache<ModuleCache.CachedModule, ModuleCac
       return byteStream.toByteArray();
     }
 
-    // CachedModule is not a record **on purpose**. There appears to be a Frgaal bug leading to invalid compilation error.
+    public record CachedModule(Module moduleIR, CompilationStage compilationStage, Source source) {
-    static class CachedModule {
-        private final Module _moduleIR;
-        private final CompilationStage _compilationStage;
-        private final Source _source;
-        public CachedModule(Module moduleIR, CompilationStage compilationStage, Source source) {
-            this._moduleIR = moduleIR;
-            this._compilationStage = compilationStage;
-            this._source = source;
-        }
-
-        Module moduleIR() {
-            return _moduleIR;
-        }
-        CompilationStage compilationStage() {
-            return _compilationStage;
-        }
-        Source source() {
-            return _source;
-        }
-
     }
 
     public record Metadata(

engine/runtime/src/main/java/org/enso/interpreter/caches/SuggestionsCache.java

@@ -1,4 +1,4 @@
-package org.enso.compiler;
+package org.enso.interpreter.caches;
 
 import buildinfo.Info;
 import com.fasterxml.jackson.annotation.JsonProperty;
@@ -115,7 +115,7 @@ public final class SuggestionsCache
   }
 
   // Suggestions class is not a record because of a Frgaal bug leading to invalid compilation error.
-  final static class Suggestions implements Serializable {
+  public final static class Suggestions implements Serializable {
 
     private final List<Suggestion> suggestions;
 

engine/runtime/src/main/java/org/enso/interpreter/runtime/Module.java

@@ -22,11 +22,11 @@ import java.util.Map;
 import java.util.UUID;
 import java.util.WeakHashMap;
 import java.util.logging.Level;
-import org.enso.compiler.ModuleCache;
 import org.enso.compiler.context.CompilerContext;
 import org.enso.compiler.context.SimpleUpdate;
 import org.enso.compiler.core.IR;
 import org.enso.compiler.core.ir.Expression;
+import org.enso.interpreter.caches.ModuleCache;
 import org.enso.interpreter.node.callable.dispatch.CallOptimiserNode;
 import org.enso.interpreter.node.callable.dispatch.LoopingCallOptimiserNode;
 import org.enso.interpreter.runtime.builtin.BuiltinFunction;

engine/runtime/src/main/java/org/enso/interpreter/runtime/TruffleCompilerContext.java

@@ -1,5 +1,6 @@
 package org.enso.interpreter.runtime;
 
+import org.enso.compiler.Passes;
 import org.enso.compiler.pass.analyse.BindingAnalysis$;
 import org.enso.compiler.context.CompilerContext;
 import org.enso.compiler.context.FreshNameSupply;
@@ -13,21 +14,17 @@ import java.io.PrintStream;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.concurrent.Future;
 import java.util.function.Consumer;
 import java.util.logging.Level;
 
-import org.enso.compiler.Cache;
 import org.enso.compiler.Compiler;
-import org.enso.compiler.ModuleCache;
 import org.enso.compiler.PackageRepository;
-import org.enso.compiler.Passes;
-import org.enso.compiler.SerializationManager;
-import org.enso.compiler.core.ir.Expression;
 import org.enso.compiler.data.BindingsMap;
 import org.enso.compiler.data.CompilerConfig;
-import org.enso.interpreter.node.ExpressionNode;
+import org.enso.editions.LibraryName;
-import org.enso.interpreter.runtime.data.Type;
+import org.enso.interpreter.caches.Cache;
-import org.enso.interpreter.runtime.scope.LocalScope;
+import org.enso.interpreter.caches.ModuleCache;
 import org.enso.interpreter.runtime.type.Types;
 import org.enso.pkg.Package;
 import org.enso.pkg.QualifiedName;
@@ -39,14 +36,16 @@ import scala.Option;
 final class TruffleCompilerContext implements CompilerContext {
 
   private final EnsoContext context;
-  private final TruffleLogger compiler;
+  private final TruffleLogger loggerCompiler;
-  private final TruffleLogger serializationManager;
+  private final TruffleLogger loggerManager;
   private final RuntimeStubsGenerator stubsGenerator;
+  private final SerializationManager serializationManager;
 
   TruffleCompilerContext(EnsoContext context) {
     this.context = context;
-    this.compiler = context.getLogger(Compiler.class);
+    this.loggerCompiler = context.getLogger(Compiler.class);
-    this.serializationManager = context.getLogger(SerializationManager.class);
+    this.loggerManager = context.getLogger(SerializationManager.class);
+    this.serializationManager = new SerializationManager(this);
     this.stubsGenerator = new RuntimeStubsGenerator(context.getBuiltins());
   }
 
@@ -70,6 +69,10 @@ final class TruffleCompilerContext implements CompilerContext {
     return context.getPackageRepository();
   }
 
+  final SerializationManager getSerializationManager() {
+    return serializationManager;
+  }
+
   @Override
   public PrintStream getErr() {
     return context.getErr();
@@ -82,12 +85,12 @@ final class TruffleCompilerContext implements CompilerContext {
 
   @Override
   public void log(Level level, String msg, Object... args) {
-    compiler.log(level, msg, args);
+    loggerCompiler.log(level, msg, args);
   }
 
   @Override
   public void logSerializationManager(Level level, String msg, Object... args) {
-    serializationManager.log(level, msg, args);
+    loggerManager.log(level, msg, args);
   }
 
   @Override
@@ -169,12 +172,10 @@ final class TruffleCompilerContext implements CompilerContext {
     }
   }
 
-  @Override
   public <T> Optional<T> loadCache(Cache<T, ?> cache) {
     return cache.load(context);
   }
 
-  @Override
   public <T> Optional<TruffleFile> saveCache(
           Cache<T, ?> cache, T entry, boolean useGlobalCacheLocations) {
     return cache.save(entry, context, useGlobalCacheLocations);
@@ -191,7 +192,8 @@ final class TruffleCompilerContext implements CompilerContext {
    */
   @Override
   public void initializeBuiltinsIr(
-          boolean irCachingEnabled, SerializationManager serializationManager,
+    Compiler compiler,
+    boolean irCachingEnabled,
     FreshNameSupply freshNameSupply, Passes passes
   ) {
     var builtins = context.getBuiltins();
@@ -207,7 +209,7 @@ final class TruffleCompilerContext implements CompilerContext {
 
       if (irCachingEnabled) {
         if (
-          serializationManager.deserialize(builtinsModule) instanceof Option<?> op &&
+          serializationManager.deserialize(compiler, builtinsModule) instanceof Option<?> op &&
           op.isDefined() &&
           op.get() instanceof Boolean b && b
         ) {
@@ -233,7 +235,7 @@ final class TruffleCompilerContext implements CompilerContext {
 
       if (irCachingEnabled && !wasLoadedFromCache(builtinsModule)) {
         serializationManager.serializeModule(
-                builtinsModule, true, true
+          compiler, builtinsModule, true, true
         );
       }
     }
@@ -250,6 +252,30 @@ final class TruffleCompilerContext implements CompilerContext {
     return option.isEmpty() ? null : option.get().asCompilerModule();
   }
 
+  @SuppressWarnings("unchecked")
+  @Override
+  public Future<Boolean> serializeLibrary(Compiler compiler, LibraryName libraryName, boolean useGlobalCacheLocations) {
+    Object res = serializationManager.serializeLibrary(compiler, libraryName, useGlobalCacheLocations);
+    return (Future<Boolean>) res;
+  }
+
+  @SuppressWarnings("unchecked")
+  @Override
+  public Future<Boolean> serializeModule(Compiler compiler, CompilerContext.Module module, boolean useGlobalCacheLocations) {
+    Object res = serializationManager.serializeModule(compiler, module, useGlobalCacheLocations, true);
+    return (Future<Boolean>) res;
+  }
+
+  @Override
+  public boolean deserializeModule(Compiler compiler, CompilerContext.Module module) {
+    var result = serializationManager.deserialize(compiler, module);
+    return result.nonEmpty();
+  }
+
+  @Override
+  public void shutdown(boolean waitForPendingJobCompletion) {
+    serializationManager.shutdown(waitForPendingJobCompletion);
+  }
 
   private final class ModuleUpdater implements Updater, AutoCloseable {
     private final Module module;
@@ -389,7 +415,6 @@ final class TruffleCompilerContext implements CompilerContext {
       return module.getDirectModulesRefs();
     }
 
-    @Override
     public ModuleCache getCache() {
       return module.getCache();
     }

engine/runtime/src/main/scala/org/enso/compiler/Compiler.scala

@@ -69,8 +69,6 @@ class Compiler(
   private val irCachingEnabled                 = !context.isIrCachingDisabled
   private val useGlobalCacheLocations          = context.isUseGlobalCacheLocations
   private val isInteractiveMode                = context.isInteractiveMode()
-  private val serializationManager: SerializationManager =
-    new SerializationManager(this)
   private val output: PrintStream =
     if (config.outputRedirect.isDefined)
       new PrintStream(config.outputRedirect.get)
@@ -105,18 +103,14 @@ class Compiler(
   /** Run the initialization sequence. */
   def initialize(): Unit = {
     context.initializeBuiltinsIr(
+      this,
       irCachingEnabled,
-      serializationManager,
       freshNameSupply,
       passes
     )
     packageRepository.initialize().left.foreach(reportPackageError)
   }
 
-  /** @return the serialization manager instance. */
-  def getSerializationManager: SerializationManager =
-    serializationManager
-
   /** @return the package repository instance. */
   def getPackageRepository(): PackageRepository =
     context.getPackageRepository
@@ -147,7 +141,7 @@ class Compiler(
   def compile(
     shouldCompileDependencies: Boolean,
     useGlobalCacheLocations: Boolean
-  ): Future[Boolean] = {
+  ): Future[java.lang.Boolean] = {
     getPackageRepository().getMainProjectPackage match {
       case None =>
         context.log(
@@ -188,9 +182,10 @@ class Compiler(
               shouldCompileDependencies
             )
 
-            serializationManager.serializeLibrary(
+            context.serializeLibrary(
+              this,
               pkg.libraryName,
-              useGlobalCacheLocations = useGlobalCacheLocations
+              useGlobalCacheLocations
             )
         }
     }
@@ -440,7 +435,8 @@ class Compiler(
             if (isInteractiveMode) {
               context.notifySerializeModule(context.getModuleName(module))
             } else {
-              serializationManager.serializeModule(
+              context.serializeModule(
+                this,
                 module,
                 useGlobalCacheLocations
               )
@@ -580,9 +576,8 @@ class Compiler(
     context.updateModule(module, _.resetScope)
 
     if (irCachingEnabled && !context.isInteractive(module)) {
-      serializationManager.deserialize(module) match {
+      if (context.deserializeModule(this, module)) {
-        case Some(_) => return
+        return
-        case _       =>
       }
     }
 
@@ -597,9 +592,9 @@ class Compiler(
   def importExportBindings(module: Module): Option[BindingsMap] = {
     if (irCachingEnabled && !context.isInteractive(module)) {
       val libraryName = Option(module.getPackage).map(_.libraryName)
-      libraryName
+      libraryName.flatMap(
-        .flatMap(packageRepository.getLibraryBindings(_, serializationManager))
+        packageRepository.getLibraryBindings(_, module.getName, context)
-        .flatMap(_.bindings.entries.get(module.getName))
+      )
     } else None
   }
 
@@ -1234,7 +1229,7 @@ class Compiler(
     *                                    jobs to complete
     */
   def shutdown(waitForPendingJobCompletion: Boolean): Unit = {
-    serializationManager.shutdown(waitForPendingJobCompletion)
+    context.shutdown(waitForPendingJobCompletion)
     shutdownParsingPool(waitForPendingJobCompletion)
   }
 

engine/runtime/src/main/scala/org/enso/compiler/PackageRepository.scala

@@ -3,7 +3,8 @@ package org.enso.compiler
 import com.oracle.truffle.api.TruffleFile
 import org.enso.editions.LibraryName
 import org.enso.compiler.context.CompilerContext
-import org.enso.pkg.{ComponentGroups, Package}
+import org.enso.compiler.data.BindingsMap
+import org.enso.pkg.{ComponentGroups, Package, QualifiedName}
 
 import scala.collection.immutable.ListSet
 import scala.jdk.OptionConverters.RichOption
@@ -110,8 +111,9 @@ trait PackageRepository {
   /** Returns a deserialized bindings map for the whole library, if available */
   def getLibraryBindings(
     libraryName: LibraryName,
-    serializationManager: SerializationManager
+    moduleName: QualifiedName,
-  ): Option[ImportExportCache.CachedBindings]
+    context: CompilerContext
+  ): Option[BindingsMap]
 
 }
 

engine/runtime/src/main/scala/org/enso/interpreter/runtime/DefaultPackageRepository.scala

@@ -1,13 +1,13 @@
 package org.enso.interpreter.runtime
 
 import org.enso.compiler.PackageRepository
-import org.enso.compiler.ImportExportCache
-import org.enso.compiler.SerializationManager
 import org.enso.compiler.context.CompilerContext
+import org.enso.compiler.data.BindingsMap
 import com.oracle.truffle.api.TruffleFile
 import com.typesafe.scalalogging.Logger
 import org.apache.commons.lang3.StringUtils
 import org.enso.editions.LibraryVersion
+import org.enso.interpreter.caches.ImportExportCache
 import org.enso.interpreter.runtime.util.TruffleFileSystem
 import org.enso.librarymanager.published.repository.LibraryManifest
 import org.enso.librarymanager.resolved.LibraryRoot
@@ -572,17 +572,21 @@ private class DefaultPackageRepository(
 
   override def getLibraryBindings(
     libraryName: LibraryName,
-    serializationManager: SerializationManager
+    moduleName: QualifiedName,
-  ): Option[ImportExportCache.CachedBindings] = {
+    context: CompilerContext
-    ensurePackageIsLoaded(libraryName).toOption.flatMap { _ =>
+  ): Option[BindingsMap] = {
+    val cache = ensurePackageIsLoaded(libraryName).toOption.flatMap { _ =>
       if (!loadedLibraryBindings.contains(libraryName)) {
         loadedPackages.get(libraryName).flatten.foreach(loadDependencies(_))
-        serializationManager
+        context
+          .asInstanceOf[TruffleCompilerContext]
+          .getSerializationManager()
           .deserializeLibraryBindings(libraryName)
           .foreach(cache => loadedLibraryBindings.addOne((libraryName, cache)))
       }
       loadedLibraryBindings.get(libraryName)
     }
+    cache.flatMap(_.bindings.entries().get(moduleName))
   }
 
   private def loadDependencies(pkg: Package[TruffleFile]): Unit = {

engine/runtime/src/main/scala/org/enso/interpreter/runtime/SerializationManager.scala

@@ -1,15 +1,20 @@
-package org.enso.compiler
+package org.enso.interpreter.runtime
 
 import com.oracle.truffle.api.source.Source
+import org.enso.compiler.Compiler
 import org.enso.compiler.core.Implicits.AsMetadata
 import org.enso.compiler.core.ir.{Module => IRModule}
 import org.enso.compiler.context.{ExportsBuilder, ExportsMap, SuggestionBuilder}
+import org.enso.compiler.context.CompilerContext
 import org.enso.compiler.context.CompilerContext.Module
 import org.enso.compiler.pass.analyse.BindingAnalysis
 import org.enso.editions.LibraryName
 import org.enso.pkg.QualifiedName
 import org.enso.polyglot.Suggestion
 import org.enso.polyglot.CompilationStage
+import org.enso.interpreter.caches.ImportExportCache
+import org.enso.interpreter.caches.ModuleCache
+import org.enso.interpreter.caches.SuggestionsCache
 
 import java.io.NotSerializableException
 import java.util
@@ -27,7 +32,11 @@ import java.util.logging.Level
 import scala.collection.mutable
 import scala.jdk.OptionConverters.RichOptional
 
-final class SerializationManager(compiler: Compiler) {
+final class SerializationManager(private val context: TruffleCompilerContext) {
+
+  def this(compiler: Compiler) = {
+    this(compiler.context.asInstanceOf[TruffleCompilerContext])
+  }
 
   import SerializationManager._
 
@@ -57,13 +66,13 @@ final class SerializationManager(compiler: Compiler) {
     TimeUnit.SECONDS,
     new LinkedBlockingDeque[Runnable](),
     (runnable: Runnable) => {
-      compiler.context.createSystemThread(runnable)
+      context.createSystemThread(runnable)
     }
   )
 
   // Make sure it is started to avoid races with language shutdown with low job
   // count.
-  if (compiler.context.isCreateThreadAllowed) {
+  if (context.isCreateThreadAllowed) {
     pool.prestartAllCoreThreads()
   }
 
@@ -91,6 +100,7 @@ final class SerializationManager(compiler: Compiler) {
     *         `true` if `module` has been successfully serialized, `false` otherwise
     */
   def serializeModule(
+    compiler: Compiler,
     module: Module,
     useGlobalCacheLocations: Boolean,
     useThreadPool: Boolean = true
@@ -100,7 +110,7 @@ final class SerializationManager(compiler: Compiler) {
         "Cannot serialize synthetic module [" + module.getName + "]"
       );
     }
-    compiler.context.logSerializationManager(
+    context.logSerializationManager(
       debugLogLevel,
       "Requesting serialization for module [{0}].",
       module.getName
@@ -112,7 +122,7 @@ final class SerializationManager(compiler: Compiler) {
     duplicatedIr.preorder.foreach(_.passData.prepareForSerialization(compiler))
 
     val task = doSerializeModule(
-      module.getCache,
+      getCache(module),
       duplicatedIr,
       module.getCompilationStage,
       module.getName,
@@ -130,7 +140,7 @@ final class SerializationManager(compiler: Compiler) {
         CompletableFuture.completedFuture(task.call())
       } catch {
         case e: Throwable =>
-          compiler.context.logSerializationManager(
+          context.logSerializationManager(
             debugLogLevel,
             s"Serialization task failed in module [${module.getName}].",
             e
@@ -141,18 +151,19 @@ final class SerializationManager(compiler: Compiler) {
   }
 
   def serializeLibrary(
+    compiler: Compiler,
     libraryName: LibraryName,
     useGlobalCacheLocations: Boolean
   ): Future[Boolean] = {
-    compiler.context.logSerializationManager(
+    context.logSerializationManager(
       Level.INFO,
       "Requesting serialization for library [{0}].",
       libraryName
     )
 
     val task: Callable[Boolean] =
-      doSerializeLibrary(libraryName, useGlobalCacheLocations)
+      doSerializeLibrary(compiler, libraryName, useGlobalCacheLocations)
-    if (compiler.context.isCreateThreadAllowed) {
+    if (context.isCreateThreadAllowed) {
       isWaitingForSerialization.synchronized {
         val future = pool.submit(task)
         isWaitingForSerialization.put(libraryName.toQualifiedName, future)
@@ -163,7 +174,7 @@ final class SerializationManager(compiler: Compiler) {
         CompletableFuture.completedFuture(task.call())
       } catch {
         case e: Throwable =>
-          compiler.context.logSerializationManager(
+          context.logSerializationManager(
             debugLogLevel,
             s"Serialization task failed for library [$libraryName].",
             e
@@ -174,6 +185,7 @@ final class SerializationManager(compiler: Compiler) {
   }
 
   private def doSerializeLibrary(
+    compiler: Compiler,
     libraryName: LibraryName,
     useGlobalCacheLocations: Boolean
   ): Callable[Boolean] = () => {
@@ -181,7 +193,7 @@ final class SerializationManager(compiler: Compiler) {
       Thread.sleep(100)
     }
 
-    compiler.context.logSerializationManager(
+    context.logSerializationManager(
       debugLogLevel,
       "Running serialization for bindings [{0}].",
       libraryName
@@ -190,7 +202,8 @@ final class SerializationManager(compiler: Compiler) {
     val bindingsCache = new ImportExportCache.CachedBindings(
       libraryName,
       new ImportExportCache.MapToBindings(
-        compiler.packageRepository
+        context
+          .getPackageRepository()
           .getModulesForLibrary(libraryName)
           .map { module =>
             val ir = module.getIr
@@ -203,7 +216,8 @@ final class SerializationManager(compiler: Compiler) {
           }
           .toMap
       ),
-      compiler.packageRepository
+      context
+        .getPackageRepository()
         .getPackageForLibraryJava(libraryName)
         .map(_.listSourcesJava())
     )
@@ -211,7 +225,7 @@ final class SerializationManager(compiler: Compiler) {
       val result =
         try {
           val cache = new ImportExportCache(libraryName)
-          val file = compiler.context.saveCache(
+          val file = context.saveCache(
             cache,
             bindingsCache,
             useGlobalCacheLocations
@@ -219,14 +233,14 @@ final class SerializationManager(compiler: Compiler) {
           file.isPresent
         } catch {
           case e: NotSerializableException =>
-            compiler.context.logSerializationManager(
+            context.logSerializationManager(
               Level.SEVERE,
               s"Could not serialize bindings [$libraryName].",
               e
             )
             throw e
           case e: Throwable =>
-            compiler.context.logSerializationManager(
+            context.logSerializationManager(
               Level.SEVERE,
               s"Serialization of bindings `$libraryName` failed: ${e.getMessage}`",
               e
@@ -234,7 +248,11 @@ final class SerializationManager(compiler: Compiler) {
             throw e
         }
 
-      doSerializeLibrarySuggestions(libraryName, useGlobalCacheLocations)
+      doSerializeLibrarySuggestions(
+        compiler,
+        libraryName,
+        useGlobalCacheLocations
+      )
 
       result
     } finally {
@@ -243,6 +261,7 @@ final class SerializationManager(compiler: Compiler) {
   }
 
   private def doSerializeLibrarySuggestions(
+    compiler: Compiler,
     libraryName: LibraryName,
     useGlobalCacheLocations: Boolean
   ): Boolean = {
@@ -252,7 +271,7 @@ final class SerializationManager(compiler: Compiler) {
 
     try {
       val libraryModules =
-        compiler.packageRepository.getModulesForLibrary(libraryName)
+        context.getPackageRepository().getModulesForLibrary(libraryName)
       libraryModules
         .flatMap { module =>
           val suggestions = SuggestionBuilder(module, compiler)
@@ -272,12 +291,13 @@ final class SerializationManager(compiler: Compiler) {
         new SuggestionsCache.CachedSuggestions(
           libraryName,
           new SuggestionsCache.Suggestions(suggestions),
-          compiler.packageRepository
+          context
+            .getPackageRepository()
             .getPackageForLibraryJava(libraryName)
             .map(_.listSourcesJava())
         )
       val cache = new SuggestionsCache(libraryName)
-      val file = compiler.context.saveCache(
+      val file = context.saveCache(
         cache,
         cachedSuggestions,
         useGlobalCacheLocations
@@ -285,14 +305,14 @@ final class SerializationManager(compiler: Compiler) {
       file.isPresent
     } catch {
       case e: NotSerializableException =>
-        compiler.context.logSerializationManager(
+        context.logSerializationManager(
           Level.SEVERE,
           s"Could not serialize suggestions [$libraryName].",
           e
         )
         throw e
       case e: Throwable =>
-        compiler.context.logSerializationManager(
+        context.logSerializationManager(
           Level.SEVERE,
           s"Serialization of suggestions `$libraryName` failed: ${e.getMessage}`",
           e
@@ -312,16 +332,16 @@ final class SerializationManager(compiler: Compiler) {
         Thread.sleep(100)
       }
       val cache = new SuggestionsCache(libraryName)
-      compiler.context.loadCache(cache).toScala match {
+      context.loadCache(cache).toScala match {
         case result @ Some(_: SuggestionsCache.CachedSuggestions) =>
-          compiler.context.logSerializationManager(
+          context.logSerializationManager(
             Level.FINE,
             "Restored suggestions for library [{0}].",
             libraryName
           )
           result
         case None =>
-          compiler.context.logSerializationManager(
+          context.logSerializationManager(
             Level.FINE,
             "Unable to load suggestions for library [{0}].",
             libraryName
@@ -342,16 +362,16 @@ final class SerializationManager(compiler: Compiler) {
         Thread.sleep(100)
       }
       val cache = new ImportExportCache(libraryName)
-      compiler.context.loadCache(cache).toScala match {
+      context.loadCache(cache).toScala match {
         case result @ Some(_: ImportExportCache.CachedBindings) =>
-          compiler.context.logSerializationManager(
+          context.logSerializationManager(
             Level.FINE,
             "Restored bindings for library [{0}].",
             libraryName
           )
           result
         case _ =>
-          compiler.context.logSerializationManager(
+          context.logSerializationManager(
             Level.FINEST,
             "Unable to load bindings for library [{0}].",
             libraryName
@@ -375,6 +395,7 @@ final class SerializationManager(compiler: Compiler) {
     *         cache could not be deserialized.
     */
   def deserialize(
+    compiler: Compiler,
     module: Module
   ): Option[Boolean] = {
     if (isWaitingForSerialization(module)) {
@@ -385,47 +406,45 @@ final class SerializationManager(compiler: Compiler) {
         Thread.sleep(100)
       }
 
-      compiler.context.loadCache(module.getCache).toScala match {
+      context.loadCache(getCache(module)).toScala match {
         case Some(loadedCache) =>
           val relinkedIrChecks =
-            loadedCache
+            loadedCache.moduleIR.preorder
-              .moduleIR()
+              .map(_.passData.restoreFromSerialization(compiler))
-              .preorder
+          context.updateModule(
-              .map(_.passData.restoreFromSerialization(this.compiler))
-          compiler.context.updateModule(
             module,
             { u =>
-              u.ir(loadedCache.moduleIR())
+              u.ir(loadedCache.moduleIR)
-              u.compilationStage(loadedCache.compilationStage())
+              u.compilationStage(loadedCache.compilationStage)
               u.loadedFromCache(true)
             }
           )
-          compiler.context.logSerializationManager(
+          context.logSerializationManager(
             debugLogLevel,
             "Restored IR from cache for module [{0}] at stage [{1}].",
             module.getName,
-            loadedCache.compilationStage()
+            loadedCache.compilationStage
           )
 
           if (!relinkedIrChecks.contains(false)) {
-            compiler.context.updateModule(module, _.hasCrossModuleLinks(true))
+            context.updateModule(module, _.hasCrossModuleLinks(true))
-            compiler.context.logSerializationManager(
+            context.logSerializationManager(
               debugLogLevel,
               "Restored links (early phase) in module [{0}].",
               module.getName
             )
             Some(true)
           } else {
-            compiler.context.logSerializationManager(
+            context.logSerializationManager(
               debugLogLevel,
               "Could not restore links (early phase) in module [{0}].",
               module.getName
             )
-            compiler.context.updateModule(module, _.hasCrossModuleLinks(false))
+            context.updateModule(module, _.hasCrossModuleLinks(false))
             Some(false)
           }
         case None =>
-          compiler.context.logSerializationManager(
+          context.logSerializationManager(
             debugLogLevel,
             "Unable to load a cache for module [{0}].",
             module.getName
@@ -524,7 +543,7 @@ final class SerializationManager(compiler: Compiler) {
           isWaitingForSerialization.size
         }
         val jobCount = waitingCount + isSerializing.size
-        compiler.context.logSerializationManager(
+        context.logSerializationManager(
           debugLogLevel,
           "Waiting for #{0} serialization jobs to complete.",
           jobCount
@@ -551,7 +570,7 @@ final class SerializationManager(compiler: Compiler) {
 
       pool.shutdownNow()
       Thread.sleep(100)
-      compiler.context.logSerializationManager(
+      context.logSerializationManager(
         debugLogLevel,
         "Serialization manager has been shut down."
       )
@@ -591,7 +610,7 @@ final class SerializationManager(compiler: Compiler) {
       Thread.sleep(100)
     }
 
-    compiler.context.logSerializationManager(
+    context.logSerializationManager(
       debugLogLevel,
       "Running serialization for module [{0}].",
       name
@@ -602,7 +621,7 @@ final class SerializationManager(compiler: Compiler) {
         if (stage.isAtLeast(CompilationStage.AFTER_STATIC_PASSES)) {
           CompilationStage.AFTER_STATIC_PASSES
         } else stage
-      compiler.context
+      context
         .saveCache(
           cache,
           new ModuleCache.CachedModule(ir, fixedStage, source),
@@ -612,14 +631,14 @@ final class SerializationManager(compiler: Compiler) {
         .orElse(false)
     } catch {
       case e: NotSerializableException =>
-        compiler.context.logSerializationManager(
+        context.logSerializationManager(
           Level.SEVERE,
           s"Could not serialize module [$name].",
           e
         )
         throw e
       case e: Throwable =>
-        compiler.context.logSerializationManager(
+        context.logSerializationManager(
           Level.SEVERE,
           s"Serialization of module `$name` failed: ${e.getMessage}`",
           e
@@ -663,6 +682,10 @@ final class SerializationManager(compiler: Compiler) {
       )
       .asScala
   }
+
+  private def getCache(module: Module): ModuleCache = {
+    module.asInstanceOf[TruffleCompilerContext.Module].getCache
+  }
 }
 
 object SerializationManager {
@@ -682,4 +705,7 @@ object SerializationManager {
       QualifiedName(List(libraryName.namespace), libraryName.name)
   }
 
+  def apply(context: CompilerContext): SerializationManager = {
+    context.asInstanceOf[TruffleCompilerContext].getSerializationManager()
+  }
 }

engine/runtime/src/test/java/org/enso/compiler/SerdeCompilerTest.java

@@ -18,6 +18,7 @@ import java.util.logging.LogRecord;
 import java.util.logging.SimpleFormatter;
 import org.enso.compiler.core.ir.Module;
 import org.enso.interpreter.runtime.EnsoContext;
+import org.enso.interpreter.runtime.SerializationManager$;
 import org.enso.pkg.PackageManager;
 import org.enso.polyglot.LanguageInfo;
 import org.enso.polyglot.MethodNames;
@@ -54,13 +55,14 @@ public class SerdeCompilerTest {
       var result = compiler.run(module);
       assertEquals("Two library modules are compiled", result.compiledModules().size(), 2);
       assertEquals(result.compiledModules().exists(m -> m == module), true);
-      var serializationManager = ensoContext.getCompiler().getSerializationManager();
+      var serializationManager =
+          SerializationManager$.MODULE$.apply(ensoContext.getCompiler().context());
       var futures = new ArrayList<Future<?>>();
       result
           .compiledModules()
           .foreach(
               (m) -> {
-                var future = serializationManager.serializeModule(m, true, true);
+                var future = serializationManager.serializeModule(compiler, m, true, true);
                 futures.add(future);
                 return null;
               });

engine/runtime/src/test/java/org/enso/compiler/SerializationManagerTest.java

@@ -10,7 +10,9 @@ import java.util.function.Supplier;
 import java.util.stream.Stream;
 import org.apache.commons.io.FileUtils;
 import org.enso.editions.LibraryName;
+import org.enso.interpreter.caches.SuggestionsCache;
 import org.enso.interpreter.runtime.EnsoContext;
+import org.enso.interpreter.runtime.SerializationManager;
 import org.enso.interpreter.runtime.util.TruffleFileSystem;
 import org.enso.interpreter.test.InterpreterContext;
 import org.enso.pkg.Package;

engine/runtime/src/test/java/org/enso/compiler/SerializerTest.java

@@ -9,6 +9,7 @@ import java.io.OutputStream;
 import java.nio.file.Paths;
 import java.util.concurrent.TimeUnit;
 import org.enso.interpreter.runtime.EnsoContext;
+import org.enso.interpreter.runtime.SerializationManager;
 import org.enso.pkg.PackageManager;
 import org.enso.polyglot.LanguageInfo;
 import org.enso.polyglot.MethodNames;
@@ -57,10 +58,10 @@ public class SerializerTest {
     assertEquals(result.compiledModules().exists(m -> m == module), true);
     var serializationManager = new SerializationManager(ensoContext.getCompiler());
     var useThreadPool = compiler.context().isCreateThreadAllowed();
-    var future = serializationManager.serializeModule(module, true, useThreadPool);
+    var future = serializationManager.serializeModule(compiler, module, true, useThreadPool);
     var serialized = future.get(5, TimeUnit.SECONDS);
     assertEquals(serialized, true);
-    var deserialized = serializationManager.deserialize(module);
+    var deserialized = serializationManager.deserialize(compiler, module);
     assertEquals(deserialized.isDefined() && (Boolean) deserialized.get(), true);
     serializationManager.shutdown(true);
     ctx.leave();