Do not retain transitive data in AndroidLocalTestBase.

The argument strings and artifacts were both transitive and flattened, causing O(N^2) memory consumption.

PiperOrigin-RevId: 156083738

3 files changed, 82 insertions, 66 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/rules/android/AndroidLibrary.java b/src/main/java/com/google/devtools/build/lib/rules/android/AndroidLibrary.java
index eae1f2dfb4..a6d7f39718 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/android/AndroidLibrary.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/android/AndroidLibrary.java
@@ -67,7 +67,9 @@ public abstract class AndroidLibrary implements RuleConfiguredTargetFactory {
       return null;
     }
     checkResourceInlining(ruleContext);
-    NestedSetBuilder<Aar> transitiveAars = collectTransitiveAars(ruleContext);
+    NestedSetBuilder<Aar> transitiveAars = NestedSetBuilder.naiveLinkOrder();
+    NestedSetBuilder<Artifact> transitiveAarArtifacts = NestedSetBuilder.stableOrder();
+    collectTransitiveAars(ruleContext, transitiveAars, transitiveAarArtifacts);
 
     NestedSetBuilder<Artifact> proguardConfigsbuilder = NestedSetBuilder.stableOrder();
     proguardConfigsbuilder.addTransitive(new ProguardLibrary(ruleContext).collectProguardSpecs());
@@ -148,12 +150,12 @@ public abstract class AndroidLibrary implements RuleConfiguredTargetFactory {
       // applicationManifest has already been checked for nullness above in this method
       ApplicationManifest applicationManifest = androidSemantics.getManifestForRule(ruleContext);
       aar = Aar.create(aarOut, applicationManifest.getManifest());
-      transitiveAars.add(aar);
+      addAarToProvider(aar, transitiveAars, transitiveAarArtifacts);
     } else if (AndroidCommon.getAndroidResources(ruleContext) != null) {
       primaryResources = Iterables.getOnlyElement(
           AndroidCommon.getAndroidResources(ruleContext).getDirectAndroidResources());
       aar = Aar.create(aarOut, primaryResources.getManifest());
-      transitiveAars.add(aar);
+      addAarToProvider(aar, transitiveAars, transitiveAarArtifacts);
     } else {
       // there are no local resources and resources attribute was not specified either
       aar = null;
@@ -234,12 +236,26 @@ public abstract class AndroidLibrary implements RuleConfiguredTargetFactory {
     if (!JavaCommon.isNeverLink(ruleContext)) {
       builder.add(
           AndroidLibraryAarProvider.class,
-          AndroidLibraryAarProvider.create(aar, transitiveAars.build()));
+          AndroidLibraryAarProvider.create(
+              aar, transitiveAars.build(), transitiveAarArtifacts.build()));
     }
 
     return builder.build();
   }
 
