[internal] BSP: switch to aggregate build targets

src/python/pants/backend/java/bsp/rules.py

@@ -6,42 +6,32 @@
 from dataclasses import dataclass
 
 from pants.backend.java.bsp.spec import JavacOptionsItem, JavacOptionsParams, JavacOptionsResult
-from pants.backend.java.compile.javac import compute_output_jar_filename
-from pants.backend.java.dependency_inference.symbol_mapper import AllJavaTargets
 from pants.backend.java.target_types import JavaSourceField
 from pants.base.build_root import BuildRoot
-from pants.bsp.context import BSPContext
+from pants.base.specs import AddressSpecs
 from pants.bsp.protocol import BSPHandlerMapping
-from pants.bsp.spec.base import (
-    BuildTarget,
-    BuildTargetCapabilities,
-    BuildTargetIdentifier,
-    StatusCode,
-)
+from pants.bsp.spec.base import BuildTargetIdentifier, StatusCode
 from pants.bsp.util_rules.compile import BSPCompileFieldSet, BSPCompileResult
 from pants.bsp.util_rules.lifecycle import BSPLanguageSupport
-from pants.bsp.util_rules.targets import BSPBuildTargets, BSPBuildTargetsRequest
-from pants.build_graph.address import Address, AddressInput
+from pants.bsp.util_rules.targets import (
+    BSPBuildTargetsMetadataRequest,
+    BSPBuildTargetsMetadataResult,
+    BSPBuildTargetsNew,
+)
 from pants.engine.addresses import Addresses
 from pants.engine.fs import CreateDigest, DigestEntries
 from pants.engine.internals.native_engine import EMPTY_DIGEST, AddPrefix, Digest
 from pants.engine.internals.selectors import Get, MultiGet
 from pants.engine.rules import collect_rules, rule
-from pants.engine.target import (
-    CoarsenedTargets,
-    Dependencies,
-    DependenciesRequest,
-    Target,
-    WrappedTarget,
-)
-from pants.engine.unions import UnionMembership, UnionRule
-from pants.jvm.bsp.spec import JvmBuildTarget
+from pants.engine.target import CoarsenedTargets, FieldSet, Targets
+from pants.engine.unions import UnionRule
 from pants.jvm.compile import (
     ClasspathEntryRequest,
     ClasspathEntryRequestFactory,
     FallibleClasspathEntry,
 )
 from pants.jvm.resolve.key import CoursierResolveKey
+from pants.jvm.target_types import JvmResolveField
 
 LANGUAGE_ID = "java"
 
@@ -53,62 +43,27 @@ class JavaBSPLanguageSupport(BSPLanguageSupport):
     can_compile = True
 
 
-class JavaBSPBuildTargetsRequest(BSPBuildTargetsRequest):
-    pass
-
-
 @dataclass(frozen=True)
-class ResolveJavaBSPBuildTargetRequest:
-    target: Target
+class JavaMetadataFieldSet(FieldSet):
+    required_fields = (JavaSourceField, JvmResolveField)
 
+    source: JavaSourceField
+    resolve: JvmResolveField
 
-@rule
-async def bsp_resolve_one_java_build_target(
-    request: ResolveJavaBSPBuildTargetRequest,
-    union_membership: UnionMembership,
-) -> BuildTarget:
-    dep_addrs = await Get(Addresses, DependenciesRequest(request.target[Dependencies]))
-    impls = union_membership.get(BSPCompileFieldSet)
-
-    reported_deps = []
-    for dep_addr in dep_addrs:
-        if dep_addr == request.target.address:
-            continue
-
-        wrapped_dep_tgt = await Get(WrappedTarget, Address, dep_addr)
-        dep_tgt = wrapped_dep_tgt.target
-        for impl in impls:
-            if impl.is_applicable(dep_tgt):
-                reported_deps.append(BuildTargetIdentifier.from_address(dep_tgt.address))
-                break
-
-    return BuildTarget(
-        id=BuildTargetIdentifier.from_address(request.target.address),
-        display_name=str(request.target.address),
-        base_directory=None,
-        tags=(),
-        capabilities=BuildTargetCapabilities(
-            can_compile=True,
-        ),
-        language_ids=(LANGUAGE_ID,),
-        dependencies=tuple(reported_deps),
-        data_kind="jvm",
-        data=JvmBuildTarget(),
-    )
+
+class JavaBSPBuildTargetsMetadataRequest(BSPBuildTargetsMetadataRequest):
+    language_id = LANGUAGE_ID
+    can_merge_metadata_from = ()
+    field_set_type = JavaMetadataFieldSet
 
 
 @rule
