Remap labels that include a repository name that appear in $(location x).

RELNOTES: None.
PiperOrigin-RevId: 201588988

3 files changed, 36 insertions, 16 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/analysis/LocationExpander.java b/src/main/java/com/google/devtools/build/lib/analysis/LocationExpander.java
index 40c77720ee..0a736a3f41 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/LocationExpander.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/LocationExpander.java
@@ -28,6 +28,7 @@ import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.analysis.configuredtargets.RuleConfiguredTarget.Mode;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.LabelSyntaxException;
+import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.packages.BuildType;
 import com.google.devtools.build.lib.packages.OutputFile;
 import com.google.devtools.build.lib.packages.RuleErrorConsumer;
@@ -68,23 +69,25 @@ public final class LocationExpander {
 
   private final RuleErrorConsumer ruleErrorConsumer;
   private final ImmutableMap<String, LocationFunction> functions;
+  private final ImmutableMap<RepositoryName, RepositoryName> repositoryMapping;
 
   @VisibleForTesting
   LocationExpander(
       RuleErrorConsumer ruleErrorConsumer,
-      Map<String, LocationFunction> functions) {
+      Map<String, LocationFunction> functions,
+      ImmutableMap<RepositoryName, RepositoryName> repositoryMapping) {
     this.ruleErrorConsumer = ruleErrorConsumer;
     this.functions = ImmutableMap.copyOf(functions);
+    this.repositoryMapping = repositoryMapping;
   }
 
   private LocationExpander(
       RuleErrorConsumer ruleErrorConsumer,
       Label root,
       Supplier<Map<Label, Collection<Artifact>>> locationMap,
-      boolean execPaths) {
-    this(
-        ruleErrorConsumer,
-        allLocationFunctions(root, locationMap, execPaths));
+      boolean execPaths,
+      ImmutableMap<RepositoryName, RepositoryName> repositoryMapping) {
+    this(ruleErrorConsumer, allLocationFunctions(root, locationMap, execPaths), repositoryMapping);
   }
 
   /**
@@ -108,7 +111,8 @@ public final class LocationExpander {
         // Use a memoizing supplier to avoid eagerly building the location map.
         Suppliers.memoize(
             () -> LocationExpander.buildLocationMap(ruleContext, labelMap, allowData)),
-        execPaths);
+        execPaths,
+        ruleContext.getRule().getPackage().getRepositoryMapping());
   }
 
   /**
@@ -209,7 +213,7 @@ public final class LocationExpander {
       // (2) Call appropriate function to obtain string replacement.
       String functionValue = value.substring(nextWhitespace + 1, end).trim();
       try {
-        String replacement = functions.get(fname).apply(functionValue);
+        String replacement = functions.get(fname).apply(functionValue, repositoryMapping);
         result.append(replacement);
       } catch (IllegalStateException ise) {
         reporter.report(ise.getMessage());
@@ -243,15 +247,19 @@ public final class LocationExpander {
     }
 
     /**
-     * Looks up the label-like string in the locationMap and returns the resolved path string.
+     * Looks up the label-like string in the locationMap and returns the resolved path string. If
+     * the label-like string begins with a repository name, the repository name may be remapped
+     * using the {@code repositoryMapping}.
      *
      * @param arg The label-like string to be expanded, e.g. ":foo" or "//foo:bar"
+     * @param repositoryMapping map of {@code RepositoryName}s defined in the main workspace
      * @return The expanded value
      */
-    public String apply(String arg) {
+    public String apply(
+        String arg, ImmutableMap<RepositoryName, RepositoryName> repositoryMapping) {
       Label label;
       try {
-        label = root.getRelative(arg);
+        label = root.getRelativeWithRemapping(arg, repositoryMapping);
       } catch (LabelSyntaxException e) {
         throw new IllegalStateException(
             String.format(
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/LocationTemplateContext.java b/src/main/java/com/google/devtools/build/lib/analysis/LocationTemplateContext.java
index faa5dff18e..d1a889a3e8 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/LocationTemplateContext.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/LocationTemplateContext.java
@@ -23,6 +23,7 @@ import com.google.devtools.build.lib.analysis.LocationExpander.LocationFunction;
 import com.google.devtools.build.lib.analysis.stringtemplate.ExpansionException;
 import com.google.devtools.build.lib.analysis.stringtemplate.TemplateContext;
 import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.cmdline.RepositoryName;
 import java.util.Collection;
 import java.util.Map;
 import javax.annotation.Nullable;
@@ -47,14 +48,17 @@ import javax.annotation.Nullable;
 final class LocationTemplateContext implements TemplateContext {
   private final TemplateContext delegate;
   private final ImmutableMap<String, LocationFunction> functions;
+  private final ImmutableMap<RepositoryName, RepositoryName> repositoryMapping;
 
   private LocationTemplateContext(
       TemplateContext delegate,
       Label root,
       Supplier<Map<Label, Collection<Artifact>>> locationMap,
-      boolean execPaths) {
+      boolean execPaths,
+      ImmutableMap<RepositoryName, RepositoryName> repositoryMapping) {
     this.delegate = delegate;
     this.functions = LocationExpander.allLocationFunctions(root, locationMap, execPaths);
+    this.repositoryMapping = repositoryMapping;
   }
 
   public LocationTemplateContext(
@@ -69,7 +73,8 @@ final class LocationTemplateContext implements TemplateContext {
         // Use a memoizing supplier to avoid eagerly building the location map.
         Suppliers.memoize(
             () -> LocationExpander.buildLocationMap(ruleContext, labelMap, allowData)),
-        execPaths);
+        execPaths,
+        ruleContext.getRule().getPackage().getRepositoryMapping());
   }
 
   @Override
@@ -82,7 +87,7 @@ final class LocationTemplateContext implements TemplateContext {
     try {
       LocationFunction f = functions.get(name);
       if (f != null) {
-        return f.apply(param);
+        return f.apply(param, repositoryMapping);
       }
     } catch (IllegalStateException e) {
       throw new ExpansionException(e.getMessage(), e);
diff --git a/src/main/java/com/google/devtools/build/lib/packages/Package.java b/src/main/java/com/google/devtools/build/lib/packages/Package.java
index 3d84851a22..cb7c427ae7 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/Package.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/Package.java
@@ -255,16 +255,23 @@ public class Package {
   }
 
   /**
-   * Returns the workspace mappings for the repository with the given absolute name.
+   * Returns the repository mapping for the requested external repository.
    *
    * @throws LabelSyntaxException if repository is not a valid {@link RepositoryName}
+   * @throws UnsupportedOperationException if called from any package other than the //external
+   *     package
    */
-  public ImmutableMap<RepositoryName, RepositoryName> getRepositoryMapping(
-      String repository) throws LabelSyntaxException {
+  public ImmutableMap<RepositoryName, RepositoryName> getRepositoryMapping(String repository)
+      throws LabelSyntaxException, UnsupportedOperationException {
     RepositoryName repositoryName = RepositoryName.create(repository);
     return getRepositoryMapping(repositoryName);
   }
 
+  /** Get the repository mapping for this package. */
+  public ImmutableMap<RepositoryName, RepositoryName> getRepositoryMapping() {
+    return repositoryMapping;
+  }
+
   /**
    * Gets the global name for a repository within an external repository.
    *