Improve VSCode/BSP integration (#2514)

#2415 only tested the Bloop codepath, and left the BSP code path broken. This fixes it Tested manually by setting `switch build server` in VSCode to BSP and then clicking around in the following examples: 1. `example/basic/1-simple-scala` 2. `example/misc/3-import-file-ivy` 3. `example/misc/4-mill-build-folder` The following navigations work: 1. between `foo` and `foo.test` 2. between `Foo` and the Java and Scala std libs 3. between `Foo` and Scalatags 4. between `build.sc` and Mill library code 5. between `build.sc` and Scalatags 6. between `build.sc` and `mill-build/src/` 7. between `mill-build/build.sc` and Mill library code 8. between `mill-build/build.sc` and Java and Scala std libs I re-used some code from `GenIdeaImpl` to gather the source jars of mill bundled modules in `MillBuildServer`. TBH the code in both `MillBuildServer` and `GenIdeaImpl` is a mess and could use some housekeeping, and `MillBuildServer` could also use some unit/e2e tests to verify basic the basic shape of that `workspaceBuildTargets`, `buildTargetSources`, `buildTargetDependencySources` and others return in a few simple example cases. But that can probably come in a follow up
com-lihaoyi · May 13, 2023 · 19ec00b · 19ec00b
1 parent 009a4da
commit 19ec00b
Show file tree

Hide file tree

Showing 8 changed files with 137 additions and 106 deletions.
diff --git a/bsp/worker/src/mill/bsp/worker/MillBuildServer.scala b/bsp/worker/src/mill/bsp/worker/MillBuildServer.scala
@@ -61,11 +61,11 @@ import mill.define.Segment.Label
 import mill.define.{Args, Discover, ExternalModule, Module, Segments, Task}
 import mill.eval.Evaluator
 import mill.main.{BspServerResult, MainModule}
-import mill.scalalib.{JavaModule, SemanticDbJavaModule, TestModule}
+import mill.scalalib.{GenIdeaImpl, JavaModule, SemanticDbJavaModule, TestModule}
 import mill.scalalib.bsp.{BspModule, JvmBuildTarget, ScalaBuildTarget}
 import mill.scalalib.internal.ModuleUtils
-import mill.runner.MillBuildRootModule
-import os.Path
+import mill.runner.{MillBuildBootstrap, MillBuildRootModule}
+import os.{Path, root}
 
 import java.io.PrintStream
 import java.util.concurrent.CompletableFuture
@@ -102,33 +102,34 @@ class MillBuildServer(
       def init(): Unit = synchronized {
         idToModule match {
           case None =>
-            val modules: Seq[Module] =
-              ModuleUtils.transitiveModules(evaluator.rootModule)
-            val map = modules.collect {
-              case m: MillBuildRootModule =>
-                val uri = sanitizeUri(m.millSourcePath)
-                val id = new BuildTargetIdentifier(uri)
-                for (millBuild <- millBuildRootModules) {
-                  log.debug(s"mill-build segments: ${millBuild.millModuleSegments.render}")
+            val modules: Seq[(Module, Seq[Module])] = rootModules
+              .map(rootModule => (rootModule, ModuleUtils.transitiveModules(rootModule)))
+
+            val map = modules
+              .flatMap { case (rootModule, otherModules) =>
+                (Seq(rootModule) ++ otherModules).collect {
+                  case m: BspModule =>
+                    val uri = sanitizeUri(
+                      rootModule.millSourcePath / m.millModuleSegments.parts
+                    )
+
+                    val id = new BuildTargetIdentifier(uri)
+
+                    (id, m)
                 }
-                (id, m)
-              case m: BspModule =>
-                val uri =
-                  sanitizeUri(millBuildRootModules.head.millSourcePath / m.millModuleSegments.parts)
-                val id = new BuildTargetIdentifier(uri)
-                (id, m)
-            }.toMap
+              }
+              .toMap
             idToModule = Some(map)
-            modulesToId = Some(map.map(_.swap).toMap)
-            log.debug(s"BspModules: ${map}")
+            modulesToId = Some(map.map(_.swap))
+            log.debug(s"BspModules: ${map.mapValues(_.bspDisplayName)}")
           case _ => // already init
         }
       }
       private[MillBuildServer] var idToModule: Option[Map[BuildTargetIdentifier, BspModule]] = None
       private[MillBuildServer] var modulesToId: Option[Map[BspModule, BuildTargetIdentifier]] = None
     }
 
-    lazy val millBuildRootModules: Seq[mill.runner.MillBuildRootModule] = {
+    lazy val rootModules: Seq[mill.main.RootModule] = {
       val evaluated = new mill.runner.MillBuildBootstrap(
         projectRoot = evaluator.rootModule.millSourcePath,
         home = os.home,
@@ -141,23 +142,26 @@ class MillBuildServer(
         logger = evaluator.baseLogger
       ).evaluate()
 
-      val allModules =
-        evaluated.result.bootstrapModuleOpt.toSeq ++
-          evaluated.result.frames
-            .zipWithIndex
-            // The depth=0 frame is for running user code, not for the build
-            .drop(1)
-            .flatMap {
-              case (f, i) =>
-                f.classLoaderOpt
-                  .flatMap(
-                    mill.runner.MillBuildBootstrap
-                      .getRootModule(_, i, evaluator.rootModule.millSourcePath)
-                      .toOption
-                  )
-            }
+      val rootModules0 = evaluated.result.frames
+        .flatMap(_.classLoaderOpt)
+        .zipWithIndex
+        .map { case (c, i) =>
+          MillBuildBootstrap
+            .getRootModule(c, i, evaluator.rootModule.millSourcePath)
+            .fold(sys.error(_), identity(_))
+        }
+
+      val bootstrapModule = evaluated.result.bootstrapModuleOpt.map(m =>
+        MillBuildBootstrap
+          .getChildRootModule(
+            m,
+            evaluated.result.frames.length,
+            evaluator.rootModule.millSourcePath
+          )
+          .fold(sys.error(_), identity(_))
+      )
 
-      allModules.collect { case m: mill.runner.MillBuildRootModule => m }
+      rootModules0 ++ bootstrapModule
     }
     def bspModulesById: Map[BuildTargetIdentifier, BspModule] = {
       internal.init()
@@ -167,11 +171,6 @@ class MillBuildServer(
       internal.init()
       internal.modulesToId.get
     }
-
-    /** Convert to BSP API. */
-    implicit class BspModuleSupport(val m: BspModule) {
-      def buildTargetId: BuildTargetIdentifier = bspIdByModule(m)
-    }
   }
 
   private[this] var statePromise: Promise[State] = Promise[State]()
@@ -394,14 +393,13 @@ class MillBuildServer(
         targetIds = sourcesParams.getTargets.asScala.toSeq,
         agg = (items: Seq[SourcesItem]) => new SourcesResult(items.asJava)
       ) {
-        case (id, module: MillBuildRootModule) if clientIsIntelliJ =>
+        case (id, module: MillBuildRootModule) =>
           T.task {
-            val sources = new SourcesItem(
-              id,
-              module.dummySources().map(p => sourceItem(p.path, true)).asJava
-            )
-            sources.setRoots(Seq(sanitizeUri(evaluator.rootModule.millSourcePath)).asJava)
-            sources
+            val items =
+              module.scriptSources().map(p => sourceItem(p.path, false)) ++
+                module.sources().map(p => sourceItem(p.path, false)) ++
+                module.generatedSources().map(p => sourceItem(p.path, true))
+            new SourcesItem(id, items.asJava)
           }
         case (id, module: JavaModule) =>
           T.task {
@@ -446,13 +444,20 @@ class MillBuildServer(
         agg = (items: Seq[DependencySourcesItem]) => new DependencySourcesResult(items.asJava)
       ) {
         case (id, m: JavaModule) =>
+          val buildSources =
+            if (!m.isInstanceOf[MillBuildRootModule]) Nil
+            else mill.scalalib.Lib
+              .resolveMillBuildDeps(Nil, None, useSources = true)
+              .map(sanitizeUri(_))
+
           T.task {
             val sources = m.resolveDeps(
               T.task(m.transitiveCompileIvyDeps() ++ m.transitiveIvyDeps()),
               sources = true
             )()
+
             val unmanaged = m.unmanagedClasspath()
-            val cp = (sources ++ unmanaged).map(sanitizeUri.apply).iterator.toSeq
+            val cp = (sources ++ unmanaged).map(sanitizeUri.apply).toSeq ++ buildSources
             new DependencySourcesItem(id, cp.asJava)
           }
       }
@@ -588,7 +593,7 @@ class MillBuildServer(
       import state._
 
       val modules = bspModulesById.values.toSeq.collect { case m: JavaModule => m }
-      val millBuildTargetIds = millBuildRootModules.map(bspIdByModule(_)).toSet
+      val millBuildTargetIds = rootModules.map { case m: BspModule => bspIdByModule(m) }.toSet
 
       val params = TaskParameters.fromTestParams(testParams)
       val argsMap =
@@ -674,7 +679,7 @@ class MillBuildServer(
     completable(s"buildTargetCleanCache ${cleanCacheParams}") { state =>
       import state._
 
-      val targetIds = millBuildRootModules.map(_.buildTargetId)
+      val targetIds = rootModules.map { case b: BspModule => bspIdByModule(b) }
       val (msg, cleaned) =
         cleanCacheParams.getTargets.asScala.filter(targetIds.contains).foldLeft((
           "",

diff --git a/integration/feature/watch-source-input/test/src/WatchSourceInputTests.scala b/integration/feature/watch-source-input/test/src/WatchSourceInputTests.scala
@@ -104,8 +104,8 @@ object WatchSourceInputTests extends IntegrationTestSuite {
 
     test("sources") {
 
-      test("noshow") - testWatchSource(false)
-      test("show") - testWatchSource(true)
+      test("noshow") - retry(3) { testWatchSource(false) }
+      test("show") - retry(3) { testWatchSource(true) }
     }
 
     def testWatchInput(show: Boolean) = {
@@ -146,8 +146,8 @@ object WatchSourceInputTests extends IntegrationTestSuite {
 
     test("input") {
 
-      test("noshow") - testWatchInput(false)
-      test("show") - testWatchInput(true)
+      test("noshow") - retry(3) { testWatchInput(false) }
+      test("show") - retry(3) { testWatchInput(true) }
     }
   }
 }
diff --git a/main/util/src/mill/util/CoursierSupport.scala b/main/util/src/mill/util/CoursierSupport.scala
@@ -92,7 +92,6 @@ trait CoursierSupport {
       ] = None,
       resolveFilter: os.Path => Boolean = _ => true
   ): Result[Agg[PathRef]] = {
-
     def isLocalTestDep(dep: coursier.Dependency): Option[Seq[PathRef]] = {
       val org = dep.module.organization.value
       val name = dep.module.name.value

diff --git a/runner/src/mill/runner/MillBuildBootstrap.scala b/runner/src/mill/runner/MillBuildBootstrap.scala
@@ -369,7 +369,9 @@ object MillBuildBootstrap {
     rootModuleOrErr.filterOrElse(
       rootModule =>
         depth == 0 || rootModule.isInstanceOf[mill.runner.MillBuildRootModule],
-      s"Root module in ${recRoot(projectRoot, depth).relativeTo(projectRoot)}/build.sc must be of ${classOf[MillBuildRootModule]}"
+      s"Root module in ${recRoot(projectRoot, depth).relativeTo(projectRoot)}/build.sc must be of ${classOf[
+          MillBuildRootModule
+        ]}, not ${rootModuleOrErr.map(_.getClass)}"
     )
   }
 

diff --git a/runner/src/mill/runner/MillBuildRootModule.scala b/runner/src/mill/runner/MillBuildRootModule.scala
@@ -29,6 +29,12 @@ class MillBuildRootModule()(implicit
     baseModuleInfo: RootModule.Info,
     millBuildRootModuleInfo: MillBuildRootModule.Info
 ) extends RootModule() with ScalaModule {
+  override def bspDisplayName0: String = millBuildRootModuleInfo
+    .projectRoot
+    .relativeTo(millBuildRootModuleInfo.topLevelProjectRoot)
+    .segments
+    .++(super.bspDisplayName0.split("/"))
+    .mkString("/")
 
   override def millSourcePath = millBuildRootModuleInfo.projectRoot / os.up / "mill-build"
 
@@ -37,11 +43,13 @@ class MillBuildRootModule()(implicit
       sources: Boolean = false
   ): Task[Agg[PathRef]] =
     T.task {
-      // We need to resolve the sources to make GenIdeaExtendedTests pass for
-      // some reason, but we don't need to actually return them (???)
-      val unused = super.resolveDeps(deps, true)()
-
-      super.resolveDeps(deps, false)()
+      if (sources == true) super.resolveDeps(deps, true)()
+      else {
+        // We need to resolve the sources to make GenIdeaExtendedTests pass for
+        // some reason, but we don't need to actually return them (???)
+        val unused = super.resolveDeps(deps, true)()
+        super.resolveDeps(deps, false)()
+      }
     }
 
   override def scalaVersion = "2.13.10"

diff --git a/scalalib/src/mill/scalalib/GenIdeaImpl.scala b/scalalib/src/mill/scalalib/GenIdeaImpl.scala
@@ -86,49 +86,16 @@ case class GenIdeaImpl(
 
     val buildLibraryPaths: immutable.Seq[Path] =
       if (!fetchMillModules) Nil
-      else
-        Util.millProperty("MILL_BUILD_LIBRARIES") match {
-          case Some(found) => found.split(',').map(os.Path(_)).distinct.toList
-          case None =>
-            val moduleRepos = evaluator.evalOrThrow(
-              exceptionFactory = r =>
-                GenIdeaException(
-                  s"Failure during resolving repositories: ${Evaluator.formatFailing(r)}"
-                )
-            )(modules.map(_._2.repositoriesTask))
-
-            val repos = moduleRepos.foldLeft(Set.empty[Repository])(_ ++ _) ++ Set(
-              LocalRepositories.ivy2Local,
-              Repositories.central
-            )
-            val millDeps = BuildInfo.millEmbeddedDeps.split(",").map(d => ivy"$d").map(dep =>
-              BoundDep(Lib.depToDependency(dep, BuildInfo.scalaVersion, ""), dep.force)
-            )
-            val Result.Success(res) = scalalib.Lib.resolveDependencies(
-              repositories = repos.toList,
-              deps = millDeps,
-              sources = false,
-              mapDependencies = None,
-              customizer = None,
-              coursierCacheCustomizer = None,
-              ctx = ctx
+      else {
+        val moduleRepos = evaluator.evalOrThrow(
+          exceptionFactory = r =>
+            GenIdeaException(
+              s"Failure during resolving repositories: ${Evaluator.formatFailing(r)}"
             )
+        )(modules.map(_._2.repositoriesTask))
 
-            // Also trigger resolve sources, but don't use them (will happen implicitly by Idea)
-            {
-              scalalib.Lib.resolveDependencies(
-                repositories = repos.toList,
-                deps = millDeps,
-                sources = true,
-                mapDependencies = None,
-                customizer = None,
-                coursierCacheCustomizer = None,
-                ctx = ctx
-              )
-            }
-
-            res.items.toList.map(_.path)
-        }
+        Lib.resolveMillBuildDeps(moduleRepos.flatten, ctx, useSources = false)
+      }
 
     val buildDepsPaths = Classpath
       .allJars(evaluator.rootModule.getClass.getClassLoader)

diff --git a/scalalib/src/mill/scalalib/Lib.scala b/scalalib/src/mill/scalalib/Lib.scala
@@ -2,8 +2,9 @@ package mill
 package scalalib
 
 import coursier.util.Task
-import coursier.{Dependency, Repository, Resolution}
+import coursier.{Dependency, LocalRepositories, Repositories, Repository, Resolution}
 import mill.api.{Ctx, Loose, PathRef, Result}
+import mill.modules.Util
 import mill.scalalib.api.ZincWorkerUtil
 
 object Lib {
@@ -135,4 +136,46 @@ object Lib {
     } yield path
   }
 
+  def resolveMillBuildDeps(
+      repos0: Seq[Repository],
+      ctx: Option[mill.api.Ctx.Log],
+      useSources: Boolean
+  ): Seq[os.Path] = {
+    Util.millProperty("MILL_BUILD_LIBRARIES") match {
+      case Some(found) => found.split(',').map(os.Path(_)).distinct.toList
+      case None =>
+        val repos = repos0 ++ Set(
+          LocalRepositories.ivy2Local,
+          Repositories.central
+        )
+        val millDeps = BuildInfo.millEmbeddedDeps.split(",").map(d => ivy"$d").map(dep =>
+          BoundDep(Lib.depToDependency(dep, BuildInfo.scalaVersion, ""), dep.force)
+        )
+        val Result.Success(res) = scalalib.Lib.resolveDependencies(
+          repositories = repos.toList,
+          deps = millDeps,
+          sources = useSources,
+          mapDependencies = None,
+          customizer = None,
+          coursierCacheCustomizer = None,
+          ctx = ctx
+        )
+
+        // Also trigger resolve sources, but don't use them (will happen implicitly by Idea)
+        if (!useSources) {
+          scalalib.Lib.resolveDependencies(
+            repositories = repos.toList,
+            deps = millDeps,
+            sources = true,
+            mapDependencies = None,
+            customizer = None,
+            coursierCacheCustomizer = None,
+            ctx = ctx
+          )
+        }
+
+        res.items.toList.map(_.path)
+    }
+  }
+
 }
diff --git a/scalalib/src/mill/scalalib/bsp/BspModule.scala b/scalalib/src/mill/scalalib/bsp/BspModule.scala
@@ -8,10 +8,17 @@ import mill._
 trait BspModule extends Module {
   import BspModule._
 
+  def bspDisplayName0: String = ModuleUtils.moduleDisplayName(this)
+
+  def bspDisplayName = bspDisplayName0 match {
+    case "" => "root-module"
+    case n => n
+  }
+
   /** Use to fill most fields of `BuildTarget`. */
   @internal
   def bspBuildTarget: BspBuildTarget = BspBuildTarget(
-    displayName = Some(ModuleUtils.moduleDisplayName(this)),
+    displayName = Some(bspDisplayName),
     baseDirectory = Some(millSourcePath),
     tags = Seq(Tag.Library, Tag.Application),
     languageIds = Seq(),