Allow users to disable some modules in BSP (#6065)

This adds a `def enableBsp: Boolean` method on `JavaModule` (via
`BspModuleApi`). Users can override this method and make it return
`false` to make the BSP server hide / ignore that module.

Note that for consistency, if a module depends on another that has
`enableBsp` set to false, it gets itself disabled too. Else, that might
create situations where a module is loaded in BSP, but some of its
dependencies would just be missing because they have `enableBsp` set to
false.

Author: alexarchambault

core/api/daemon/src/mill/api/daemon/internal/bsp/BspModuleApi.scala

@@ -6,6 +6,14 @@ trait BspModuleApi extends ModuleApi {
   private[mill] def bspBuildTargetData: TaskApi[Option[(String, AnyRef)]]
   private[mill] def bspBuildTarget: BspBuildTarget
   private[mill] def bspDisplayName: String
+
+  /**
+   * Set this to false to make the Mill BSP server hide / ignore that module
+   *
+   * Beware that if a module depends via `moduleDeps` or `compileModuleDeps` on modules
+   * that have `enableBsp` set to false, it will be ignored by the Mill BSP server too
+   */
+  def enableBsp: Boolean = true
 }
 
 object BspModuleApi {

integration/ide/bsp-server/resources/project/build.mill

@@ -114,3 +114,14 @@ object sourcesNeedCompile extends scalalib.ScalaModule {
     super.sources()
   }
 }
+
+object skipped extends scalalib.ScalaModule {
+  def enableBsp = false
+  def scalaVersion = Option(System.getenv("TEST_SCALA_2_13_VERSION")).getOrElse(???)
+}
+
+// depends on skipped module, but doesn't set enableBsp to false - should be ignored from BSP anyway
+object dependsOnSkipped extends scalalib.ScalaModule {
+  def moduleDeps = Seq(skipped)
+  def scalaVersion = Option(System.getenv("TEST_SCALA_2_13_VERSION")).getOrElse(???)
+}

integration/ide/bsp-server/resources/snapshots/logging

@@ -7,10 +7,14 @@
 [bsp-init-mill-build/build.mill-59] [info] done compiling
 [bsp-init-build.mill-59] [info] compiling * Scala sources to * ...
 [bsp-init-build.mill-59] [info] done compiling
-[bsp-init-098] [info] compiling * Scala source to * ...
-[bsp-init-098] [info] done compiling
+[bsp-init-106] [info] compiling * Scala source to * ...
+[bsp-init-106] [info] done compiling
+[bsp-init] BSP disabled for target file:///workspace/dependsOnSkipped because of its dependencies file:///workspace/skipped
+[bsp-init] BSP disabled for target file:///workspace/skipped via BspModuleApi#enableBsp
 [bsp-init] SNAPSHOT
 [2-workspaceBuildTargets] Entered workspaceBuildTargets
+[bsp-init] BSP disabled for target file:///workspace/dependsOnSkipped because of its dependencies file:///workspace/skipped
+[bsp-init] BSP disabled for target file:///workspace/skipped via BspModuleApi#enableBsp
 [2-workspaceBuildTargets] Evaluating * tasks
 [2-workspaceBuildTargets] Done
 [2-workspaceBuildTargets] Evaluating 1 task

integration/ide/bsp-server/src/BspServerTests.scala

@@ -420,9 +420,11 @@ object BspServerTests extends UtestIntegrationTestSuite {
         delayedCompileFuture.get()
       }
 
+      val workspaceUri = tester.workspacePath.toURI.toASCIIString.stripSuffix("/") + "/"
       val logs = stderr.toString
         .linesWithSeparators
         .filter(_.startsWith("["))
+        .map(_.replace(workspaceUri, "file:///workspace/"))
         .mkString
 
       val expectedCancelledLine = "[7-compile] buildTargetCompile was cancelled"

libs/javalib/src/mill/javalib/JavaModule.scala

@@ -80,6 +80,8 @@ trait JavaModule
       outer.repositoriesTask()
     }
 
+    override def enableBsp: Boolean = outer.enableBsp
+
     override def resolutionCustomizer: Task[Option[coursier.Resolution => coursier.Resolution]] =
       outer.resolutionCustomizer
 

runner/bsp/worker/src/mill/bsp/worker/BspEvaluators.scala

