Share classpath NestedSet between full and header compile actions

The Starlark Java rules already add the direct JARs to the transitive classpath depset and pass the same depset to the full and the header compile action. By allowing them to bypass the logic that prepends the direct jars to the compile-time classpath, both actions retain the same `NestedSet` instead of both retaining a new one with identical elements. Closes bazelbuild#21343. PiperOrigin-RevId: 607790056 Change-Id: Ia08c834ae3e5b1151607459408cdfea85d47314f
bazel-io · Feb 16, 2024 · 31fae9e · 31fae9e
1 parent 014292b
commit 31fae9e
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 6 deletions.
diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/JavaStarlarkCommon.java b/src/main/java/com/google/devtools/build/lib/rules/java/JavaStarlarkCommon.java
@@ -141,7 +141,8 @@ public void createHeaderCompilationAction(
             .addSourceJars(Sequence.cast(sourceJars, Artifact.class, "source_jars"))
             .addSourceFiles(sourceFiles.toList(Artifact.class))
             .addDirectJars(directJars.getSet(Artifact.class))
-            .addCompileTimeClassPathEntries(compileTimeClasspath.getSet(Artifact.class))
+            .setCompileTimeClassPathEntriesWithPrependedDirectJars(
+                compileTimeClasspath.getSet(Artifact.class))
             .setStrictJavaDeps(getStrictDepsMode(Ascii.toUpperCase(strictDepsMode)))
             .setTargetLabel(targetLabel)
             .setInjectingRuleKind(
@@ -214,7 +215,8 @@ public void createCompilationAction(
             .addSourceJars(Sequence.cast(sourceJars, Artifact.class, "source_jars"))
             .addSourceFiles(Depset.noneableCast(sourceFiles, Artifact.class, "sources").toList())
             .addDirectJars(directJars.getSet(Artifact.class))
-            .addCompileTimeClassPathEntries(compileTimeClasspath.getSet(Artifact.class))
+            .setCompileTimeClassPathEntriesWithPrependedDirectJars(
+                compileTimeClasspath.getSet(Artifact.class))
             .addClassPathResources(
                 Sequence.cast(classpathResources, Artifact.class, "classpath_resources"))
             .setStrictJavaDeps(getStrictDepsMode(Ascii.toUpperCase(strictDepsMode)))

diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/JavaTargetAttributes.java b/src/main/java/com/google/devtools/build/lib/rules/java/JavaTargetAttributes.java
@@ -88,6 +88,8 @@ public static class Builder {
 
     private final NestedSetBuilder<Artifact> excludedArtifacts = NestedSetBuilder.naiveLinkOrder();
 
+    private boolean prependDirectJars = true;
+
     private boolean built = false;
 
     private final JavaSemantics semantics;
@@ -179,6 +181,24 @@ public Builder addCompileTimeClassPathEntries(NestedSet<Artifact> entries) {
       return this;
     }
 
+    /**
+     * Avoids prepending the direct jars to the compile-time classpath when building the attributes,
+     * assuming that they have already been prepended. This avoids creating a new {@link NestedSet}
+     * instance.
+     *
+     * <p>After this method is called, {@link #addDirectJar(Artifact)} and {@link
+     * #addDirectJars(NestedSet)} will throw an exception.
+     */
+    @CanIgnoreReturnValue
+    public Builder setCompileTimeClassPathEntriesWithPrependedDirectJars(
+        NestedSet<Artifact> entries) {
+      Preconditions.checkArgument(!built);
+      Preconditions.checkArgument(compileTimeClassPathBuilder.isEmpty());
+      prependDirectJars = false;
+      compileTimeClassPathBuilder.addTransitive(entries);
+      return this;
+    }
+
     @CanIgnoreReturnValue
     public Builder setTargetLabel(Label targetLabel) {
       Preconditions.checkArgument(!built);
@@ -251,13 +271,15 @@ public Builder setStrictJavaDeps(StrictDepsMode strictDeps) {
     @CanIgnoreReturnValue
     public Builder addDirectJars(NestedSet<Artifact> directJars) {
       Preconditions.checkArgument(!built);
+      Preconditions.checkArgument(prependDirectJars);
       this.directJarsBuilder.addTransitive(directJars);
       return this;
     }
 
     @CanIgnoreReturnValue
     public Builder addDirectJar(Artifact directJar) {
       Preconditions.checkArgument(!built);
+      Preconditions.checkArgument(prependDirectJars);
       this.directJarsBuilder.add(directJar);
       return this;
     }
@@ -323,10 +345,12 @@ public JavaTargetAttributes build() {
       built = true;
       NestedSet<Artifact> directJars = directJarsBuilder.build();
       NestedSet<Artifact> compileTimeClassPath =
-          NestedSetBuilder.<Artifact>naiveLinkOrder()
-              .addTransitive(directJars)
-              .addTransitive(compileTimeClassPathBuilder.build())
-              .build();
+          prependDirectJars
+              ? NestedSetBuilder.<Artifact>naiveLinkOrder()
+                  .addTransitive(directJars)
+                  .addTransitive(compileTimeClassPathBuilder.build())
+                  .build()
+              : compileTimeClassPathBuilder.build();
       return new JavaTargetAttributes(
           ImmutableSet.copyOf(sourceFiles),
           runtimeClassPath.build(),