Revert "Don't eagerly flatten a NestedSet in RepoMappingManifestAction (#18419)" (#18886)

This reverts commit e023bd3.

Author: iancha1992

src/main/java/com/google/devtools/build/lib/analysis/BUILD

@@ -995,9 +995,9 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/actions:commandline_item",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/collect/nestedset",
-        "//src/main/java/com/google/devtools/build/lib/packages",
         "//src/main/java/com/google/devtools/build/lib/util",
         "//src/main/java/net/starlark/java/eval",
+        "//third_party:auto_value",
         "//third_party:guava",
         "//third_party:jsr305",
     ],

src/main/java/com/google/devtools/build/lib/analysis/RepoMappingManifestAction.java

@@ -13,69 +13,61 @@
 // limitations under the License.
 package com.google.devtools.build.lib.analysis;
 
-import static com.google.common.collect.ImmutableSet.toImmutableSet;
-import static com.google.common.collect.ImmutableSortedMap.toImmutableSortedMap;
 import static java.nio.charset.StandardCharsets.ISO_8859_1;
 import static java.util.Comparator.comparing;
 
-import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.ImmutableSortedMap;
+import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.actions.ActionExecutionContext;
 import com.google.devtools.build.lib.actions.ActionKeyContext;
 import com.google.devtools.build.lib.actions.ActionOwner;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.Artifact.ArtifactExpander;
 import com.google.devtools.build.lib.actions.CommandLineExpansionException;
-import com.google.devtools.build.lib.actions.CommandLineItem.MapFn;
 import com.google.devtools.build.lib.actions.ExecException;
 import com.google.devtools.build.lib.analysis.actions.AbstractFileWriteAction;
 import com.google.devtools.build.lib.analysis.actions.DeterministicWriter;
-import com.google.devtools.build.lib.cmdline.RepositoryMapping;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
-import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.collect.nestedset.Order;
-import com.google.devtools.build.lib.packages.Package;
 import com.google.devtools.build.lib.util.Fingerprint;
 import java.io.PrintWriter;
-import java.util.Map.Entry;
+import java.util.List;
 import java.util.UUID;
 import javax.annotation.Nullable;
 import net.starlark.java.eval.EvalException;
 
 /** Creates a manifest file describing the repos and mappings relevant for a runfile tree. */
