POC: Make BSP modular and support dynamic extensions

bsp/src/mill/bsp/BSP.scala

@@ -2,18 +2,44 @@ package mill.bsp
 
 import mill.api.{Ctx, PathRef}
 import mill.{Agg, T}
-import mill.define.{Command, Discover, ExternalModule}
+import mill.define.{Command, Discover, ExternalModule, Task}
 import mill.main.BuildInfo
 import mill.eval.Evaluator
 import mill.util.Util.millProjectModule
-import mill.scalalib.CoursierModule
+import mill.scalalib.{BoundDep, CoursierModule, Dep, JavaModule, Lib}
+import mill.scalalib.internal.JavaModuleUtils
 
 object BSP extends ExternalModule with CoursierModule {
 
+  private def evaluator(): Evaluator =
+    Option(Evaluator.currentEvaluator.value)
+      .orElse(Option(Evaluator.allBootstrapEvaluators.value).flatMap(_.value.headOption))
+      .get
+
   lazy val millDiscover: Discover[this.type] = Discover[this.type]
 
+  private def bspExtensions: T[(Seq[String], Seq[Dep])] = T.input {
+    // As this is a runtime value which can change, we need to be in an input target
+    val modules = JavaModuleUtils.transitiveModules(evaluator().rootModule)
+      .collect { case m: JavaModule => m }
+    val extensions = modules.flatMap(m => m.bspExtensions).distinct
+    val classes = extensions.map(_.className).distinct
+    T.log.debug(s"BSP extensions: ${BspUtil.pretty(extensions)}")
+    val extensionIvyDeps = extensions.flatMap(_.ivyDeps)
+
+    (classes, extensionIvyDeps)
+  }
+
+  override def bindDependency: Task[Dep => BoundDep] = T.task { dep: Dep =>
+    Lib.depToBoundDep(dep, BuildInfo.scalaVersion, s"_mill${BuildInfo.millBinPlatform}")
+  }
+
   private def bspWorkerLibs: T[Agg[PathRef]] = T {
-    millProjectModule("mill-bsp-worker", repositoriesTask())
+    millProjectModule(
+      "mill-bsp-worker",
+      repositoriesTask(),
+      extraDeps = bspExtensions()._2.map(bindDependency().andThen(_.dep))
+    )
   }
 
   /**
@@ -32,12 +58,20 @@ object BSP extends ExternalModule with CoursierModule {
    */
   def install(jobs: Int = 1): Command[(PathRef, ujson.Value)] = T.command {
     // we create a file containing the additional jars to load
-    val libUrls = bspWorkerLibs().map(_.path.toNIO.toUri.toURL).iterator.toSeq
+//    val libUrls = bspWorkerLibs().map(_.path.toNIO.toUri).iterator.toSeq
+
+    val classes = bspExtensions()._1
+
+    val bspServerConfig = BspServerConfig(
+      classes,
+      bspWorkerLibs().toSeq
+    )
+
     val cpFile =
-      T.workspace / Constants.bspDir / s"${Constants.serverName}-${BuildInfo.millVersion}.resources"
+      T.workspace / Constants.bspDir / s"${Constants.serverName}-${BuildInfo.millVersion}.conf"
     os.write.over(
       cpFile,
-      libUrls.mkString("\n"),
+      upickle.default.write(bspServerConfig, indent = 2),
       createFolders = true
     )
     createBspConnection(jobs, Constants.serverName)
@@ -97,7 +131,8 @@ object BSP extends ExternalModule with CoursierModule {
         millVersion = BuildInfo.millVersion,
         bspVersion = Constants.bspProtocolVersion,
         languages = Constants.languages
-      )
+      ),
+      indent = 2
     )
   }
 

bsp/src/mill/bsp/BspContext.scala

@@ -68,7 +68,8 @@ private[mill] class BspContext(
         streams,
         logStream.getOrElse(streams.err),
         home / Constants.bspDir,
-        canReload
+        canReload,
+        os.pwd
       )
     }
   }

bsp/src/mill/bsp/BspServerConfig.scala

