[feature] Jetpack Compose preview support | #BAZEL-868

[fix] Explicitly build the Android Kotlin target during build

[feature] buildTarget/jvmBinaryJars request for Jetpack Compose preview

[maintenance] AndroidBuildTarget should use String instead of URI to conform with BSP

[maintenance] Move Android resource calculation to the Android plugin


Merge-request: BAZEL-MR-947
Merged-by: Lev Leontev <lev.leontev@jetbrains.com>

protocol/src/main/kotlin/org/jetbrains/bsp/AndroidBuildTarget.kt

@@ -3,7 +3,6 @@ package org.jetbrains.bsp
 import ch.epfl.scala.bsp4j.JvmBuildTarget
 import com.google.gson.annotations.JsonAdapter
 import org.eclipse.lsp4j.jsonrpc.json.adapters.EnumTypeAdapter
-import java.net.URI
 
 @JsonAdapter(EnumTypeAdapter.Factory::class)
 public enum class AndroidTargetType(public val value: Int) {
@@ -13,10 +12,10 @@ public enum class AndroidTargetType(public val value: Int) {
 }
 
 public data class AndroidBuildTarget(
-  val androidJar: URI,
+  val androidJar: String,
   val androidTargetType: AndroidTargetType,
-  val manifest: URI?,
-  val resourceFolders: List<URI>,
+  val manifest: String?,
+  val resourceFolders: List<String>,
   var jvmBuildTarget: JvmBuildTarget? = null,
   var kotlinBuildTarget: KotlinBuildTarget? = null,
 )

protocol/src/main/kotlin/org/jetbrains/bsp/BazelBuildServer.kt

@@ -19,4 +19,7 @@ public interface BazelBuildServer {
 
   @JsonRequest("buildTarget/mobileInstall")
   public fun buildTargetMobileInstall(params: MobileInstallParams): CompletableFuture<MobileInstallResult>
+
+  @JsonRequest("buildTarget/jvmBinaryJars")
+  public fun buildTargetJvmBinaryJars(params: JvmBinaryJarsParams): CompletableFuture<JvmBinaryJarsResult>
 }

protocol/src/main/kotlin/org/jetbrains/bsp/BazelBuildServerCapabilities.kt

@@ -28,6 +28,7 @@ public class BazelBuildServerCapabilities(
   public val workspaceDirectoriesProvider: Boolean = false,
   public val workspaceInvalidTargetsProvider: Boolean = false,
   public val runWithDebugProvider: Boolean = false,
+  public val jvmBinaryJarsProvider: Boolean = false,
 ) : BuildServerCapabilities() {
   init {
     this.compileProvider = compileProvider

protocol/src/main/kotlin/org/jetbrains/bsp/JvmBinaryJars.kt

@@ -0,0 +1,16 @@
+package org.jetbrains.bsp
+
+import ch.epfl.scala.bsp4j.BuildTargetIdentifier
+
+public data class JvmBinaryJarsParams(
+  val targets: List<BuildTargetIdentifier>,
+)
+
+public data class JvmBinaryJarsResult(
+  val items: List<JvmBinaryJarsItem>,
+)
+
+public data class JvmBinaryJarsItem(
+  val target: BuildTargetIdentifier,
+  val jars: List<String>,
+)

server/src/main/kotlin/org/jetbrains/bsp/bazel/server/bsp/BspServerApi.kt

@@ -54,6 +54,8 @@ import ch.epfl.scala.bsp4j.TestResult
 import ch.epfl.scala.bsp4j.WorkspaceBuildTargetsResult
 import org.eclipse.lsp4j.jsonrpc.CancelChecker
 import org.jetbrains.bsp.BazelBuildServer
+import org.jetbrains.bsp.JvmBinaryJarsParams
+import org.jetbrains.bsp.JvmBinaryJarsResult
 import org.jetbrains.bsp.MobileInstallParams
 import org.jetbrains.bsp.MobileInstallResult
 import org.jetbrains.bsp.RunWithDebugParams
@@ -362,6 +364,16 @@ class BspServerApi(private val bazelServicesBuilder: (BuildClient) -> BazelServi
     )
   }
 
+  override fun buildTargetJvmBinaryJars(params: JvmBinaryJarsParams): CompletableFuture<JvmBinaryJarsResult> {
+    return runner.handleRequest(
+      "jvmBinaryJars", { cancelChecker: CancelChecker, params: JvmBinaryJarsParams ->
+        projectSyncService.jvmBinaryJars(
+          cancelChecker, params
+        )
+      }, params
+    )
+  }
+
   override fun workspaceLibraries(): CompletableFuture<WorkspaceLibrariesResult> {
     return runner.handleRequest("libraries") { cancelChecker: CancelChecker ->
       projectSyncService.workspaceBuildLibraries(

server/src/main/kotlin/org/jetbrains/bsp/bazel/server/sync/BazelProjectMapper.kt

@@ -417,7 +417,7 @@ class BazelProjectMapper(
     val baseDirectory = label.toDirectoryUri()
     val languagePlugin = languagePluginsService.getPlugin(languages)
     val sourceSet = resolveSourceSet(target, languagePlugin)
-    val resources = resolveResources(target)
+    val resources = resolveResources(target, languagePlugin)
     val languageData = languagePlugin.resolveModule(target)
     val sourceDependencies = languagePlugin.dependencySources(target, dependencyGraph)
     val environment = environmentItem(target)
@@ -479,17 +479,8 @@ class BazelProjectMapper(
     bspClientLogger.error(message)
   }
 
-  private fun resolveResources(target: TargetInfo): Set<URI> =
-    bazelPathsResolver.resolveUris(target.resourcesList).toSet() + resolveAndroidResources(target)
-
-  private fun resolveAndroidResources(target: TargetInfo): Set<URI> {
-    if (!target.hasAndroidTargetInfo()) return emptySet()
-    val androidTargetInfo = target.androidTargetInfo
-    if (!androidTargetInfo.hasManifest()) return emptySet()
-
-    return bazelPathsResolver
-      .resolveUris(listOf(target.androidTargetInfo.manifest) + target.androidTargetInfo.resourcesList).toSet()
-  }
+  private fun resolveResources(target: TargetInfo, languagePlugin: LanguagePlugin<*>): Set<URI> =
+    bazelPathsResolver.resolveUris(target.resourcesList).toSet() + languagePlugin.resolveAdditionalResources(target)
 
   private fun buildReverseSourceMapping(modules: List<Module>): Map<URI, Label> =
     modules.asSequence().flatMap(::buildReverseSourceMappingForModule).toMap()

server/src/main/kotlin/org/jetbrains/bsp/bazel/server/sync/BspProjectMapper.kt

@@ -53,6 +53,9 @@ import ch.epfl.scala.bsp4j.WorkspaceBuildTargetsResult
 import org.eclipse.lsp4j.jsonrpc.CancelChecker
 import org.jetbrains.bsp.BazelBuildServerCapabilities
 import org.jetbrains.bsp.DirectoryItem
+import org.jetbrains.bsp.JvmBinaryJarsItem
+import org.jetbrains.bsp.JvmBinaryJarsParams
+import org.jetbrains.bsp.JvmBinaryJarsResult
 import org.jetbrains.bsp.LibraryItem
 import org.jetbrains.bsp.WorkspaceDirectoriesResult
 import org.jetbrains.bsp.WorkspaceInvalidTargetsResult
@@ -103,7 +106,8 @@ class BspProjectMapper(
             workspaceLibrariesProvider = true,
             workspaceDirectoriesProvider = true,
             workspaceInvalidTargetsProvider = true,
-            runWithDebugProvider = true
+            runWithDebugProvider = true,
+            jvmBinaryJarsProvider = true,
         )
         return InitializeBuildResult(
             Constants.NAME, Constants.VERSION, Constants.BSP_VERSION, capabilities
@@ -325,6 +329,21 @@ class BspProjectMapper(
         }
     }
 
+    fun jvmBinaryJars(project: Project, params: JvmBinaryJarsParams): JvmBinaryJarsResult {
+        fun toJvmBinaryJarsItem(module: Module): JvmBinaryJarsItem? =
+            module.javaModule?.let { javaModule ->
+                val jars = javaModule.binaryOutputs.map { it.toString() }
+                JvmBinaryJarsItem(BspMappings.toBspId(module), jars)
+            }
+
+        val jvmBinaryJarsItems = params.targets.mapNotNull { target ->
+            val label = Label(target.uri)
+            val module = project.findModule(label)
+            module?.let { toJvmBinaryJarsItem(it) }
+        }
+        return JvmBinaryJarsResult(jvmBinaryJarsItems)
+    }
+
     fun buildTargetJavacOptions(project: Project, params: JavacOptionsParams, cancelChecker: CancelChecker): JavacOptionsResult {
         val items = params.targets.collectClasspathForTargetsAndApply(project, cancelChecker) { module, ideClasspath ->
             module.javaModule?.let { toJavacOptionsItem(module, it, ideClasspath) }

server/src/main/kotlin/org/jetbrains/bsp/bazel/server/sync/ExecuteService.kt

@@ -32,11 +32,11 @@ import org.jetbrains.bsp.bazel.server.bsp.managers.BepReader
 import org.jetbrains.bsp.bazel.server.diagnostics.DiagnosticsService
 import org.jetbrains.bsp.bazel.server.paths.BazelPathsResolver
 import org.jetbrains.bsp.bazel.server.sync.BspMappings.toBspId
+import org.jetbrains.bsp.bazel.server.sync.languages.android.AdditionalAndroidBuildTargetsQuery.additionalAndroidBuildTargetsQuery
 import org.jetbrains.bsp.bazel.server.sync.model.Module
 import org.jetbrains.bsp.bazel.server.sync.model.Tag
 import org.jetbrains.bsp.bazel.workspacecontext.TargetsSpec
 import org.jetbrains.bsp.bazel.workspacecontext.WorkspaceContextProvider
-import java.util.Optional
 
 class ExecuteService(
     private val compilationManager: BazelBspCompilationManager,
@@ -167,14 +167,20 @@ class ExecuteService(
     }
 
     private fun build(cancelChecker: CancelChecker, bspIds: List<BuildTargetIdentifier>, originId: String): BazelProcessResult {
-        val targetsSpec = TargetsSpec(bspIds, emptyList())
+        val targets = bspIds + getAdditionalBuildTargets(cancelChecker, bspIds)
+        val targetsSpec = TargetsSpec(targets, emptyList())
         return compilationManager.buildTargetsWithBep(
                 cancelChecker,
             targetsSpec,
             originId
         ).processResult
     }
 
+    private fun getAdditionalBuildTargets(
+        cancelChecker: CancelChecker,
+        bspIds: List<BuildTargetIdentifier>,
+    ): List<BuildTargetIdentifier> = additionalAndroidBuildTargetsQuery(bspIds, bazelRunner, cancelChecker)
+
     private fun selectTargets(cancelChecker: CancelChecker, targets: List<BuildTargetIdentifier>): List<BuildTargetIdentifier> =
             selectModules(cancelChecker, targets).map { toBspId(it) }
 

server/src/main/kotlin/org/jetbrains/bsp/bazel/server/sync/ProjectSyncService.kt

@@ -33,6 +33,8 @@ import ch.epfl.scala.bsp4j.SourcesParams
 import ch.epfl.scala.bsp4j.SourcesResult
 import ch.epfl.scala.bsp4j.WorkspaceBuildTargetsResult
 import org.eclipse.lsp4j.jsonrpc.CancelChecker
+import org.jetbrains.bsp.JvmBinaryJarsParams
+import org.jetbrains.bsp.JvmBinaryJarsResult
 import org.jetbrains.bsp.WorkspaceDirectoriesResult
 import org.jetbrains.bsp.WorkspaceInvalidTargetsResult
 import org.jetbrains.bsp.WorkspaceLibrariesResult
@@ -113,6 +115,11 @@ class ProjectSyncService(private val bspMapper: BspProjectMapper, private val pr
     return bspMapper.jvmTestEnvironment(project, params, cancelChecker)
   }
 
+  fun jvmBinaryJars(cancelChecker: CancelChecker, params: JvmBinaryJarsParams): JvmBinaryJarsResult {
+    val project = projectProvider.get(cancelChecker)
+    return bspMapper.jvmBinaryJars(project, params)
+  }
+
   fun buildTargetJavacOptions(cancelChecker: CancelChecker, params: JavacOptionsParams): JavacOptionsResult {
     val project = projectProvider.get(cancelChecker)
     return bspMapper.buildTargetJavacOptions(project, params, cancelChecker)

server/src/main/kotlin/org/jetbrains/bsp/bazel/server/sync/languages/LanguagePlugin.kt

@@ -9,6 +9,7 @@ import java.nio.file.Path
 abstract class LanguagePlugin<T : LanguageData> {
 
     open fun calculateSourceRoot(source: Path): Path? = null
+    open fun resolveAdditionalResources(targetInfo: BspTargetInfo.TargetInfo): Set<URI> = emptySet()
     open fun prepareSync(targets: Sequence<BspTargetInfo.TargetInfo>) {}
     open fun resolveModule(targetInfo: BspTargetInfo.TargetInfo): T? = null
 

server/src/main/kotlin/org/jetbrains/bsp/bazel/server/sync/languages/android/AdditionalAndroidBuildTargetsQuery.kt

@@ -0,0 +1,40 @@
+package org.jetbrains.bsp.bazel.server.sync.languages.android
+
+import ch.epfl.scala.bsp4j.BuildTargetIdentifier
+import ch.epfl.scala.bsp4j.StatusCode
+import org.eclipse.lsp4j.jsonrpc.CancelChecker
+import org.jetbrains.bsp.bazel.bazelrunner.BazelRunner
+import org.jetbrains.bsp.bazel.bazelrunner.utils.BazelArgumentsUtils.getJoinedBazelTargets
+import org.jetbrains.bsp.bazel.bazelrunner.utils.BazelArgumentsUtils.toRawUris
+
+/**
+ * Every Kotlin Android target actually produces three targets, which we merge inside [KotlinAndroidModulesMerger].
+ * However, in order for all the dependent libraries to be unpacked properly (e.g. for Jetpack Compose preview),
+ * we still have to pass the dependent Kotlin target explicitly during build (and not just the merged target).
+ */
+object AdditionalAndroidBuildTargetsQuery {
+  fun additionalAndroidBuildTargetsQuery(
+    targets: List<BuildTargetIdentifier>,
+    bazelRunner: BazelRunner,
+    cancelChecker: CancelChecker,
+  ): List<BuildTargetIdentifier> {
+    val targetUris = toRawUris(targets)
+    val joinedTargets = getJoinedBazelTargets(targetUris)
+    val childTargetsResult = bazelRunner.commandBuilder().query()
+      .withArgument("""kind("kt_jvm_library", deps(kind("android_", $joinedTargets), 1))""")
+      .executeBazelCommand(null).waitAndGetResult(cancelChecker)
+
+    if (childTargetsResult.statusCode != StatusCode.OK) {
+      error("Could not retrieve additional android build targets")
+    }
+
+    val targetUriSet = targetUris.asSequence().map { it.trimStart('@') }.toSet()
+
+    return childTargetsResult.stdoutLines
+      .asSequence()
+      .filter { it.endsWith("_kt") }
+      .filter { it.dropLast(3) in targetUriSet }
+      .map { BuildTargetIdentifier(it) }
+      .toList()
+  }
+}

server/src/main/kotlin/org/jetbrains/bsp/bazel/server/sync/languages/android/AndroidLanguagePlugin.kt

@@ -18,10 +18,10 @@ class AndroidLanguagePlugin(
   override fun applyModuleData(moduleData: AndroidModule, buildTarget: BuildTarget) {
     val androidBuildTarget = with(moduleData) {
       AndroidBuildTarget(
-        androidJar = androidJar,
+        androidJar = androidJar.toString(),
         androidTargetType = androidTargetType,
-        manifest = manifest,
-        resourceFolders = resourceFolders,
+        manifest = manifest?.toString(),
+        resourceFolders = resourceFolders.map { it.toString() },
       )
     }
     moduleData.javaModule?.let { javaLanguagePlugin.toJvmBuildTarget(it) }?.let {
@@ -68,4 +68,14 @@ class AndroidLanguagePlugin(
 
   override fun calculateSourceRoot(source: Path): Path =
     javaLanguagePlugin.calculateSourceRoot(source)
+
+  override fun resolveAdditionalResources(targetInfo: BspTargetInfo.TargetInfo): Set<URI> {
+    if (!targetInfo.hasAndroidTargetInfo()) return emptySet()
+    val androidTargetInfo = targetInfo.androidTargetInfo
+    if (!androidTargetInfo.hasManifest()) return emptySet()
+
+    return bazelPathsResolver
+      .resolveUris(listOf(targetInfo.androidTargetInfo.manifest) + targetInfo.androidTargetInfo.resourcesList)
+      .toSet()
+  }
 }

server/src/main/kotlin/org/jetbrains/bsp/bazel/server/sync/languages/android/KotlinAndroidModulesMerger.kt

@@ -1,5 +1,6 @@
 package org.jetbrains.bsp.bazel.server.sync.languages.android
 
+import org.jetbrains.bsp.bazel.server.sync.languages.jvm.javaModule
 import org.jetbrains.bsp.bazel.server.sync.languages.kotlin.KotlinModule
 import org.jetbrains.bsp.bazel.server.sync.model.Module
 
@@ -48,7 +49,10 @@ class KotlinAndroidModulesMerger {
     val kotlinLanguageData = kotlinModule.languageData
     if (kotlinLanguageData !is KotlinModule?) return null
     val androidLanguageData = androidModule.languageData as? AndroidModule ?: return null
-    val kotlinAndroidLanguageData = androidLanguageData.copy(kotlinModule = kotlinLanguageData)
+    val javaModule = androidLanguageData.javaModule?.run {
+      copy(binaryOutputs = binaryOutputs + kotlinModule.javaModule?.binaryOutputs.orEmpty())
+    }
+    val kotlinAndroidLanguageData = androidLanguageData.copy(kotlinModule = kotlinLanguageData, javaModule = javaModule)
 
     val kotlinModuleWithoutSdk = kotlinModule.directDependencies.asSequence()
       .filterNot { it.value.endsWith("//third_party:android_sdk") }

server/src/main/kotlin/org/jetbrains/bsp/bazel/server/sync/languages/java/JavaLanguagePlugin.kt

@@ -27,9 +27,7 @@ class JavaLanguagePlugin(
     override fun resolveModule(targetInfo: TargetInfo): JavaModule? =
         targetInfo.takeIf(TargetInfo::hasJvmTargetInfo)?.jvmTargetInfo?.run {
             val mainOutput = bazelPathsResolver.resolveUri(getJars(0).getBinaryJars(0))
-            val allOutputs = jarsList.flatMap {
-                it.interfaceJarsList + it.binaryJarsList
-            }.map(bazelPathsResolver::resolveUri)
+            val binaryOutputs = jarsList.flatMap { it.binaryJarsList }.map(bazelPathsResolver::resolveUri)
             val mainClass = getMainClass(this)
             val runtimeJdk = jdkResolver.resolveJdk(targetInfo)
 
@@ -39,7 +37,7 @@ class JavaLanguagePlugin(
                 javacOptsList,
                 jvmFlagsList,
                 mainOutput,
-                allOutputs,
+                binaryOutputs,
                 mainClass,
                 argsList,
             )

server/src/main/kotlin/org/jetbrains/bsp/bazel/server/sync/languages/java/JavaModule.kt

@@ -14,7 +14,7 @@ data class JavaModule(
     val javacOpts: List<String>,
     val jvmOps: List<String>,
     val mainOutput: URI,
-    val allOutputs: List<URI>,
+    val binaryOutputs: List<URI>,
     val mainClass: String?,
     val args: List<String>,
 ) : LanguageData