Use original Evaluators to preserve their association with modules …

…in BSP to avoid using the wrong evaluator (#2692)

Should fix #2643.

The basic problem is previously BSP always used the `frame(0)` evaluator
passed in to `updateEvaluator`, which is meant specifically for
evaluating user tasks in the root `./build.sc`. However, BSP also
querying of modules and tasks both in `build.sc` as well as meta-builds
in `mill-build/build.sc` and nested meta-builds. This ends up evaluating
the meta-builds in the same `out/` folder as the main build, rather than
in the `out/mill-build/` folders they should be evaluated in.

This PR avoids using the `Evaluator` given in `updateEvaluator`, and
instead preserves the `Evaluator`s we already construct as part of
`MillBuildBootstrap` in `mill.bsp.worker.State`, and passes them as a
`Seq[Evaluator]` from the enclosing `MillBuildBootstrap` code all the
way to the Mill-BSP code so we can associate each `Module` with its
respective `Ealuator`. We then use those associated `Evaluator`s to
perform the actual evaluation throughout the
`Mill{Scala,Java,Jvm,}BuildServer` implementations. There's some
additional re-shaping, so that in case a single BSP API call wants to
query modules/tasks from multiple evaluators at once we properly group
the tasks/modules together according to their respective evaluators for
evaluation.

By passing the `Seq[Evaluator]` from the enclosing `MillBuildBootstrap`,
this also lets us remove the now-redundant call to `MillBuildBootstrap`
buried within `State.scala`. We already have all the necessary data
structures from the enclosing `MillBuildBootstrap` and its evaluators,
so we don't need to re-run the bootstrap process and re-construct all of
it from scratch just for BSP

I haven't managed to reproduce the original error, so confirming the fix
will have to wait until this lands in master and people can try it out.
But I have manually tested using VSCode to load and jump around
`example/basic/1-simple-scala/` and
`example/scalabuilds/10-scala-realistic/`, and it seems all the
references and modules work and are set up correctly. So at least we
have some confidence that this doesn't cause any regressions.

This should also supersede
#2648, since we now no longer
need to re-run the `MillBuildBootstrap`/`RootModuleFinder` logic at all
since we just take the existing list of `Evaluator`s from the enclosing
bootstrap process

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,