@@ -0,0 +1,13 @@
+package mill.bsp
+
+import mill.api.PathRef
+
+case class BspServerConfig(
+    extensions: Seq[String],
+    classpath: Seq[PathRef]
+)
+
+object BspServerConfig {
+  implicit val jsonify: upickle.default.ReadWriter[BspServerConfig] =
+    upickle.default.macroRW
+}

bsp/src/mill/bsp/BspUtil.scala

@@ -0,0 +1,7 @@
+package mill.bsp
+
+private[bsp] trait BspUtil {
+  def pretty = pprint.PPrinter(defaultHeight = 10000)
+}
+
+object BspUtil extends BspUtil

bsp/src/mill/bsp/BspWorker.scala

@@ -4,58 +4,69 @@ import mill.api.{Ctx, Logger, SystemStreams}
 import os.Path
 
 import java.io.PrintStream
-import java.net.URL
 
 private trait BspWorker {
   def startBspServer(
       streams: SystemStreams,
       logStream: PrintStream,
       logDir: os.Path,
-      canReload: Boolean
+      canReload: Boolean,
+      projectDir: os.Path
   ): Either[String, BspServerHandle]
 }
 
 private object BspWorker {
 
   private[this] var worker: Option[BspWorker] = None
 
+  def readConfig(workspace: os.Path): Either[String, BspServerConfig] = {
+    val configFile =
+      workspace / Constants.bspDir / s"${Constants.serverName}-${mill.main.BuildInfo.millVersion}.conf"
+    if (!os.exists(configFile)) return Left(
+      s"""Could not find config file: ${configFile}
+         |You need to run `mill mill.bsp.BSP/install` before you can use the BSP server""".stripMargin
+    )
+
+    val config = upickle.default.read[BspServerConfig](
+      os.read(configFile)
+    )
+
+    // TODO: if outdated, we could regenerate the resource file and re-load the worker
+    Right(config)
+  }
+
   def apply(
       workspace: os.Path,
       home0: os.Path,
-      log: Logger,
-      workerLibs: Option[Seq[URL]] = None
+      log: Logger
+//      workerLibs: Option[Seq[URL]] = None
   ): Either[String, BspWorker] = {
     worker match {
       case Some(x) => Right(x)
       case None =>
-        val urls = workerLibs.map { urls =>
-          log.debug("Using direct submitted worker libs")
-          urls
-        }.getOrElse {
-          // load extra classpath entries from file
-          val cpFile =
-            workspace / Constants.bspDir / s"${Constants.serverName}-${mill.main.BuildInfo.millVersion}.resources"
-          if (!os.exists(cpFile)) return Left(
-            "You need to run `mill mill.bsp.BSP/install` before you can use the BSP server"
-          )
+//        val urls = workerLibs.map { urls =>
+//          log.debug("Using direct submitted worker libs")
+//          urls
+//        }.getOrElse {
+        // load extra config from file
+        readConfig(workspace).map { config =>
+          log.debug(s"BSP Server config: ${BspUtil.pretty(config)}")
 
-          // TODO: if outdated, we could regenerate the resource file and re-load the worker
+          val urls = config.classpath.map(_.path.toNIO.toUri.toURL)
 
-          // read the classpath from resource file
-          log.debug(s"Reading worker classpath from file: ${cpFile}")
-          os.read(cpFile).linesIterator.map(u => new URL(u)).toSeq
-        }
-
-        // create classloader with bsp.worker and deps
-        val cl = mill.api.ClassLoader.create(urls, getClass().getClassLoader())(
-          new Ctx.Home { override def home: Path = home0 }
-        )
+          // create classloader with bsp.worker and deps
+          val cl = mill.api.ClassLoader.create(urls, getClass().getClassLoader())(
+            new Ctx.Home {
+              override def home: Path = home0
+            }
+          )
 
-        val workerCls = cl.loadClass(Constants.bspWorkerImplClass)
-        val ctr = workerCls.getConstructor()
-        val workerImpl = ctr.newInstance().asInstanceOf[BspWorker]
-        worker = Some(workerImpl)
-        Right(workerImpl)
+          val workerCls = cl.loadClass(Constants.bspWorkerImplClass)
+          val ctr = workerCls.getConstructor()
+          val workerImpl = ctr.newInstance().asInstanceOf[BspWorker]
+          worker = Some(workerImpl)
+          workerImpl
+        }
     }
   }
 

bsp/src/mill/bsp/spi/MillBuildServerBase.scala

@@ -0,0 +1,42 @@
+package mill.bsp.spi
+
+import mill.define.Task
+import mill.eval.Evaluator
+import mill.scalalib.bsp.{BspModule, BspUri}
+
+import java.util.concurrent.CompletableFuture
+import scala.reflect.ClassTag
+
+trait MillBuildServerBase {
+
+  /**
+   * Write a debug message to the log file.
+   * @param msg
+   */
+  def debug(msg: String): Unit
+
+  /**
+   * Given a function that take input of type T and return output of type V,
+   * apply the function on the given inputs and return a completable future of
+   * the result. If the execution of the function raises an Exception, complete
+   * the future exceptionally. Also complete exceptionally if the server was not
+   * yet initialized.
+   */
+  protected def completable[V](
+      hint: String,
+      checkInitialized: Boolean = true
+  )(f: State => V): CompletableFuture[V]
+
+  def completableTasks[T, V, W: ClassTag](
+      hint: String,
+      targetIds: State => Seq[BspUri],
+      tasks: BspModule => Task[W]
+  )(f: (Evaluator, State, BspUri, BspModule, W) => T)(agg: java.util.List[T] => V)
+      : CompletableFuture[V]
+
+  /**
+   * `true` if semanticDb generation is enabled for this build.
+   * @return
+   */
+  def enableSemanticDb: Boolean
+}

bsp/worker/src/mill/bsp/worker/State.scala → bsp/src/mill/bsp/spi/State.scala

@@ -1,25 +1,26 @@
-package mill.bsp.worker
+package mill.bsp.spi
 
-import ch.epfl.scala.bsp4j.BuildTargetIdentifier
-import mill.scalalib.bsp.BspModule
-import mill.scalalib.internal.JavaModuleUtils
 import mill.define.Module
 import mill.eval.Evaluator
+import mill.scalalib.bsp.{BspModule, BspUri}
+import mill.scalalib.internal.JavaModuleUtils
+
+class State private[bsp] (evaluators: Seq[Evaluator], debug: String => Unit) {
 
-private class State(evaluators: Seq[Evaluator], debug: String => Unit) {
-  lazy val bspModulesById: Map[BuildTargetIdentifier, (BspModule, Evaluator)] = {
+  /** Mapping of BSP target identifier to the Mill module and evaluator. */
+  lazy val bspModulesById: Map[BspUri, (BspModule, Evaluator)] = {
     val modules: Seq[(Module, Seq[Module], Evaluator)] = evaluators
       .map(ev => (ev.rootModule, JavaModuleUtils.transitiveModules(ev.rootModule), ev))
 
     val map = modules
       .flatMap { case (rootModule, otherModules, eval) =>
         (Seq(rootModule) ++ otherModules).collect {
           case m: BspModule =>
-            val uri = Utils.sanitizeUri(
+            val uri = BspUri(
               rootModule.millSourcePath / m.millModuleSegments.parts
             )
 
-            (new BuildTargetIdentifier(uri), (m, eval))
+            (uri, (m, eval))
         }
       }
       .toMap
@@ -28,8 +29,10 @@ private class State(evaluators: Seq[Evaluator], debug: String => Unit) {
     map
   }
 
+  /** All root modules (at different meta-levels) of the project. */
   lazy val rootModules: Seq[mill.define.BaseModule] = evaluators.map(_.rootModule)
 
-  lazy val bspIdByModule: Map[BspModule, BuildTargetIdentifier] =
+  /** Mapping of Mill Modules to BSP target identifiers. */
+  lazy val bspIdByModule: Map[BspModule, BspUri] =
     bspModulesById.view.mapValues(_._1).map(_.swap).toMap
 }