Attempt to make main repo mapping inverse more efficient

src/main/java/com/google/devtools/build/lib/cmdline/RepositoryMapping.java

@@ -14,53 +14,77 @@
 
 package com.google.devtools.build.lib.cmdline;
 
-import com.google.auto.value.AutoValue;
+import com.google.common.base.Objects;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Maps;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Optional;
 import javax.annotation.Nullable;
 
 /**
- * A class to distinguish repository mappings for repos from WORKSPACE and Bzlmod.
+ * Stores the mapping from apparent repo name to canonical repo name, from the viewpoint of an
+ * "owner repo".
  *
  * <p>For repositories from the WORKSPACE file, if the requested repo doesn't exist in the mapping,
  * we fallback to the requested name. For repositories from Bzlmod, we return null to let the caller
- * decide what to do. This class won't be needed if one day we don't define external repositories in
- * the WORKSPACE file since {@code fallback} would always be false.
+ * decide what to do.
  */
-@AutoValue
-public abstract class RepositoryMapping {
-
- // Always fallback to the requested name
+public class RepositoryMapping {
+ /* A repo mapping that always falls back to using the apparent name as the canonical name. */
  public static final RepositoryMapping ALWAYS_FALLBACK = createAllowingFallback(ImmutableMap.of());
 
+ private final ImmutableMap<String, RepositoryName> entries;
+ @Nullable private final RepositoryName ownerRepo;
+
+ private RepositoryMapping(
+ Map<String, RepositoryName> entries, @Nullable RepositoryName ownerRepo) {
+ this.entries = ImmutableMap.copyOf(entries);
+ this.ownerRepo = ownerRepo;
+ }
+
  /** Returns all the entries in this repo mapping. */
- public abstract ImmutableMap<String, RepositoryName> entries();
+ public final ImmutableMap<String, RepositoryName> entries() {
+ return entries;
+ }
 
  /**
  * The owner repo of this repository mapping. It is for providing useful debug information when
  * repository mapping fails due to enforcing strict dependency, therefore it's only recorded when
- * we don't fallback to the requested repo name.
+ * we don't fall back to the requested repo name.
  */
  @Nullable
- abstract RepositoryName ownerRepo();
+ public final RepositoryName ownerRepo() {
+ return ownerRepo;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (!(o instanceof RepositoryMapping)) {
+ return false;
+ }
+ RepositoryMapping that = (RepositoryMapping) o;
+ return Objects.equal(entries, that.entries) && Objects.equal(ownerRepo, that.ownerRepo);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hashCode(entries, ownerRepo);
+ }
 
  public static RepositoryMapping create(
  Map<String, RepositoryName> entries, RepositoryName ownerRepo) {
- return createInternal(
+ return new RepositoryMapping(
  Preconditions.checkNotNull(entries), Preconditions.checkNotNull(ownerRepo));
  }
 
  public static RepositoryMapping createAllowingFallback(Map<String, RepositoryName> entries) {
- return createInternal(Preconditions.checkNotNull(entries), null);
- }
-
- private static RepositoryMapping createInternal(
- Map<String, RepositoryName> entries, RepositoryName ownerRepo) {
- return new AutoValue_RepositoryMapping(ImmutableMap.copyOf(entries), ownerRepo);
+ return new RepositoryMapping(Preconditions.checkNotNull(entries), null);
  }
 
  /**
@@ -70,7 +94,7 @@ private static RepositoryMapping createInternal(
  public RepositoryMapping withAdditionalMappings(Map<String, RepositoryName> additionalMappings) {
  HashMap<String, RepositoryName> allMappings = new HashMap<>(additionalMappings);
  allMappings.putAll(entries());
- return createInternal(allMappings, ownerRepo());
+ return new RepositoryMapping(allMappings, ownerRepo());
  }
 
  /**
@@ -134,6 +158,26 @@ public RepositoryMapping composeWith(RepositoryMapping other) {
  RepositoryName mappedName = other.get(entry.getValue().getName());
  entries.put(entry.getKey(), mappedName.isVisible() ? mappedName : entry.getValue());
  }
- return createInternal(entries, null);
+ return new RepositoryMapping(entries, null);
+ }
+
+ /**
+ * Returns a new {@link RepositoryMapping} instance with identical contents, except that the
+ * inverse mapping is cached, causing {@link #getInverse} to be much more efficient. This is
+ * particularly important for the main repo mapping, as it's often used to generate display-form
+ * labels ({@link Label#getDisplayForm}).
+ */
+ public RepositoryMapping withCachedInverseMap() {
+ var inverse = Maps.<RepositoryName, String>newHashMapWithExpectedSize(entries.size());
+ for (Map.Entry<String, RepositoryName> entry : entries.entrySet()) {
+ inverse.putIfAbsent(entry.getValue(), entry.getKey());
+ }
+ var inverseCopy = ImmutableMap.copyOf(inverse);
+ return new RepositoryMapping(entries, ownerRepo) {
+ @Override
+ public Optional<String> getInverse(RepositoryName postMappingName) {
+ return Optional.ofNullable(inverseCopy.get(postMappingName));
+ }
+ };
  }
 }

src/main/java/com/google/devtools/build/lib/skyframe/RepositoryMappingFunction.java

@@ -128,14 +128,17 @@ public SkyValue compute(SkyKey skyKey, Environment env)
  .withAdditionalMappings(
  ImmutableMap.of(
  externalPackageValue.getPackage().getWorkspaceName(), RepositoryName.MAIN))
- .withAdditionalMappings(additionalMappings);
+ .withAdditionalMappings(additionalMappings)
+ .withCachedInverseMap();
  }
 
  // Try and see if this is a repo generated from a Bazel module.
  Optional<RepositoryMappingValue> mappingValue =
  computeForBazelModuleRepo(repositoryName, bazelDepGraphValue);
  if (mappingValue.isPresent()) {
- return mappingValue.get();
+ return repositoryName.isMain()
+ ? mappingValue.get().withCachedInverseMap()
+ : mappingValue.get();
  }
 
  // Now try and see if this is a repo generated from a module extension.

src/main/java/com/google/devtools/build/lib/skyframe/RepositoryMappingValue.java

@@ -107,6 +107,17 @@ public final RepositoryMappingValue withAdditionalMappings(Map<String, Repositor
  getAssociatedModuleVersion());
  }
 
+ /**
+ * Replaces the inner {@link #getRepositoryMapping() repository mapping} with one returned by
+ * calling its {@link RepositoryMapping#withCachedInverseMap} method.
+ */
+ public final RepositoryMappingValue withCachedInverseMap() {
+ return new AutoValue_RepositoryMappingValue(
+ getRepositoryMapping().withCachedInverseMap(),
+ getAssociatedModuleName(),
+ getAssociatedModuleVersion());
+ }
+
  /** Returns the {@link Key} for {@link RepositoryMappingValue}s. */
  public static Key key(RepositoryName repositoryName) {
  return RepositoryMappingValue.Key.create(