Use original Evaluators to preserve their association with modules in BSP to avoid using the wrong evaluator

bsp/src/mill/bsp/BSP.scala

@@ -49,9 +49,10 @@ object BSP extends ExternalModule with CoursierModule {
    * @param ev The Evaluator
    * @return The server result, indicating if mill should re-run this command or just exit.
    */
-  def startSession(ev: Evaluator): Command[BspServerResult] = T.command {
+  def startSession(allBootstrapEvaluators: Evaluator.AllBootstrapEvaluators)
+      : Command[BspServerResult] = T.command {
     T.log.errorStream.println("BSP/startSession: Starting BSP session")
-    val res = BspContext.bspServerHandle.runSession(ev)
+    val res = BspContext.bspServerHandle.runSession(allBootstrapEvaluators.value)
     T.log.errorStream.println(s"BSP/startSession: Finished BSP session, result: ${res}")
     res
   }

bsp/src/mill/bsp/BspServerHandle.scala

@@ -11,7 +11,7 @@ trait BspServerHandle {
    * Runs a new session with the given evaluator. This one blocks until the session ends.
    * @return The reason which the session ended, possibly indictating the wish for restart (e.g. in case of workspace reload).
    */
-  def runSession(evaluator: Evaluator): BspServerResult
+  def runSession(evaluators: Seq[Evaluator]): BspServerResult
 
   /**
    * The result of the latest started session. Once a new session was started but not finished, this may be [[None]].

bsp/worker/src/mill/bsp/worker/BspWorkerImpl.scala

@@ -56,9 +56,9 @@ private class BspWorkerImpl() extends BspWorker {
       val bspServerHandle = new BspServerHandle {
         private[this] var lastResult0: Option[BspServerResult] = None
 
-        override def runSession(evaluator: Evaluator): BspServerResult = {
+        override def runSession(evaluators: Seq[Evaluator]): BspServerResult = {
           lastResult0 = None
-          millServer.updateEvaluator(Option(evaluator))
+          millServer.updateEvaluator(Option(evaluators))
           val onReload = Promise[BspServerResult]()
           millServer.onSessionEnd = Some { serverResult =>
             if (!onReload.isCompleted) {

bsp/worker/src/mill/bsp/worker/MillBuildServer.scala

@@ -58,7 +58,7 @@ import com.google.gson.JsonObject
 import mill.T
 import mill.api.{DummyTestReporter, Result, Strict}
 import mill.define.Segment.Label
-import mill.define.{Args, Discover, ExternalModule, Task}
+import mill.define.{Args, BaseModule, Discover, ExternalModule, Task}
 import mill.eval.Evaluator
 import mill.main.MainModule
 import mill.scalalib.{JavaModule, SemanticDbJavaModule, TestModule}
@@ -94,23 +94,15 @@ private class MillBuildServer(
   protected var clientIsIntelliJ = false
 
   private[this] var statePromise: Promise[State] = Promise[State]()
-  var evaluatorOpt: Option[Evaluator] = None
-  def evaluator = evaluatorOpt.get
 
-  def updateEvaluator(evaluator: Option[Evaluator]): Unit = {
-    debug(s"Updating Evaluator: $evaluator")
+  def updateEvaluator(evaluatorsOpt: Option[Seq[Evaluator]]): Unit = {
+    debug(s"Updating Evaluator: $evaluatorsOpt")
     if (statePromise.isCompleted) statePromise = Promise[State]() // replace the promise
-    evaluatorOpt = evaluator
-    evaluatorOpt.foreach(e =>
+    evaluatorsOpt.foreach { evaluators =>
       statePromise.success(
-        new State(
-          e.rootModule.millSourcePath,
-          e.baseLogger,
-          debug,
-          e.disableCallgraphInvalidation
-        )
+        new State(evaluators, debug)
       )
-    )
+    }
   }
 
   def debug(msg: String) = logStream.println(msg)
@@ -207,7 +199,7 @@ private class MillBuildServer(
       "workspaceBuildTargets",
       targetIds = _.bspModulesById.keySet.toSeq,
       tasks = { case m: JavaModule => T.task { m.bspBuildTargetData() } }
-    ) { (state, id, m, bspBuildTargetData) =>
+    ) { (ev, state, id, m, bspBuildTargetData) =>
       val s = m.bspBuildTarget
       val deps = m match {
         case jm: JavaModule =>
@@ -296,7 +288,7 @@ private class MillBuildServer(
           }
       }
     ) {
-      case (state, id, module, items) => new SourcesItem(id, items.asJava)
+      case (ev, state, id, module, items) => new SourcesItem(id, items.asJava)
     } {
       new SourcesResult(_)
     }
@@ -306,18 +298,23 @@ private class MillBuildServer(
   override def buildTargetInverseSources(p: InverseSourcesParams)
       : CompletableFuture[InverseSourcesResult] = {
     completable(s"buildtargetInverseSources ${p}") { state =>
-      val tasks = state.bspModulesById.iterator.collect {
-        case (id, m: JavaModule) =>
+      val tasksEvaluators = state.bspModulesById.iterator.collect {
+        case (id, (m: JavaModule, ev)) =>
           T.task {
             val src = m.allSourceFiles()
             val found = src.map(sanitizeUri).contains(
               p.getTextDocument.getUri
             )
-            if (found) Seq((id)) else Seq()
-          }
+            if (found) Seq(id) else Seq()
+          } -> ev
       }.toSeq
 
-      val ids = evaluator.evalOrThrow()(tasks).flatten
+      val ids = tasksEvaluators
+        .groupMap(_._2)(_._1)
+        .flatMap { case (ev, ts) => ev.evalOrThrow()(ts) }
+        .flatten
+        .toSeq
+
       new InverseSourcesResult(ids.asJava)
     }
   }
@@ -342,7 +339,7 @@ private class MillBuildServer(
         }
       }
     ) {
-      case (state, id, m: JavaModule, (resolveDepsSources, unmanagedClasspath)) =>
+      case (ev, state, id, m: JavaModule, (resolveDepsSources, unmanagedClasspath)) =>
         val buildSources =
           if (!m.isInstanceOf[MillBuildRootModule]) Nil
           else mill.scalalib.Lib
@@ -368,6 +365,7 @@ private class MillBuildServer(
       }
     ) {
       case (
+            ev,
             state,
             id,
             m: JavaModule,
@@ -394,7 +392,7 @@ private class MillBuildServer(
         case _ => T.task { Nil }
       }
     ) {
-      case (state, id, m, resources) =>
+      case (ev, state, id, m, resources) =>
         val resourcesUrls = resources.map(_.path).filter(os.exists).map(sanitizeUri)
         new ResourcesItem(id, resourcesUrls.asJava)
 
@@ -408,22 +406,28 @@ private class MillBuildServer(
     completable(s"buildTargetCompile ${p}") { state =>
       val params = TaskParameters.fromCompileParams(p)
       val taskId = params.hashCode()
-      val compileTasks = params.getTargets.distinct.map(state.bspModulesById).map {
-        case m: SemanticDbJavaModule if clientWantsSemanticDb => m.compiledClassesAndSemanticDbFiles
-        case m: JavaModule => m.compile
-        case m => T.task {
+      val compileTasksEvs = params.getTargets.distinct.map(state.bspModulesById).map {
+        case (m: SemanticDbJavaModule, ev) if clientWantsSemanticDb =>
+          (m.compiledClassesAndSemanticDbFiles, ev)
+        case (m: JavaModule, ev) => (m.compile, ev)
+        case (m, ev) => T.task {
             Result.Failure(
               s"Don't know how to compile non-Java target ${m.bspBuildTarget.displayName}"
             )
-          }
+          } -> ev
       }
 
-      val result = evaluator.evaluate(
-        compileTasks,
-        Utils.getBspLoggedReporterPool(p.getOriginId, state.bspIdByModule, client),
-        DummyTestReporter,
-        new MillBspLogger(client, taskId, evaluator.baseLogger)
-      )
+      val result = compileTasksEvs
+        .groupMap(_._2)(_._1)
+        .map { case (ev, ts) =>
+          ev.evaluate(
+            ts,
+            Utils.getBspLoggedReporterPool(p.getOriginId, state.bspIdByModule, client),
+            DummyTestReporter,
+            new MillBspLogger(client, taskId, ev.baseLogger)
+          )
+        }
+        .toSeq
       val compileResult = new CompileResult(Utils.getStatusCode(result))
       compileResult.setOriginId(p.getOriginId)
       compileResult // TODO: See in what form IntelliJ expects data about products of compilation in order to set data field
@@ -432,25 +436,25 @@ private class MillBuildServer(
   override def buildTargetOutputPaths(params: OutputPathsParams)
       : CompletableFuture[OutputPathsResult] =
     completable(s"buildTargetOutputPaths ${params}") { state =>
-      val outItems = new OutputPathItem(
-        // Spec says, a directory must end with a forward slash
-        sanitizeUri(evaluator.outPath) + "/",
-        OutputPathItemKind.DIRECTORY
-      )
-
-      val extItems = new OutputPathItem(
-        // Spec says, a directory must end with a forward slash
-        sanitizeUri(evaluator.externalOutPath) + "/",
-        OutputPathItemKind.DIRECTORY
-      )
-
       val items = for {
         target <- params.getTargets.asScala
-        module <- state.bspModulesById.get(target)
+        (module, ev) <- state.bspModulesById.get(target)
       } yield {
         val items =
-          if (module.millOuterCtx.external) List(extItems)
-          else List(outItems)
+          if (module.millOuterCtx.external) List(
+            new OutputPathItem(
+              // Spec says, a directory must end with a forward slash
+              sanitizeUri(ev.externalOutPath) + "/",
+              OutputPathItemKind.DIRECTORY
+            )
+          )
+          else List(
+            new OutputPathItem(
+              // Spec says, a directory must end with a forward slash
+              sanitizeUri(ev.outPath) + "/",
+              OutputPathItemKind.DIRECTORY
+            )
+          )
         new OutputPathsItem(target, items.asJava)
       }
 
@@ -460,15 +464,16 @@ private class MillBuildServer(
   override def buildTargetRun(runParams: RunParams): CompletableFuture[RunResult] =
     completable(s"buildTargetRun ${runParams}") { state =>
       val params = TaskParameters.fromRunParams(runParams)
-      val module = params.getTargets.map(state.bspModulesById).collectFirst {
-        case m: JavaModule => m
+      val (module, ev) = params.getTargets.map(state.bspModulesById).collectFirst {
+        case (m: JavaModule, ev) => (m, ev)
       }.get
+
       val args = params.getArguments.getOrElse(Seq.empty[String])
       val runTask = module.run(T.task(Args(args)))
-      val runResult = evaluator.evaluate(
+      val runResult = ev.evaluate(
         Strict.Agg(runTask),
         Utils.getBspLoggedReporterPool(runParams.getOriginId, state.bspIdByModule, client),
-        logger = new MillBspLogger(client, runTask.hashCode(), evaluator.baseLogger)
+        logger = new MillBspLogger(client, runTask.hashCode(), ev.baseLogger)
       )
       val response = runResult.results(runTask) match {
         case _: Result.Success[Any] => new RunResult(StatusCode.OK)
@@ -510,7 +515,7 @@ private class MillBuildServer(
         .filter(millBuildTargetIds.contains)
         .foldLeft(StatusCode.OK) { (overallStatusCode, targetId) =>
           state.bspModulesById(targetId) match {
-            case testModule: TestModule =>
+            case (testModule: TestModule, ev) =>
               val testTask = testModule.testLocal(argsMap(targetId.getUri): _*)
 
               // notifying the client that the testing of this build target started
@@ -529,17 +534,17 @@ private class MillBuildServer(
                   Seq.empty[String]
                 )
 
-              val results = evaluator.evaluate(
+              val results = ev.evaluate(
                 Strict.Agg(testTask),
                 Utils.getBspLoggedReporterPool(
                   testParams.getOriginId,
                   state.bspIdByModule,
                   client
                 ),
                 testReporter,
-                new MillBspLogger(client, testTask.hashCode, evaluator.baseLogger)
+                new MillBspLogger(client, testTask.hashCode, ev.baseLogger)
               )
-              val statusCode = Utils.getStatusCode(results)
+              val statusCode = Utils.getStatusCode(Seq(results))
 
               // Notifying the client that the testing of this build target ended
               val taskFinishParams =
@@ -582,16 +587,16 @@ private class MillBuildServer(
           true
         )) {
           case ((msg, cleaned), targetId) =>
-            val module = state.bspModulesById(targetId)
+            val (module, ev) = state.bspModulesById(targetId)
             val mainModule = new MainModule {
               override implicit def millDiscover: Discover[_] = Discover[this.type]
             }
             val compileTargetName = (module.millModuleSegments ++ Label("compile")).render
             debug(s"about to clean: ${compileTargetName}")
-            val cleanTask = mainModule.clean(evaluator, Seq(compileTargetName): _*)
-            val cleanResult = evaluator.evaluate(
+            val cleanTask = mainModule.clean(ev, Seq(compileTargetName): _*)
+            val cleanResult = ev.evaluate(
               Strict.Agg(cleanTask),
-              logger = new MillBspLogger(client, cleanTask.hashCode, evaluator.baseLogger)
+              logger = new MillBspLogger(client, cleanTask.hashCode, ev.baseLogger)
             )
             if (cleanResult.failing.keyCount > 0) (
               msg + s" Target ${compileTargetName} could not be cleaned. See message from mill: \n" +
@@ -602,7 +607,7 @@ private class MillBuildServer(
               false
             )
             else {
-              val outPaths = evaluator.pathsResolver.resolveDest(
+              val outPaths = ev.pathsResolver.resolveDest(
                 module.millModuleSegments ++ Label("compile")
               )
               val outPathSeq = Seq(outPaths.dest, outPaths.meta, outPaths.log)
@@ -626,14 +631,23 @@ private class MillBuildServer(
       hint: String,
       targetIds: State => Seq[BuildTargetIdentifier],
       tasks: BspModule => Task[W]
-  )(f: (State, BuildTargetIdentifier, BspModule, W) => T)(agg: java.util.List[T] => V)
+  )(f: (Evaluator, State, BuildTargetIdentifier, BspModule, W) => T)(agg: java.util.List[T] => V)
       : CompletableFuture[V] =
     completable(hint) { state: State =>
       val ids = targetIds(state)
-      val tasksSeq = ids.map(m => tasks(state.bspModulesById(m)))
-      val evaluated = evaluator.evalOrThrow()(tasksSeq)
-      val res = evaluated.zip(ids).map { case (v, i) => f(state, i, state.bspModulesById(i), v) }
-      agg(res.asJava)
+      val tasksSeq = ids.map { id =>
+        val (m, ev) = state.bspModulesById(id)
+        (tasks(m), (ev, id))
+      }
+
+      val evaluated = tasksSeq
+        .groupMap(_._2)(_._1)
+        .map { case ((ev, id), ts) =>
+          ev.evalOrThrow()(ts)
+            .map { v => f(ev, state, id, state.bspModulesById(id)._1, v) }
+        }
+
+      agg(evaluated.flatten.toSeq.asJava)
     }
 
   /**

bsp/worker/src/mill/bsp/worker/MillJavaBuildServer.scala

@@ -29,8 +29,8 @@ private trait MillJavaBuildServer extends JavaBuildServer { this: MillBuildServe
         T.task { (classesPathTask(), m.javacOptions(), m.bspCompileClasspath()) }
       }
     ) {
-      case (state, id, m: JavaModule, (classesPath, javacOptions, bspCompileClasspath)) =>
-        val pathResolver = evaluator.pathsResolver
+      case (ev, state, id, m: JavaModule, (classesPath, javacOptions, bspCompileClasspath)) =>
+        val pathResolver = ev.pathsResolver
         val options = javacOptions
         val classpath =
           bspCompileClasspath.map(_.resolve(pathResolver)).map(sanitizeUri)

bsp/worker/src/mill/bsp/worker/MillJvmBuildServer.scala

@@ -60,6 +60,7 @@ private trait MillJvmBuildServer extends JvmBuildServer { this: MillBuildServer
       }
     ) {
       case (
+            ev,
             state,
             id,
             m: JavaModule,