-public final class RepoMappingManifestAction extends AbstractFileWriteAction {
-
+public class RepoMappingManifestAction extends AbstractFileWriteAction {
   private static final UUID MY_UUID = UUID.fromString("458e351c-4d30-433d-b927-da6cddd4737f");
 
-  // Uses MapFn's args parameter just like Fingerprint#addString to compute a cacheable fingerprint
-  // of just the repo name and mapping of a given Package.
-  private static final MapFn<Package> REPO_AND_MAPPING_DIGEST_FN =
-      (pkg, args) -> {
-        args.accept(pkg.getPackageIdentifier().getRepository().getName());
+  private final ImmutableList<Entry> entries;
+  private final String workspaceName;
 
-        var mapping = pkg.getRepositoryMapping().entries();
-        args.accept(String.valueOf(mapping.size()));
-        mapping.forEach(
-            (apparentName, canonicalName) -> {
-              args.accept(apparentName);
-              args.accept(canonicalName.getName());
-            });
-      };
+  /** An entry in the repo mapping manifest file. */
+  @AutoValue
+  public abstract static class Entry {
+    public static Entry of(
+        RepositoryName sourceRepo, String targetRepoApparentName, RepositoryName targetRepo) {
+      return new AutoValue_RepoMappingManifestAction_Entry(
+          sourceRepo, targetRepoApparentName, targetRepo);
+    }
 
-  private final NestedSet<Package> transitivePackages;
-  private final NestedSet<Artifact> runfilesArtifacts;
-  private final String workspaceName;
+    public abstract RepositoryName sourceRepo();
+
+    public abstract String targetRepoApparentName();
+
+    public abstract RepositoryName targetRepo();
+  }
 
   public RepoMappingManifestAction(
-      ActionOwner owner,
-      Artifact output,
-      NestedSet<Package> transitivePackages,
-      NestedSet<Artifact> runfilesArtifacts,
-      String workspaceName) {
+      ActionOwner owner, Artifact output, List<Entry> entries, String workspaceName) {
     super(owner, NestedSetBuilder.emptySet(Order.STABLE_ORDER), output, /*makeExecutable=*/ false);
-    this.transitivePackages = transitivePackages;
-    this.runfilesArtifacts = runfilesArtifacts;
+    this.entries =
+        ImmutableList.sortedCopyOf(
+            comparing((Entry e) -> e.sourceRepo().getName())
+                .thenComparing(Entry::targetRepoApparentName),
+            entries);
     this.workspaceName = workspaceName;
   }
 
@@ -86,7 +78,7 @@ public String getMnemonic() {
 
   @Override
   protected String getRawProgressMessage() {
-    return "Writing repo mapping manifest for " + getOwner().getLabel();
+    return "writing repo mapping manifest for " + getOwner().getLabel();
   }
 
   @Override
@@ -96,61 +88,35 @@ protected void computeKey(
       Fingerprint fp)
       throws CommandLineExpansionException, EvalException, InterruptedException {
     fp.addUUID(MY_UUID);
-    actionKeyContext.addNestedSetToFingerprint(REPO_AND_MAPPING_DIGEST_FN, fp, transitivePackages);
-    actionKeyContext.addNestedSetToFingerprint(fp, runfilesArtifacts);
     fp.addString(workspaceName);
+    for (Entry entry : entries) {
+      fp.addString(entry.sourceRepo().getName());
+      fp.addString(entry.targetRepoApparentName());
+      fp.addString(entry.targetRepo().getName());
+    }
   }
 
   @Override
   public DeterministicWriter newDeterministicWriter(ActionExecutionContext ctx)
       throws InterruptedException, ExecException {
     return out -> {
-      PrintWriter writer = new PrintWriter(out, /* autoFlush= */ false, ISO_8859_1);
-
-      ImmutableSet<RepositoryName> reposContributingRunfiles =
-          runfilesArtifacts.toList().stream()
-              .filter(a -> a.getOwner() != null)
-              .map(a -> a.getOwner().getRepository())
-              .collect(toImmutableSet());
-      transitivePackages.toList().stream()
-          .collect(
-              toImmutableSortedMap(
-                  comparing(RepositoryName::getName),
-                  pkg -> pkg.getPackageIdentifier().getRepository(),
-                  Package::getRepositoryMapping,
-                  // All packages in a given repository have the same repository mapping, so the
-                  // particular way of resolving duplicates does not matter.
-                  (first, second) -> first))
-          .forEach(
-              (repoName, mapping) ->
-                  writeRepoMapping(writer, reposContributingRunfiles, repoName, mapping));
+      PrintWriter writer = new PrintWriter(out, /*autoFlush=*/ false, ISO_8859_1);
+      for (Entry entry : entries) {
+        if (entry.targetRepoApparentName().isEmpty()) {
+          // The apparent repo name can only be empty for the main repo. We skip this line as
+          // Rlocation paths can't reference an empty apparent name anyway.
+          continue;
+        }
+        // The canonical name of the main repo is the empty string, which is not a valid name for a
+        // directory, so the "workspace name" is used the name of the directory under the runfiles
+        // tree for it.
+        String targetRepoDirectoryName =
+            entry.targetRepo().isMain() ? workspaceName : entry.targetRepo().getName();
+        writer.format(
+            "%s,%s,%s\n",
+            entry.sourceRepo().getName(), entry.targetRepoApparentName(), targetRepoDirectoryName);
+      }
       writer.flush();
     };
   }
-
-  private void writeRepoMapping(
-      PrintWriter writer,
-      ImmutableSet<RepositoryName> reposContributingRunfiles,
-      RepositoryName repoName,
-      RepositoryMapping repoMapping) {
-    for (Entry<String, RepositoryName> mappingEntry :
-        ImmutableSortedMap.copyOf(repoMapping.entries()).entrySet()) {
-      if (mappingEntry.getKey().isEmpty()) {
-        // The apparent repo name can only be empty for the main repo. We skip this line as
-        // Rlocation paths can't reference an empty apparent name anyway.
-        continue;
-      }
-      if (!reposContributingRunfiles.contains(mappingEntry.getValue())) {
-        // We only write entries for repos that actually contribute runfiles.
-        continue;
-      }
-      // The canonical name of the main repo is the empty string, which is not a valid name for a
-      // directory, so the "workspace name" is used the name of the directory under the runfiles
-      // tree for it.
-      String targetRepoDirectoryName =
-          mappingEntry.getValue().isMain() ? workspaceName : mappingEntry.getValue().getName();
-      writer.format(
-          "%s,%s,%s\n", repoName.getName(), mappingEntry.getKey(), targetRepoDirectoryName);
-    }
-  }
 }

src/main/java/com/google/devtools/build/lib/analysis/RunfilesSupport.java

@@ -14,27 +14,34 @@
 
 package com.google.devtools.build.lib.analysis;
 
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.actions.ActionEnvironment;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.CommandLine;
+import com.google.devtools.build.lib.analysis.RepoMappingManifestAction.Entry;
 import com.google.devtools.build.lib.analysis.SourceManifestAction.ManifestType;
 import com.google.devtools.build.lib.analysis.actions.ActionConstructionContext;
 import com.google.devtools.build.lib.analysis.actions.SymlinkTreeAction;
 import com.google.devtools.build.lib.analysis.config.BuildConfigurationValue;
 import com.google.devtools.build.lib.analysis.config.RunUnder;
+import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
+import com.google.devtools.build.lib.packages.Package;
 import com.google.devtools.build.lib.packages.TargetUtils;
 import com.google.devtools.build.lib.packages.Type;
 import com.google.devtools.build.lib.packages.semantics.BuildLanguageOptions;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.util.Collection;
+import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
@@ -551,9 +558,45 @@ private static Artifact createRepoMappingManifestAction(
             new RepoMappingManifestAction(
                 ruleContext.getActionOwner(),
                 repoMappingManifest,
-                ruleContext.getTransitivePackagesForRunfileRepoMappingManifest(),
-                runfiles.getAllArtifacts(),
+                collectRepoMappings(
+                    Preconditions.checkNotNull(
+                        ruleContext.getTransitivePackagesForRunfileRepoMappingManifest()),
+                    runfiles),
                 ruleContext.getWorkspaceName()));
     return repoMappingManifest;
   }
+
+  /** Returns the list of entries (unsorted) that should appear in the repo mapping manifest. */
+  private static ImmutableList<Entry> collectRepoMappings(
+      NestedSet<Package> transitivePackages, Runfiles runfiles) {
+    // NOTE: It might appear that the flattening of `transitivePackages` is better suited to the
+    // execution phase rather than here in the analysis phase, but we can't do that since it would
+    // necessitate storing `transitivePackages` in an action, which breaks skyframe serialization
+    // since packages cannot be serialized here.
+
+    ImmutableSet<RepositoryName> reposContributingRunfiles =
+        runfiles.getAllArtifacts().toList().stream()
+            .filter(a -> a.getOwner() != null)
+            .map(a -> a.getOwner().getRepository())
+            .collect(toImmutableSet());
+    Set<RepositoryName> seenRepos = new HashSet<>();
+    ImmutableList.Builder<Entry> entries = ImmutableList.builder();
+    for (Package pkg : transitivePackages.toList()) {
+      if (!seenRepos.add(pkg.getPackageIdentifier().getRepository())) {
+        // Any package from the same repo would have the same repo mapping.
+        continue;
+      }
+      for (Map.Entry<String, RepositoryName> repoMappingEntry :
+          pkg.getRepositoryMapping().entries().entrySet()) {
+        if (reposContributingRunfiles.contains(repoMappingEntry.getValue())) {
+          entries.add(
+              Entry.of(
+                  pkg.getPackageIdentifier().getRepository(),
+                  repoMappingEntry.getKey(),
+                  repoMappingEntry.getValue()));
+        }
+      }
+    }
+    return entries.build();
+  }
 }