@@ -6,6 +6,7 @@ import mill.api.daemon.internal.{BaseModuleApi, EvaluatorApi, JavaModuleApi, Mod
 import mill.api.daemon.Watchable
 
 import java.nio.file.Path
+import java.util.concurrent.ConcurrentHashMap
 
 private[mill] class BspEvaluators(
     workspaceDir: os.Path,
@@ -15,30 +16,72 @@ private[mill] class BspEvaluators(
 ) {
 
   /**
-   * Compute all transitive modules from module children and via moduleDeps + compileModuleDeps
+   * Compute all transitive modules from module children via moduleDirectChildren
    */
   def transitiveModules(module: ModuleApi): Seq[ModuleApi] = {
     Seq(module) ++ module.moduleDirectChildren.flatMap(transitiveModules)
   }
 
-  lazy val bspModulesIdList0: Seq[(BuildTargetIdentifier, (BspModuleApi, EvaluatorApi))] = {
-    val modules: Seq[(ModuleApi, Seq[ModuleApi], EvaluatorApi)] = evaluators
-      .map(ev => (ev.rootModule, transitiveModules(ev.rootModule), ev))
+  private val transitiveDependencyModules0 = new ConcurrentHashMap[ModuleApi, Seq[ModuleApi]]
+  private val transitiveModulesEnableBsp0 =
+    new ConcurrentHashMap[ModuleApi, Option[Seq[BspModuleApi]]]
 
-    val regularModules = modules
-      .flatMap { case (rootModule, modules, eval) =>
-        modules.collect {
-          case m: BspModuleApi =>
-            val uri = Utils.sanitizeUri(
-              (os.Path(rootModule.moduleDirJava) / m.moduleSegments.parts).toNIO
-            )
+  // Compute all transitive dependency modules via moduleDeps + compileModuleDeps
+  private def transitiveDependencyModules(module: ModuleApi): Seq[ModuleApi] = {
+    if (!transitiveDependencyModules0.contains(module)) {
+      val directDependencies = module match {
+        case jm: JavaModuleApi => jm.recursiveModuleDeps ++ jm.compileModuleDepsChecked
+        case _ => Nil
+      }
+      val value = Seq(module) ++ directDependencies.flatMap(transitiveDependencyModules)
+      transitiveDependencyModules0.putIfAbsent(module, value)
+    }
 
-            (new BuildTargetIdentifier(uri), (m, eval))
-        }
+    transitiveDependencyModules0.get(module)
+  }
+
+  private def transitiveModulesEnableBsp(module: ModuleApi): Option[Seq[BspModuleApi]] = {
+    if (!transitiveModulesEnableBsp0.contains(module)) {
+      val disabledTransitiveModules = transitiveDependencyModules(module).collect {
+        case b: BspModuleApi if !b.enableBsp => b
       }
-    regularModules
+      val value =
+        if (disabledTransitiveModules.isEmpty) None
+        else Some(disabledTransitiveModules)
+      transitiveModulesEnableBsp0.putIfAbsent(module, value)
+    }
+
+    transitiveModulesEnableBsp0.get(module)
   }
 
+  private def moduleUri(rootModule: ModuleApi, module: ModuleApi) = Utils.sanitizeUri(
+    (os.Path(rootModule.moduleDirJava) / module.moduleSegments.parts).toNIO
+  )
+
+  lazy val bspModulesIdList0: Seq[(BuildTargetIdentifier, (BspModuleApi, EvaluatorApi))] =
+    for {
+      eval <- evaluators
+      bspModule <- transitiveModules(eval.rootModule).collect { case m: BspModuleApi => m }
+      uri = moduleUri(eval.rootModule, bspModule)
+      if {
+        if (bspModule.enableBsp)
+          transitiveModulesEnableBsp(bspModule) match {
+            case Some(disabledTransitiveModules) =>
+              val uris = disabledTransitiveModules.map(moduleUri(eval.rootModule, _))
+              eval.baseLogger.warn(
+                s"BSP disabled for target $uri because of its dependencies ${uris.mkString(", ")}"
+              )
+              false
+            case None =>
+              true
+          }
+        else {
+          eval.baseLogger.info(s"BSP disabled for target $uri via BspModuleApi#enableBsp")
+          false
+        }
+      }
+    } yield (new BuildTargetIdentifier(uri), (bspModule, eval))
+
   val nonScriptSources = evaluators.flatMap { ev =>
     val bspSourceTasks: Seq[TaskApi[(sources: Seq[Path], generatedSources: Seq[Path])]] =
       transitiveModules(ev.rootModule)