Add search/import command (#1310)

Add `search/import` request returning the info
required for module import.

docs/language-server/protocol-language-server.md

@@ -34,6 +34,7 @@ transport formats, please look [here](./protocol-architecture).
   - [`FieldUpdate`](#fieldupdate)
   - [`SuggestionArgumentUpdate`](#suggestionargumentupdate)
   - [`SuggestionsDatabaseUpdate`](#suggestionsdatabaseupdate)
+  - [`Export`](#export)
   - [`File`](#file)
   - [`DirectoryTree`](#directorytree)
   - [`FileAttributes`](#fileattributes)
@@ -118,6 +119,7 @@ transport formats, please look [here](./protocol-architecture).
   - [`search/getSuggestionsDatabaseVersion`](#searchgetsuggestionsdatabaseversion)
   - [`search/suggestionsDatabaseUpdate`](#searchsuggestionsdatabaseupdate)
   - [`search/completion`](#searchcompletion)
+  - [`search/import`](#searchimport)
 - [Input/Output Operations](#input-output-operations)
   - [`io/redirectStandardOutput`](#ioredirectstdardoutput)
   - [`io/suppressStandardOutput`](#iosuppressstdardoutput)
@@ -537,6 +539,37 @@ interface Modify {
 }
 ```
 
+### `Export`
+
+The info about module re-export.
+
+#### Format
+
+```typescript
+type Export = Qualified | Unqualified;
+
+interface Qualified {
+  /**
+   * The module that re-exports the given module.
+   */
+  module: String;
+
+  /**
+   * The new name of the given module if it was renamed in the export clause.
+   *
+   * I.e. `X` in `export A.B as X`.
+   */
+  alias?: String;
+}
+
+interface Unqualified {
+  /**
+   * The module name that re-exports the given module.
+   */
+  module: String;
+}
+```
+
 ### `File`
 
 A representation of a file on disk.
@@ -3038,6 +3071,56 @@ Sent from client to the server to receive the autocomplete suggestion.
 - [`ModuleNameNotResolvedError`](#modulenamenotresolvederror) the module name
   cannot be extracted from the provided file path parameter
 
+### `search/import`
+
+Sent from client to the server to receive the information required for module
+import.
+
+- **Type:** Request
+- **Direction:** Client -> Server
+- **Connection:** Protocol
+- **Visibility:** Public
+
+#### Parameters
+
+```typescript
+{
+  /**
+   * The id of suggestion to import.
+   */
+  id: SuggestionId;
+}
+```
+
+#### Result
+
+```typescript
+{
+  /**
+   * The definition module of the suggestion.
+   */
+  module: String;
+
+  /**
+   * The name of the resolved suggestion.
+   */
+  symbol: String;
+
+  /**
+   * The list of modules that re-export the suggestion. Modules are ordered
+   * from the least to most nested.
+   */
+  exports: Export[];
+}
+```
+
+#### Errors
+
+- [`SuggestionsDatabaseError`](#suggestionsdatabaseerror) an error accessing the
+  suggestions database
+- [`SuggestionNotFoundError`](#suggestionnotfounderror) the requested suggestion
+  was not found in the suggestions database
+
 ## Input/Output Operations
 
 The input/output portion of the language server API deals with redirecting
@@ -3513,3 +3596,14 @@ Signals that the module name can not be resolved for the given file.
   "message" : "Module name can't be resolved for the given file"
 }
 ```
+
+### `SuggestionNotFoundError`
+
+Signals that the requested suggestion was not found.
+
+```typescript
+"error" : {
+  "code" : 7004,
+  "message" : "Requested suggestion was not found"
+}
+```

engine/language-server/src/main/scala/org/enso/languageserver/protocol/json/JsonConnectionController.scala

@@ -41,6 +41,7 @@ import org.enso.languageserver.search.SearchApi.{
   Completion,
   GetSuggestionsDatabase,
   GetSuggestionsDatabaseVersion,
+  Import,
   InvalidateSuggestionsDatabase
 }
 import org.enso.languageserver.runtime.VisualisationApi.{
@@ -287,6 +288,7 @@ class JsonConnectionController(
         .props(requestTimeout, suggestionsHandler),
       Completion -> search.CompletionHandler
         .props(requestTimeout, suggestionsHandler),
+      Import -> search.ImportHandler.props(requestTimeout, suggestionsHandler),
       AttachVisualisation -> AttachVisualisationHandler
         .props(rpcSession.clientId, requestTimeout, contextRegistry),
       DetachVisualisation -> DetachVisualisationHandler

engine/language-server/src/main/scala/org/enso/languageserver/protocol/json/JsonRpc.scala

@@ -58,6 +58,7 @@ object JsonRpc {
     .registerRequest(GetSuggestionsDatabaseVersion)
     .registerRequest(InvalidateSuggestionsDatabase)
     .registerRequest(Completion)
+    .registerRequest(Import)
     .registerRequest(RenameProject)
     .registerNotification(ForceReleaseCapability)
     .registerNotification(GrantCapability)

engine/language-server/src/main/scala/org/enso/languageserver/requesthandler/search/ImportHandler.scala

@@ -0,0 +1,79 @@
+package org.enso.languageserver.requesthandler.search
+
+import akka.actor.{Actor, ActorLogging, ActorRef, Cancellable, Props, Status}
+import org.enso.jsonrpc.Errors.ServiceError
+import org.enso.jsonrpc._
+import org.enso.languageserver.requesthandler.RequestTimeout
+import org.enso.languageserver.search.SearchApi.{
+  Import,
+  SuggestionsDatabaseError
+}
+import org.enso.languageserver.search.{SearchFailureMapper, SearchProtocol}
+import org.enso.languageserver.util.UnhandledLogging
+
+import scala.concurrent.duration.FiniteDuration
+
+/** A request handler for `search/import` command.
+  *
+  * @param timeout request timeout
+  * @param suggestionsHandler a reference to the suggestions handler
+  */
+class ImportHandler(
+  timeout: FiniteDuration,
+  suggestionsHandler: ActorRef
+) extends Actor
+    with ActorLogging
+    with UnhandledLogging {
+
+  import context.dispatcher
+
+  override def receive: Receive = requestStage
+
+  private def requestStage: Receive = {
+    case Request(Import, id, Import.Params(suggestionId)) =>
+      suggestionsHandler ! SearchProtocol.Import(suggestionId)
+      val cancellable =
+        context.system.scheduler.scheduleOnce(timeout, self, RequestTimeout)
+      context.become(responseStage(id, sender(), cancellable))
+  }
+
+  private def responseStage(
+    id: Id,
+    replyTo: ActorRef,
+    cancellable: Cancellable
+  ): Receive = {
+    case Status.Failure(ex) =>
+      log.error(ex, "Search import error")
+      replyTo ! ResponseError(Some(id), SuggestionsDatabaseError)
+      cancellable.cancel()
+      context.stop(self)
+
+    case RequestTimeout =>
+      log.error(s"Request $id timed out")
+      replyTo ! ResponseError(Some(id), ServiceError)
+      context.stop(self)
+
+    case msg: SearchProtocol.SearchFailure =>
+      replyTo ! ResponseError(Some(id), SearchFailureMapper.mapFailure(msg))
+
+    case SearchProtocol.ImportResult(module, symbol, exports) =>
+      replyTo ! ResponseResult(
+        Import,
+        id,
+        Import.Result(module, symbol, exports)
+      )
+      cancellable.cancel()
+      context.stop(self)
+  }
+}
+
+object ImportHandler {
+
+  /** Creates configuration object used to create a [[ImportHandler]].
+    *
+    * @param timeout request timeout
+    * @param suggestionsHandler a reference to the suggestions handler
+    */
+  def props(timeout: FiniteDuration, suggestionsHandler: ActorRef): Props =
+    Props(new ImportHandler(timeout, suggestionsHandler))
+}

engine/language-server/src/main/scala/org/enso/languageserver/search/SearchApi.scala

@@ -3,12 +3,12 @@ package org.enso.languageserver.search
 import org.enso.jsonrpc.{Error, HasParams, HasResult, Method, Unused}
 import org.enso.languageserver.filemanager.Path
 import org.enso.languageserver.search.SearchProtocol.{
+  Export,
   SuggestionDatabaseEntry,
   SuggestionId,
   SuggestionKind,
   SuggestionsDatabaseUpdate
 }
-
 import org.enso.text.editing.model.Position
 
 /** The execution JSON RPC API provided by the language server.
@@ -90,6 +90,20 @@ object SearchApi {
     }
   }
 
+  case object Import extends Method("search/import") {
+
+    case class Params(id: Long)
+
+    case class Result(module: String, symbol: String, exports: Seq[Export])
+
+    implicit val hasParams = new HasParams[this.type] {
+      type Params = Import.Params
+    }
+    implicit val hasResult = new HasResult[this.type] {
+      type Result = Import.Result
+    }
+  }
+
   case object SuggestionsDatabaseError
       extends Error(7001, "Suggestions database error")
 
@@ -98,4 +112,7 @@ object SearchApi {
 
   case object ModuleNameNotResolvedError
       extends Error(7003, "Module name can't be resolved for the given file")
+
+  case object SuggestionNotFoundError
+      extends Error(7004, "Requested suggestion was not found")
 }

engine/language-server/src/main/scala/org/enso/languageserver/search/SearchFailureMapper.scala

@@ -6,7 +6,8 @@ import org.enso.languageserver.search.SearchProtocol.{
   FileSystemError,
   ModuleNameNotResolvedError,
   ProjectNotFoundError,
-  SearchFailure
+  SearchFailure,
+  SuggestionNotFoundError
 }
 
 object SearchFailureMapper {
@@ -21,6 +22,7 @@ object SearchFailureMapper {
       case FileSystemError(e)            => FileSystemFailureMapper.mapFailure(e)
       case ProjectNotFoundError          => SearchApi.ProjectNotFoundError
       case ModuleNameNotResolvedError(_) => SearchApi.ModuleNameNotResolvedError
+      case SuggestionNotFoundError       => SearchApi.SuggestionNotFoundError
     }
 
 }

engine/language-server/src/main/scala/org/enso/languageserver/search/SearchProtocol.scala

@@ -391,6 +391,86 @@ object SearchProtocol {
     */
   case class CompletionResult(currentVersion: Long, results: Seq[SuggestionId])
 
+  /** The request returning the info about the suggestion import.
+    *
+    * @param id the requested suggestion id
+    */
+  case class Import(id: SuggestionId)
+
+  /** The request returning the info about the suggestion import.
+    *
+    * @param suggestion the requested suggestion
+    */
+  case class ImportSuggestion(suggestion: Suggestion)
+
+  /** Base trait for export statements. */
+  sealed trait Export {
+    def module: String
+  }
+  object Export {
+
+    /** Qualified module re-export.
+      *
+      * @param module the module name that exports the given module
+      * @param alias new module name if the module was renamed in the export
+      * clause
+      */
+    case class Qualified(module: String, alias: Option[String]) extends Export
+
+    /** Unqualified module export.
+      *
+      * @param module the module name that exports the given module
+      */
+    case class Unqualified(module: String) extends Export
+
+    private object CodecType {
+
+      val Qualified = "Qualified"
+
+      val Unqualified = "Unqualified"
+    }
+
+    implicit val encoder: Encoder[Export] =
+      Encoder.instance {
+        case qualified: Qualified =>
+          Encoder[Export.Qualified]
+            .apply(qualified)
+            .deepMerge(Json.obj(CodecField.Type -> CodecType.Qualified.asJson))
+            .dropNullValues
+
+        case unqualified: Unqualified =>
+          Encoder[Export.Unqualified]
+            .apply(unqualified)
+            .deepMerge(
+              Json.obj(CodecField.Type -> CodecType.Unqualified.asJson)
+            )
+            .dropNullValues
+      }
+
+    implicit val decoder: Decoder[Export] =
+      Decoder.instance { cursor =>
+        cursor.downField(CodecField.Type).as[String].flatMap {
+          case CodecType.Qualified =>
+            Decoder[Export.Qualified].tryDecode(cursor)
+
+          case CodecType.Unqualified =>
+            Decoder[Export.Unqualified].tryDecode(cursor)
+        }
+      }
+  }
+
+  /** The result of the import request.
+    *
+    * @param module the definition module of the symbol
+    * @param symbol the resolved symbol
+    * @param exports the list of re-exports
+    */
+  case class ImportResult(
+    module: String,
+    symbol: String,
+    exports: Seq[Export]
+  )
+
   /** The request to invalidate the modules index. */
   case object InvalidateModulesIndex
 
@@ -409,6 +489,9 @@ object SearchProtocol {
   /** Signals that the project not found in the root directory. */
   case object ProjectNotFoundError extends SearchFailure
 
+  /** Signals that the requested suggestion was not found. */
+  case object SuggestionNotFoundError extends SearchFailure
+
   /** Signals that the module name can not be resolved for the given file.
     *
     * @param file the file path

engine/language-server/src/main/scala/org/enso/languageserver/search/SuggestionsHandler.scala

@@ -18,7 +18,10 @@ import org.enso.languageserver.event.InitializedEvent
 import org.enso.languageserver.filemanager.{FileDeletedEvent, Path}
 import org.enso.languageserver.refactoring.ProjectNameChangedEvent
 import org.enso.languageserver.search.SearchProtocol._
-import org.enso.languageserver.search.handler.InvalidateModulesIndexHandler
+import org.enso.languageserver.search.handler.{
+  ImportModuleHandler,
+  InvalidateModulesIndexHandler
+}
 import org.enso.languageserver.session.SessionRouter.DeliverToJsonController
 import org.enso.languageserver.util.UnhandledLogging
 import org.enso.pkg.PackageManager
@@ -229,6 +232,17 @@ final class SuggestionsHandler(
         )
         .pipeTo(sender())
 
+    case Import(suggestionId) =>
+      val action = for {
+        result <- suggestionsRepo.select(suggestionId)
+      } yield result
+        .map(SearchProtocol.ImportSuggestion)
+        .getOrElse(SearchProtocol.SuggestionNotFoundError)
+
+      val handler = context.system
+        .actorOf(ImportModuleHandler.props(timeout, runtimeConnector))
+      action.pipeTo(handler)(sender())
+
     case FileDeletedEvent(path) =>
       getModuleName(projectName, path)
         .fold(

engine/language-server/src/main/scala/org/enso/languageserver/search/handler/ImportModuleHandler.scala

@@ -0,0 +1,86 @@
+package org.enso.languageserver.search.handler
+
+import java.util.UUID
+
+import akka.actor.{Actor, ActorLogging, ActorRef, Cancellable, Props}
+import org.enso.languageserver.requesthandler.RequestTimeout
+import org.enso.languageserver.runtime.RuntimeFailureMapper
+import org.enso.languageserver.search.SearchProtocol
+import org.enso.languageserver.util.UnhandledLogging
+import org.enso.polyglot.runtime.Runtime.Api
+
+import scala.concurrent.duration.FiniteDuration
+
+/** A request handler for import module command.
+  *
+  * @param timeout request timeout
+  * @param runtime reference to the runtime connector
+  */
+final class ImportModuleHandler(
+  timeout: FiniteDuration,
+  runtime: ActorRef
+) extends Actor
+    with ActorLogging
+    with UnhandledLogging {
+
+  import context.dispatcher
+
+  override def receive: Receive = requestStage
+
+  private def requestStage: Receive = {
+    case SearchProtocol.ImportSuggestion(suggestion) =>
+      runtime ! Api.Request(
+        UUID.randomUUID(),
+        Api.ImportSuggestionRequest(suggestion)
+      )
+      val cancellable =
+        context.system.scheduler.scheduleOnce(timeout, self, RequestTimeout)
+      context.become(responseStage(sender(), cancellable))
+
+    case msg: SearchProtocol.SearchFailure =>
+      sender() ! msg
+      context.stop(self)
+  }
+
+  private def responseStage(
+    replyTo: ActorRef,
+    cancellable: Cancellable
+  ): Receive = {
+    case RequestTimeout =>
+      replyTo ! RequestTimeout
+      context.stop(self)
+
+    case Api.Response(_, Api.ImportSuggestionResponse(module, sym, exports)) =>
+      replyTo ! SearchProtocol.ImportResult(
+        module,
+        sym,
+        exports.map(toSearchExport)
+      )
+      cancellable.cancel()
+      context.stop(self)
+
+    case Api.Response(_, error: Api.Error) =>
+      replyTo ! RuntimeFailureMapper.mapApiError(error)
+      cancellable.cancel()
+      context.stop(self)
+  }
+
+  private def toSearchExport(export: Api.Export): SearchProtocol.Export =
+    export match {
+      case Api.Export.Unqualified(module) =>
+        SearchProtocol.Export.Unqualified(module)
+      case Api.Export.Qualified(module, alias) =>
+        SearchProtocol.Export.Qualified(module, alias)
+    }
+}
+
+object ImportModuleHandler {
+
+  /** Creates a configuration object used to create [[ImportModuleHandler]].
+    *
+    * @param timeout request timeout
+    * @param runtime reference to the runtime conector
+    */
+  def props(timeout: FiniteDuration, runtime: ActorRef): Props =
+    Props(new ImportModuleHandler(timeout, runtime))
+}

engine/polyglot-api/src/main/scala/org/enso/polyglot/runtime/Runtime.scala

@@ -180,6 +180,14 @@ object Runtime {
       new JsonSubTypes.Type(
         value = classOf[Api.InvalidateModulesIndexResponse],
         name  = "invalidateModulesIndexResponse"
+      ),
+      new JsonSubTypes.Type(
+        value = classOf[Api.ImportSuggestionRequest],
+        name  = "importSuggestionRequest"
+      ),
+      new JsonSubTypes.Type(
+        value = classOf[Api.ImportSuggestionResponse],
+        name  = "importSuggestionResponse"
       )
     )
   )
@@ -561,6 +569,39 @@ object Runtime {
 
     }
 
+    @JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "type")
+    @JsonSubTypes(
+      Array(
+        new JsonSubTypes.Type(
+          value = classOf[Export.Qualified],
+          name  = "exportQualified"
+        ),
+        new JsonSubTypes.Type(
+          value = classOf[Export.Unqualified],
+          name  = "exportUnqualified"
+        )
+      )
+    )
+    sealed trait Export {
+      def module: String
+    }
+    object Export {
+
+      /** Qualified module re-export.
+        *
+        * @param module the module name that exports the given module
+        * @param alias new module name if the module was renamed in the export
+        * clause
+        */
+      case class Qualified(module: String, alias: Option[String]) extends Export
+
+      /** Unqualified module export.
+        *
+        * @param module the module name that exports the given module
+        */
+      case class Unqualified(module: String) extends Export
+    }
+
     /** The notification about the execution status.
       *
       * @param contextId the context's id
@@ -900,6 +941,25 @@ object Runtime {
     /** Signals that the module indexes has been invalidated. */
     case class InvalidateModulesIndexResponse() extends ApiResponse
 
+    /** A request to return info needed to import the suggestion.
+      *
+      * @param suggestion the suggestion to import
+      */
+    case class ImportSuggestionRequest(suggestion: Suggestion)
+        extends ApiRequest
+
+    /** The result of the import request.
+      *
+      * @param module the definition module of the symbol
+      * @param symbol the resolved symbol
+      * @param exports the list of exports of the symbol
+      */
+    case class ImportSuggestionResponse(
+      module: String,
+      symbol: String,
+      exports: Seq[Export]
+    ) extends ApiResponse
+
     private lazy val mapper = {
       val factory = new CBORFactory()
       val mapper  = new ObjectMapper(factory) with ScalaObjectMapper

engine/runtime/src/main/scala/org/enso/interpreter/instrument/command/CommandFactory.scala

@@ -47,6 +47,9 @@ object CommandFactory {
       case payload: Api.InvalidateModulesIndexRequest =>
         new InvalidateModulesIndexCmd(request.requestId, payload)
 
+      case payload: Api.ImportSuggestionRequest =>
+        new ImportSuggestionCmd(request.requestId, payload)
+
       case Api.ShutDownRuntimeServer() =>
         throw new IllegalArgumentException(
           "ShutDownRuntimeServer request is not convertible to command object"

engine/runtime/src/main/scala/org/enso/interpreter/instrument/command/ImportSuggestionCmd.scala

@@ -0,0 +1,126 @@
+package org.enso.interpreter.instrument.command
+
+import org.enso.compiler.data.BindingsMap
+import org.enso.compiler.pass.analyse.BindingAnalysis
+import org.enso.interpreter.instrument.execution.RuntimeContext
+import org.enso.interpreter.runtime.Module
+import org.enso.polyglot.Suggestion
+import org.enso.polyglot.runtime.Runtime.Api
+
+import scala.concurrent.{ExecutionContext, Future}
+
+/** A command that gathers info required for suggestion import.
+  *
+  * @param maybeRequestId an option with request id
+  * @param request a request for suggestion import
+  */
+final class ImportSuggestionCmd(
+  maybeRequestId: Option[Api.RequestId],
+  val request: Api.ImportSuggestionRequest
+) extends Command(maybeRequestId) {
+
+  import ImportSuggestionCmd._
+
+  /** Executes a request.
+    *
+    * @param ctx contains suppliers of services to perform a request
+    * @param ec execution context
+    */
+  override def execute(implicit
+    ctx: RuntimeContext,
+    ec: ExecutionContext
+  ): Future[Unit] = Future {
+    val suggestion = request.suggestion
+    reply(
+      Api.ImportSuggestionResponse(
+        suggestion.module,
+        suggestion.name,
+        findExports.sortBy(_.depth).map(_.export)
+      )
+    )
+  }
+
+  /** Find re-exports of the given symbol.
+    *
+    * @param ctx contains suppliers of services to perform a request
+    */
+  private def findExports(implicit ctx: RuntimeContext): Seq[ExportResult] = {
+    val suggestion = request.suggestion
+    val topScope =
+      ctx.executionService.getContext.getCompiler.context.getTopScope
+    val builder = Vector.newBuilder[ExportResult]
+
+    topScope.getModules
+      .stream()
+      .filter(isCompiled)
+      .forEach { module =>
+        module.getIr.getMetadata(BindingAnalysis).foreach { bindings =>
+          builder ++= getQualifiedExport(
+            module,
+            suggestion,
+            bindings
+          )
+          builder ++= getUnqualifiedExport(module, suggestion, bindings)
+        }
+      }
+
+    builder.result()
+  }
+
+  /** Extract the qualified export from the bindings map. */
+  private def getQualifiedExport(
+    module: Module,
+    suggestion: Suggestion,
+    bindings: BindingsMap
+  ): Option[ExportResult] = {
+    bindings.resolvedExports
+      .find(_.module.getName.toString == suggestion.module)
+      .filter(_.exportedAs.isDefined)
+      .map { exportedModule =>
+        val qualified = Api.Export.Qualified(
+          module.getName.toString,
+          exportedModule.exportedAs
+        )
+        ExportResult(qualified, getDepth(module))
+      }
+  }
+
+  /** Extract the unqualified export from the bindings map. */
+  private def getUnqualifiedExport(
+    module: Module,
+    suggestion: Suggestion,
+    bindings: BindingsMap
+  ): Option[ExportResult] = {
+    bindings.exportedSymbols.get(suggestion.name).flatMap { resolvedExports =>
+      resolvedExports
+        .find(_.module.getName.toString == suggestion.module)
+        .map { _ =>
+          val unqualified = Api.Export.Unqualified(module.getName.toString)
+          ExportResult(unqualified, getDepth(module))
+        }
+    }
+  }
+
+  private def getDepth(module: Module): Int =
+    module.getName.path.size
+
+  private def isCompiled(module: Module): Boolean =
+    module.getIr != null
+}
+
+object ImportSuggestionCmd {
+
+  /** Module that exports target symbol.
+    *
+    * @param name the module name
+    */
+  private case class ExportingModule(name: String)
+
+  /** An intermediate result of exports resolution.
+    *
+    * @param export the module export
+    * @param depth how nested is the exporting module
+    */
+  private case class ExportResult(export: Api.Export, depth: Int)
+
+}

engine/runtime/src/main/scala/org/enso/interpreter/instrument/job/EnsureCompiledJob.scala

@@ -39,12 +39,7 @@ class EnsureCompiledJob(protected val files: Iterable[File])
   override def run(implicit ctx: RuntimeContext): CompilationStatus = {
     ctx.locking.acquireWriteCompilationLock()
     try {
-      val modules = files.flatMap { file =>
-        ctx.executionService.getContext.getModuleForFile(file).toScala
-      }
-      ensureIndexedModules(modules)
-      ensureIndexedImports(modules)
-      ensureCompiledScope()
+      ensureCompiledFiles(files)
     } finally {
       ctx.locking.releaseWriteCompilationLock()
     }
@@ -53,54 +48,78 @@ class EnsureCompiledJob(protected val files: Iterable[File])
   /** Run the scheduled compilation and invalidation logic, and send the
     * suggestion updates.
     *
-    * @param modules the list of modules to compile.
+    * @param files the list of files to compile.
     * @param ctx the runtime context
     */
-  protected def ensureIndexedModules(
-    modules: Iterable[Module]
-  )(implicit ctx: RuntimeContext): Unit = {
-    modules
-      .foreach { module =>
-        compile(module)
-        val changeset = applyEdits(new File(module.getPath))
-        compile(module).foreach { module =>
+  protected def ensureCompiledFiles(
+    files: Iterable[File]
+  )(implicit ctx: RuntimeContext): CompilationStatus = {
+    val modules = files.flatMap { file =>
+      ctx.executionService.getContext.getModuleForFile(file).toScala
+    }
+    val moduleCompilationStatus = modules.flatMap { module =>
+      ensureCompiledModule(module) +: ensureCompiledImports(module)
+    }
+    val scopeCompilationStatus = ensureCompiledScope()
+    (moduleCompilationStatus ++ scopeCompilationStatus).maxOption
+      .getOrElse(CompilationStatus.Success)
+  }
+
+  /** Run the scheduled compilation and invalidation logic, and send the
+    * suggestion updates.
+    *
+    * @param module the module to compile.
+    * @param ctx the runtime context
+    */
+  private def ensureCompiledModule(
+    module: Module
+  )(implicit ctx: RuntimeContext): CompilationStatus = {
+    compile(module)
+    val changeset = applyEdits(new File(module.getPath))
+    compile(module)
+      .map {
+        case Some(module) =>
           runInvalidationCommands(
-            buildCacheInvalidationCommands(changeset, module.getLiteralSource)
+            buildCacheInvalidationCommands(
+              changeset,
+              module.getLiteralSource
+            )
           )
           analyzeModule(module, changeset)
-        }
+          runCompilationDiagnostics(module)
+        case None =>
+          CompilationStatus.Success
       }
+      .getOrElse(CompilationStatus.Failure)
   }
 
   /** Compile the imported modules and send the suggestion updates.
     *
-    * @param modules the list of modules to analyze.
+    * @param module the modules to analyze.
     * @param ctx the runtime context
     */
-  protected def ensureIndexedImports(
-    modules: Iterable[Module]
-  )(implicit ctx: RuntimeContext): Unit = {
-    modules.foreach { module =>
-      compile(module).foreach { module =>
-        val importedModules =
-          new ImportResolver(ctx.executionService.getContext.getCompiler)
-            .mapImports(module)
-            .filter(_.getName != module.getName)
-        ctx.executionService.getLogger.finest(
-          s"Module ${module.getName} imports ${importedModules.map(_.getName)}"
-        )
-        importedModules.foreach(analyzeImport)
-      }
-    }
+
+  private def ensureCompiledImports(module: Module)(implicit
+    ctx: RuntimeContext
+  ): Seq[CompilationStatus] = {
+    val importedModules =
+      new ImportResolver(ctx.executionService.getContext.getCompiler)
+        .mapImports(module)
+        .filter(_.getName != module.getName)
+    ctx.executionService.getLogger.finest(
+      s"Module ${module.getName} imports ${importedModules.map(_.getName)}"
+    )
+    importedModules.foreach(analyzeImport)
+    importedModules.map(runCompilationDiagnostics)
   }
 
   /** Compile all modules in the scope and send the extracted suggestions.
     *
     * @param ctx the runtime context
     */
-  protected def ensureCompiledScope()(implicit
+  private def ensureCompiledScope()(implicit
     ctx: RuntimeContext
-  ): CompilationStatus = {
+  ): Iterable[CompilationStatus] = {
     val modulesInScope =
       ctx.executionService.getContext.getTopScope.getModules.asScala
     ctx.executionService.getLogger
@@ -118,13 +137,13 @@ class EnsureCompiledJob(protected val files: Iterable[File])
               )
             )
             CompilationStatus.Failure
-          case Right(module) =>
+          case Right(Some(module)) =>
             analyzeModuleInScope(module)
             runCompilationDiagnostics(module)
+          case Right(None) =>
+            CompilationStatus.Success
         }
       }
-      .maxOption
-      .getOrElse(CompilationStatus.Success)
   }
 
   private def analyzeImport(
@@ -296,17 +315,19 @@ class EnsureCompiledJob(protected val files: Iterable[File])
     */
   private def compile(
     module: Module
-  )(implicit ctx: RuntimeContext): Either[Throwable, Module] = {
+  )(implicit ctx: RuntimeContext): Either[Throwable, Option[Module]] = {
     val prevStage = module.getCompilationStage
     val compilationResult = Either.catchNonFatal {
       module.compileScope(ctx.executionService.getContext).getModule
     }
-    if (prevStage != module.getCompilationStage) {
-      ctx.executionService.getLogger.finest(
-        s"Compiled ${module.getName} $prevStage->${module.getCompilationStage}"
-      )
+    compilationResult.map { compiledModule =>
+      if (prevStage != compiledModule.getCompilationStage) {
+        ctx.executionService.getLogger.finest(
+          s"Compiled ${module.getName} $prevStage->${module.getCompilationStage}"
+        )
+        Some(compiledModule)
+      } else None
     }
-    compilationResult
   }
 
   /** Apply pending edits to the file.

engine/runtime/src/main/scala/org/enso/interpreter/instrument/job/EnsureCompiledStackJob.scala

@@ -5,7 +5,7 @@ import java.io.File
 import org.enso.compiler.pass.analyse.CachePreferenceAnalysis
 import org.enso.interpreter.instrument.{CacheInvalidation, InstrumentFrame}
 import org.enso.interpreter.instrument.execution.RuntimeContext
-import org.enso.interpreter.runtime.Module
+import org.enso.interpreter.instrument.job.EnsureCompiledJob.CompilationStatus
 import org.enso.polyglot.runtime.Runtime.Api
 
 import scala.jdk.OptionConverters._
@@ -19,10 +19,10 @@ class EnsureCompiledStackJob(stack: Iterable[InstrumentFrame])(implicit
 ) extends EnsureCompiledJob(EnsureCompiledStackJob.extractFiles(stack)) {
 
   /** @inheritdoc */
-  override protected def ensureIndexedModules(
-    modules: Iterable[Module]
-  )(implicit ctx: RuntimeContext): Unit = {
-    super.ensureIndexedModules(modules)
+  override protected def ensureCompiledFiles(
+    files: Iterable[File]
+  )(implicit ctx: RuntimeContext): CompilationStatus = {
+    val compilationStatus = super.ensureCompiledFiles(files)
     getCacheMetadata(stack).foreach { metadata =>
       CacheInvalidation.run(
         stack,
@@ -32,6 +32,7 @@ class EnsureCompiledStackJob(stack: Iterable[InstrumentFrame])(implicit
         )
       )
     }
+    compilationStatus
   }
 
   private def getCacheMetadata(