-async def bsp_resolve_all_java_build_targets(
-    _: JavaBSPBuildTargetsRequest,
-    all_java_targets: AllJavaTargets,
-    bsp_context: BSPContext,
-) -> BSPBuildTargets:
-    if LANGUAGE_ID not in bsp_context.client_params.capabilities.language_ids:
-        return BSPBuildTargets()
-    build_targets = await MultiGet(
-        Get(BuildTarget, ResolveJavaBSPBuildTargetRequest(tgt)) for tgt in all_java_targets
+async def bsp_resolve_java_metadata(
+    _: JavaBSPBuildTargetsMetadataRequest,
+) -> BSPBuildTargetsMetadataResult:
+    return BSPBuildTargetsMetadataResult(
+        can_compile=True,
     )
-    return BSPBuildTargets(targets=tuple(build_targets))
 
 
 # -----------------------------------------------------------------------------------------------
@@ -137,23 +92,25 @@ class HandleJavacOptionsResult:
 async def handle_bsp_java_options_request(
     request: HandleJavacOptionsRequest,
     build_root: BuildRoot,
+    bsp_build_targets: BSPBuildTargetsNew,
 ) -> HandleJavacOptionsResult:
-    wrapped_target = await Get(WrappedTarget, AddressInput, request.bsp_target_id.address_input)
-    coarsened_targets = await Get(CoarsenedTargets, Addresses([wrapped_target.target.address]))
-    assert len(coarsened_targets) == 1
-    coarsened_target = coarsened_targets[0]
-    resolve = await Get(CoursierResolveKey, CoarsenedTargets([coarsened_target]))
-    output_file = compute_output_jar_filename(coarsened_target)
+    bsp_target_name = request.bsp_target_id.uri[len("pants:") :]
+    if bsp_target_name not in bsp_build_targets.targets_mapping:
+        raise ValueError(f"Invalid BSP target name: {request.bsp_target_id}")
+    targets = await Get(
+        Targets,
+        AddressSpecs,
+        bsp_build_targets.targets_mapping[bsp_target_name].specs.address_specs,
+    )
+
+    coarsened_targets = await Get(CoarsenedTargets, Addresses(tgt.address for tgt in targets))
+    resolve = await Get(CoursierResolveKey, CoarsenedTargets, coarsened_targets)
 
     return HandleJavacOptionsResult(
         JavacOptionsItem(
             target=request.bsp_target_id,
             options=(),
-            classpath=(
-                build_root.pathlib_path.joinpath(
-                    f".pants.d/bsp/jvm/resolves/{resolve.name}/lib/{output_file}"
-                ).as_uri(),
-            ),
+            classpath=(),
             class_directory=build_root.pathlib_path.joinpath(
                 f".pants.d/bsp/jvm/resolves/{resolve.name}/classes"
             ).as_uri(),
@@ -216,7 +173,7 @@ def rules():
     return (
         *collect_rules(),
         UnionRule(BSPLanguageSupport, JavaBSPLanguageSupport),
-        UnionRule(BSPBuildTargetsRequest, JavaBSPBuildTargetsRequest),
+        UnionRule(BSPBuildTargetsMetadataRequest, JavaBSPBuildTargetsMetadataRequest),
         UnionRule(BSPHandlerMapping, JavacOptionsHandlerMapping),
         UnionRule(BSPCompileFieldSet, JavaBSPCompileFieldSet),
     )

src/python/pants/backend/scala/bsp/rules.py

@@ -1,5 +1,7 @@
 # Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
+from __future__ import annotations
+
 import dataclasses
 import logging
 import os
@@ -12,36 +14,26 @@
     ScalacOptionsResult,
     ScalaPlatform,
 )
-from pants.backend.scala.compile.scalac import compute_output_jar_filename
-from pants.backend.scala.dependency_inference.symbol_mapper import AllScalaTargets
 from pants.backend.scala.subsystems.scala import ScalaSubsystem
 from pants.backend.scala.target_types import ScalaSourceField
 from pants.base.build_root import BuildRoot
-from pants.bsp.context import BSPContext
+from pants.base.specs import AddressSpecs
 from pants.bsp.protocol import BSPHandlerMapping
-from pants.bsp.spec.base import (
-    BuildTarget,
-    BuildTargetCapabilities,
-    BuildTargetIdentifier,
-    StatusCode,
-)
+from pants.bsp.spec.base import BuildTarget, BuildTargetIdentifier, StatusCode
 from pants.bsp.util_rules.compile import BSPCompileFieldSet, BSPCompileResult
 from pants.bsp.util_rules.lifecycle import BSPLanguageSupport
-from pants.bsp.util_rules.targets import BSPBuildTargets, BSPBuildTargetsRequest
-from pants.build_graph.address import Address, AddressInput
+from pants.bsp.util_rules.targets import (
+    BSPBuildTargetsMetadataRequest,
+    BSPBuildTargetsMetadataResult,
+    BSPBuildTargetsNew,
+)
 from pants.engine.addresses import Addresses
 from pants.engine.fs import EMPTY_DIGEST, AddPrefix, CreateDigest, Digest, DigestEntries
-from pants.engine.internals.native_engine import MergeDigests, Snapshot
+from pants.engine.internals.native_engine import Snapshot
 from pants.engine.internals.selectors import Get, MultiGet
 from pants.engine.rules import collect_rules, rule
-from pants.engine.target import (
-    CoarsenedTargets,
-    Dependencies,
-    DependenciesRequest,
-    Target,
-    WrappedTarget,
-)
-from pants.engine.unions import UnionMembership, UnionRule
+from pants.engine.target import CoarsenedTargets, FieldSet, Target, Targets
+from pants.engine.unions import UnionRule
 from pants.jvm.compile import (
     ClasspathEntryRequest,
     ClasspathEntryRequestFactory,
@@ -63,8 +55,18 @@ class ScalaBSPLanguageSupport(BSPLanguageSupport):
     can_compile = True
 
 
-class ScalaBSPBuildTargetsRequest(BSPBuildTargetsRequest):
-    pass
+@dataclass(frozen=True)
+class ScalaMetadataFieldSet(FieldSet):
+    required_fields = (ScalaSourceField, JvmResolveField)
+
+    source: ScalaSourceField
+    resolve: JvmResolveField
+
+
+class ScalaBSPBuildTargetsMetadataRequest(BSPBuildTargetsMetadataRequest):
+    language_id = LANGUAGE_ID
+    can_merge_metadata_from = ("java",)
+    field_set_type = ScalaMetadataFieldSet
 
 
 @dataclass(frozen=True)
@@ -121,75 +123,37 @@ async def materialize_scala_runtime_jars(
 
 
 @rule
-async def bsp_resolve_one_scala_build_target(
-    request: ResolveScalaBSPBuildTargetRequest,
+async def bsp_resolve_scala_metadata(
+    request: ScalaBSPBuildTargetsMetadataRequest,
     jvm: JvmSubsystem,
     scala: ScalaSubsystem,
-    union_membership: UnionMembership,
     build_root: BuildRoot,
-) -> ResolveScalaBSPBuildTargetResult:
-    resolve = request.target[JvmResolveField].normalized_value(jvm)
+) -> BSPBuildTargetsMetadataResult:
+    resolves = {fs.resolve.normalized_value(jvm) for fs in request.field_sets}
+    if len(resolves) > 1:
+        raise ValueError("Cannot provide Scala metadata for multiple resolves.")
+    resolve = list(resolves)[0]
     scala_version = scala.version_for_resolve(resolve)
 
-    dep_addrs, scala_runtime = await MultiGet(
-        Get(Addresses, DependenciesRequest(request.target[Dependencies])),
-        Get(MaterializeScalaRuntimeJarsResult, MaterializeScalaRuntimeJarsRequest(scala_version)),
+    scala_runtime = await Get(
+        MaterializeScalaRuntimeJarsResult, MaterializeScalaRuntimeJarsRequest(scala_version)
     )
-    impls = union_membership.get(BSPCompileFieldSet)
-
-    reported_deps = []
-    for dep_addr in dep_addrs:
-        if dep_addr == request.target.address:
-            continue
-
-        wrapped_dep_tgt = await Get(WrappedTarget, Address, dep_addr)
-        dep_tgt = wrapped_dep_tgt.target
-        for impl in impls:
-            if impl.is_applicable(dep_tgt):
-                reported_deps.append(BuildTargetIdentifier.from_address(dep_tgt.address))
-                break
 
     scala_jar_uris = tuple(
         build_root.pathlib_path.joinpath(".pants.d/bsp").joinpath(p).as_uri()
         for p in scala_runtime.content.files
     )
-    bsp_target = BuildTarget(
-        id=BuildTargetIdentifier.from_address(request.target.address),
-        display_name=str(request.target.address),
-        base_directory=None,
-        tags=(),
-        capabilities=BuildTargetCapabilities(
-            can_compile=True,
-        ),
-        language_ids=(LANGUAGE_ID,),
-        dependencies=tuple(reported_deps),
-        data_kind="scala",
-        data=ScalaBuildTarget(
-            scala_organization="unknown",
+
+    return BSPBuildTargetsMetadataResult(
+        metadata=ScalaBuildTarget(
+            scala_organization="org.scala-lang",
             scala_version=scala_version,
             scala_binary_version=".".join(scala_version.split(".")[0:2]),
             platform=ScalaPlatform.JVM,
             jars=scala_jar_uris,
         ),
-    )
-    return ResolveScalaBSPBuildTargetResult(bsp_target, scala_runtime=scala_runtime.content)
-
-
-@rule
-async def bsp_resolve_all_scala_build_targets(
-    _: ScalaBSPBuildTargetsRequest,
-    all_scala_targets: AllScalaTargets,
-    bsp_context: BSPContext,
-) -> BSPBuildTargets:
-    if LANGUAGE_ID not in bsp_context.client_params.capabilities.language_ids:
-        return BSPBuildTargets()
-    build_targets = await MultiGet(
-        Get(ResolveScalaBSPBuildTargetResult, ResolveScalaBSPBuildTargetRequest(tgt))
-        for tgt in all_scala_targets
-    )
-    output_digest = await Get(Digest, MergeDigests([d.scala_runtime.digest for d in build_targets]))
-    return BSPBuildTargets(
-        targets=tuple(btgt.build_target for btgt in build_targets), digest=output_digest
+        can_compile=True,
+        digest=scala_runtime.content.digest,
     )
 
 
@@ -219,23 +183,24 @@ class HandleScalacOptionsResult:
 async def handle_bsp_scalac_options_request(
     request: HandleScalacOptionsRequest,
     build_root: BuildRoot,
+    bsp_build_targets: BSPBuildTargetsNew,
 ) -> HandleScalacOptionsResult:
-    wrapped_target = await Get(WrappedTarget, AddressInput, request.bsp_target_id.address_input)
-    coarsened_targets = await Get(CoarsenedTargets, Addresses([wrapped_target.target.address]))
-    assert len(coarsened_targets) == 1
-    coarsened_target = coarsened_targets[0]
-    resolve = await Get(CoursierResolveKey, CoarsenedTargets([coarsened_target]))
-    output_file = compute_output_jar_filename(coarsened_target)
+    bsp_target_name = request.bsp_target_id.uri[len("pants:") :]
+    if bsp_target_name not in bsp_build_targets.targets_mapping:
+        raise ValueError(f"Invalid BSP target name: {request.bsp_target_id}")
+    targets = await Get(
+        Targets,
+        AddressSpecs,
+        bsp_build_targets.targets_mapping[bsp_target_name].specs.address_specs,
+    )
+    coarsened_targets = await Get(CoarsenedTargets, Addresses(tgt.address for tgt in targets))
+    resolve = await Get(CoursierResolveKey, CoarsenedTargets, coarsened_targets)
 
     return HandleScalacOptionsResult(
         ScalacOptionsItem(
             target=request.bsp_target_id,
             options=(),
-            classpath=(
-                build_root.pathlib_path.joinpath(
-                    f".pants.d/bsp/jvm/resolves/{resolve.name}/lib/{output_file}"
-                ).as_uri(),
-            ),
+            classpath=(),
             class_directory=build_root.pathlib_path.joinpath(
                 f".pants.d/bsp/jvm/resolves/{resolve.name}/classes"
             ).as_uri(),
@@ -299,7 +264,7 @@ def rules():
     return (
         *collect_rules(),
         UnionRule(BSPLanguageSupport, ScalaBSPLanguageSupport),
-        UnionRule(BSPBuildTargetsRequest, ScalaBSPBuildTargetsRequest),
+        UnionRule(BSPBuildTargetsMetadataRequest, ScalaBSPBuildTargetsMetadataRequest),
         UnionRule(BSPHandlerMapping, ScalacOptionsHandlerMapping),
         UnionRule(BSPCompileFieldSet, ScalaBSPCompileFieldSet),
     )