src/test/java/com/google/devtools/build/lib/analysis/BUILD

@@ -361,14 +361,16 @@ java_test(
     srcs = ["RunfilesRepoMappingManifestTest.java"],
     deps = [
         "//src/main/java/com/google/devtools/build/lib/actions",
-        "//src/main/java/com/google/devtools/build/lib/actions:commandline_item",
+        "//src/main/java/com/google/devtools/build/lib/analysis:blaze_directories",
         "//src/main/java/com/google/devtools/build/lib/analysis:repo_mapping_manifest_action",
+        "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:resolution",
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:resolution_impl",
         "//src/main/java/com/google/devtools/build/lib/bazel/repository:repository_options",
         "//src/main/java/com/google/devtools/build/lib/skyframe:precomputed_value",
-        "//src/main/java/com/google/devtools/build/lib/util",
+        "//src/main/java/com/google/devtools/build/lib/skyframe:sky_functions",
         "//src/main/java/com/google/devtools/build/lib/vfs",
-        "//src/main/java/net/starlark/java/eval",
+        "//src/main/java/com/google/devtools/build/skyframe",
+        "//src/main/java/com/google/devtools/build/skyframe:skyframe-objects",
         "//src/test/java/com/google/devtools/build/lib/analysis/util",
         "//src/test/java/com/google/devtools/build/lib/bazel/bzlmod:util",
         "//third_party:guava",