From 31fae9e8e6687cbaf0dfe55a466696210c80be96 Mon Sep 17 00:00:00 2001
From: Fabian Meumertzheim <fabian@meumertzhe.im>
Date: Fri, 16 Feb 2024 13:48:18 -0800
Subject: [PATCH] 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 #21343.

PiperOrigin-RevId: 607790056
Change-Id: Ia08c834ae3e5b1151607459408cdfea85d47314f
---
 .../lib/rules/java/JavaStarlarkCommon.java    |  6 ++--
 .../lib/rules/java/JavaTargetAttributes.java  | 32 ++++++++++++++++---
 2 files changed, 32 insertions(+), 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
index 52654a50bad464..3f82911e020668 100644
--- 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
index e78eabdf36f08d..57e63b28bd7936 100644
--- 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,6 +271,7 @@ public Builder setStrictJavaDeps(StrictDepsMode strictDeps) {
     @CanIgnoreReturnValue
     public Builder addDirectJars(NestedSet<Artifact> directJars) {
       Preconditions.checkArgument(!built);
+      Preconditions.checkArgument(prependDirectJars);
       this.directJarsBuilder.addTransitive(directJars);
       return this;
     }
@@ -258,6 +279,7 @@ public Builder addDirectJars(NestedSet<Artifact> directJars) {
     @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(),