+  private void addAarToProvider(
+      Aar aar,
+      NestedSetBuilder<Aar> transitiveAars,
+      NestedSetBuilder<Artifact> transitiveAarArtifacts) {
+    transitiveAars.add(aar);
+    if (aar.getAar() != null) {
+      transitiveAarArtifacts.add(aar.getAar());
+    }
+    if (aar.getManifest() != null) {
+      transitiveAarArtifacts.add(aar.getManifest());
+    }
+  }
+
   private void checkResourceInlining(RuleContext ruleContext) {
     AndroidResourcesProvider resources = AndroidCommon.getAndroidResources(ruleContext);
     if (resources == null) {
@@ -256,13 +272,15 @@ public abstract class AndroidLibrary implements RuleConfiguredTargetFactory {
     }
   }
 
-  private NestedSetBuilder<Aar> collectTransitiveAars(RuleContext ruleContext) {
-    NestedSetBuilder<Aar> builder = NestedSetBuilder.naiveLinkOrder();
+  private void collectTransitiveAars(
+      RuleContext ruleContext,
+      NestedSetBuilder<Aar> transitiveAars,
+      NestedSetBuilder<Artifact> transitiveAarArtifacts) {
     for (AndroidLibraryAarProvider library : AndroidCommon.getTransitivePrerequisites(
         ruleContext, Mode.TARGET, AndroidLibraryAarProvider.class)) {
-      builder.addTransitive(library.getTransitiveAars());
+      transitiveAars.addTransitive(library.getTransitiveAars());
+      transitiveAarArtifacts.addTransitive(library.getTransitiveAarArtifacts());
     }
-    return builder;
   }
 
   private NestedSetBuilder<Artifact> collectTransitiveResourceJars(RuleContext ruleContext) {
diff --git a/src/main/java/com/google/devtools/build/lib/rules/android/AndroidLibraryAarProvider.java b/src/main/java/com/google/devtools/build/lib/rules/android/AndroidLibraryAarProvider.java
index ff9236e9ce..d4a8a8afdd 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/android/AndroidLibraryAarProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/android/AndroidLibraryAarProvider.java
@@ -28,14 +28,19 @@ import javax.annotation.Nullable;
 @Immutable
 public abstract class AndroidLibraryAarProvider implements TransitiveInfoProvider {
 
-  public static AndroidLibraryAarProvider create(@Nullable Aar aar, NestedSet<Aar> transitiveAars) {
-    return new AutoValue_AndroidLibraryAarProvider(aar, transitiveAars);
+  public static AndroidLibraryAarProvider create(
+      @Nullable Aar aar,
+      NestedSet<Aar> transitiveAars,
+      NestedSet<Artifact> transitiveAarArtifacts) {
+    return new AutoValue_AndroidLibraryAarProvider(aar, transitiveAars, transitiveAarArtifacts);
   }
 
   @Nullable public abstract Aar getAar();
 
   public abstract NestedSet<Aar> getTransitiveAars();
 
+  public abstract NestedSet<Artifact> getTransitiveAarArtifacts();
+
   /** The .aar file and associated AndroidManifest.xml contributed by a single target. */
   @AutoValue
   @Immutable
diff --git a/src/main/java/com/google/devtools/build/lib/rules/android/AndroidLocalTestBase.java b/src/main/java/com/google/devtools/build/lib/rules/android/AndroidLocalTestBase.java
index 1ef476f87f..ddd4e84c2c 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/android/AndroidLocalTestBase.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/android/AndroidLocalTestBase.java
@@ -16,7 +16,7 @@ package com.google.devtools.build.lib.rules.android;
 import static com.google.devtools.build.lib.packages.BuildType.LABEL_LIST;
 import static com.google.devtools.build.lib.rules.java.DeployArchiveBuilder.Compression.COMPRESSED;
 
-import com.google.common.base.Joiner;
+import com.google.common.base.Function;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
@@ -31,6 +31,7 @@ import com.google.devtools.build.lib.analysis.Runfiles;
 import com.google.devtools.build.lib.analysis.RunfilesProvider;
 import com.google.devtools.build.lib.analysis.RunfilesSupport;
 import com.google.devtools.build.lib.analysis.TransitiveInfoCollection;
+import com.google.devtools.build.lib.analysis.actions.CustomCommandLine;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.rules.RuleConfiguredTargetFactory;
@@ -54,12 +55,9 @@ import com.google.devtools.build.lib.rules.java.JavaTargetAttributes;
 import com.google.devtools.build.lib.rules.java.SingleJarActionBuilder;
 import com.google.devtools.build.lib.rules.java.proto.GeneratedExtensionRegistryProvider;
 import com.google.devtools.build.lib.syntax.Type;
-import com.google.devtools.build.lib.util.Preconditions;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.util.ArrayList;
-import java.util.LinkedHashSet;
 import java.util.List;
-import java.util.Set;
 
 /**
  * An base implementation for the "android_local_test" rule.
@@ -183,13 +181,46 @@ public abstract class AndroidLocalTestBase implements RuleConfiguredTargetFactor
                 ruleContext.getPrerequisites(
                     "deps", Mode.TARGET, AndroidLibraryAarProvider.class)));
 
-    Runfiles defaultRunfiles =
-        collectDefaultRunfiles(ruleContext, javaCommon, filesToBuild, androidAarProviders);
+    NestedSetBuilder<Aar> transitiveAarsBuilder = NestedSetBuilder.naiveLinkOrder();
+    NestedSetBuilder<Aar> strictAarsBuilder = NestedSetBuilder.naiveLinkOrder();
+    NestedSetBuilder<Artifact> transitiveAarArtifactsBuilder = NestedSetBuilder.stableOrder();
+    for (AndroidLibraryAarProvider aarProvider : androidAarProviders) {
+      transitiveAarsBuilder.addTransitive(aarProvider.getTransitiveAars());
+      transitiveAarArtifactsBuilder.addTransitive(aarProvider.getTransitiveAarArtifacts());
+      if (aarProvider.getAar() != null) {
+        strictAarsBuilder.add(aarProvider.getAar());
+      }
+    }
+    NestedSet<Aar> transitiveAars = transitiveAarsBuilder.build();
+    NestedSet<Aar> strictAars = strictAarsBuilder.build();
+    NestedSet<Artifact> transitiveAarArtifacts = transitiveAarArtifactsBuilder.build();
 
-    ImmutableList<String> cmdLineArgs =
-        ImmutableList.of(
-            "--android_libraries=" + getTransitiveLibrariesArg(androidAarProviders),
-            "--strict_libraries=" + getStrictLibrariesArg(androidAarProviders));
+    Runfiles defaultRunfiles =
+        collectDefaultRunfiles(ruleContext, javaCommon, filesToBuild, transitiveAarArtifacts);
+
+    CustomCommandLine cmdLineArgs =
+        CustomCommandLine.builder()
+            .addJoinValues(
+                "--android_libraries",
+                ",",
+                transitiveAars,
+                new Function<Aar, String>() {
+                  @Override
+                  public String apply(Aar aar) {
+                    return aarCmdLineArg(aar);
+                  }
+                })
+            .addJoinValues(
+                "--strict_libraries",
+                ",",
+                strictAars,
+                new Function<Aar, String>() {
+                  @Override
+                  public String apply(Aar aar) {
+                    return aarCmdLineArg(aar);
+                  }
+                })
+            .build();
 
     RunfilesSupport runfilesSupport =
         RunfilesSupport.withExecutable(ruleContext, defaultRunfiles, executable, cmdLineArgs);
@@ -267,6 +298,12 @@ public abstract class AndroidLocalTestBase implements RuleConfiguredTargetFactor
         .build();
   }
 
+  private static String aarCmdLineArg(Aar aar) {
+    return aar.getManifest().getRootRelativePathString()
+        + ":"
+        + aar.getAar().getRootRelativePathString();
+  }
+
   protected abstract JavaSemantics createJavaSemantics();
 
   protected abstract void addExtraProviders(
@@ -337,7 +374,7 @@ public abstract class AndroidLocalTestBase implements RuleConfiguredTargetFactor
       RuleContext ruleContext,
       JavaCommon javaCommon,
       NestedSet<Artifact> filesToBuild,
-      Iterable<AndroidLibraryAarProvider> androidLibraryAarProviders) {
+      NestedSet<Artifact> transitiveAarArtifacts) {
     Runfiles.Builder builder = new Runfiles.Builder(ruleContext.getWorkspaceName());
     builder.addTransitiveArtifacts(filesToBuild);
     builder.addArtifacts(javaCommon.getJavaCompilationArtifacts().getRuntimeJars());
@@ -358,17 +395,7 @@ public abstract class AndroidLocalTestBase implements RuleConfiguredTargetFactor
 
     builder.addTargets(depsForRunfiles, JavaRunfilesProvider.TO_RUNFILES);
     builder.addTargets(depsForRunfiles, RunfilesProvider.DEFAULT_RUNFILES);
-
-    for (AndroidLibraryAarProvider aarProvider : androidLibraryAarProviders) {
-      if (aarProvider.getAar() != null) {
-        builder.addArtifact(aarProvider.getAar().getAar());
-        builder.addArtifact(aarProvider.getAar().getManifest());
-      }
-      for (Aar aar : aarProvider.getTransitiveAars()) {
-        builder.addArtifact(aar.getAar());
-        builder.addArtifact(aar.getManifest());
-      }
-    }
+    builder.addTransitiveArtifacts(transitiveAarArtifacts);
 
     if (ruleContext.getConfiguration().isCodeCoverageEnabled()) {
       Artifact instrumentedJar = javaCommon.getJavaCompilationArtifacts().getInstrumentedJar();
@@ -410,38 +437,4 @@ public abstract class AndroidLocalTestBase implements RuleConfiguredTargetFactor
     }
     return testClass;
   }
-
-  protected static String getTransitiveLibrariesArg(
-      Iterable<AndroidLibraryAarProvider> aarProviders) {
-    Set<String> args = new LinkedHashSet<>();
-
-    for (AndroidLibraryAarProvider aarProvider : aarProviders) {
-      for (Aar aar : aarProvider.getTransitiveAars()) {
-        Preconditions.checkNotNull(aar.getAar());
-        Preconditions.checkNotNull(aar.getManifest());
-        args.add(
-            aar.getManifest().getRootRelativePathString()
-                + ":"
-                + aar.getAar().getRootRelativePathString());
-      }
-    }
-    return Joiner.on(",").join(args);
-  }
-
-  protected static String getStrictLibrariesArg(Iterable<AndroidLibraryAarProvider> aarProviders) {
-    Set<String> args = new LinkedHashSet<>();
-
-    for (AndroidLibraryAarProvider aarProvider : aarProviders) {
-      Aar aar = aarProvider.getAar();
-      if (aar != null) {
-        Preconditions.checkNotNull(aar.getAar());
-        Preconditions.checkNotNull(aar.getManifest());
-        args.add(
-            aar.getManifest().getRootRelativePathString()
-                + ":"
-                + aar.getAar().getRootRelativePathString());
-      }
-    }
-    return Joiner.on(",").join(args);
-  }
 }