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

RELNOTES: None.
PiperOrigin-RevId: 201588988

6 files changed, 131 insertions, 30 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.
    *
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/LocationExpanderTest.java b/src/test/java/com/google/devtools/build/lib/analysis/LocationExpanderTest.java
index c334cce1c7..92f5616f90 100644
--- a/src/test/java/com/google/devtools/build/lib/analysis/LocationExpanderTest.java
+++ b/src/test/java/com/google/devtools/build/lib/analysis/LocationExpanderTest.java
@@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat;
 
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.analysis.LocationExpander.LocationFunction;
+import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.packages.AbstractRuleErrorConsumer;
 import com.google.devtools.build.lib.packages.RuleErrorConsumer;
 import java.util.ArrayList;
@@ -60,7 +61,6 @@ public class LocationExpanderTest {
   }
 
   private LocationExpander makeExpander(RuleErrorConsumer ruleErrorConsumer) throws Exception {
-
     LocationFunction f1 = new LocationFunctionBuilder("//a", false)
         .setExecPaths(false)
         .add("//a", "/exec/src/a")
@@ -75,7 +75,8 @@ public class LocationExpanderTest {
         ruleErrorConsumer,
         ImmutableMap.<String, LocationFunction>of(
             "location", f1,
-            "locations", f2));
+            "locations", f2),
+        ImmutableMap.of());
   }
 
   private String expand(String input) throws Exception {
@@ -123,4 +124,24 @@ public class LocationExpanderTest {
     assertThat(value).isEqualTo("foo $(location a) $(location a");
     assertThat(capture.warnsOrErrors).containsExactly("ERROR: unterminated $(location) expression");
   }
+
+  @Test
+  public void expansionWithRepositoryMapping() throws Exception {
+    LocationFunction f1 = new LocationFunctionBuilder("//a", false)
+        .setExecPaths(false)
+        .add("@bar//a", "/exec/src/a")
+        .build();
+
+    ImmutableMap<RepositoryName, RepositoryName> repositoryMapping = ImmutableMap.of(
+        RepositoryName.create("@foo"),
+        RepositoryName.create("@bar"));
+
+    LocationExpander locationExpander = new LocationExpander(
+        new Capture(),
+        ImmutableMap.<String, LocationFunction>of("location", f1),
+        repositoryMapping);
+
+    String value = locationExpander.expand("$(location @foo//a)");
+    assertThat(value).isEqualTo("src/a");
+  }
 }
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/LocationFunctionTest.java b/src/test/java/com/google/devtools/build/lib/analysis/LocationFunctionTest.java
index c6ce084d7a..258095fa95 100644
--- a/src/test/java/com/google/devtools/build/lib/analysis/LocationFunctionTest.java
+++ b/src/test/java/com/google/devtools/build/lib/analysis/LocationFunctionTest.java
@@ -18,10 +18,12 @@ import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.fail;
 
 import com.google.common.base.Suppliers;
+import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.ArtifactRoot;
 import com.google.devtools.build.lib.analysis.LocationExpander.LocationFunction;
 import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.vfs.FileSystem;
 import com.google.devtools.build.lib.vfs.Root;
 import com.google.devtools.build.lib.vfs.inmemoryfs.InMemoryFileSystem;
@@ -42,23 +44,23 @@ public class LocationFunctionTest {
   public void absoluteAndRelativeLabels() throws Exception {
     LocationFunction func =
         new LocationFunctionBuilder("//foo", false).add("//foo", "/exec/src/bar").build();
-    assertThat(func.apply("//foo")).isEqualTo("src/bar");
-    assertThat(func.apply(":foo")).isEqualTo("src/bar");
-    assertThat(func.apply("foo")).isEqualTo("src/bar");
+    assertThat(func.apply("//foo", ImmutableMap.of())).isEqualTo("src/bar");
+    assertThat(func.apply(":foo", ImmutableMap.of())).isEqualTo("src/bar");
+    assertThat(func.apply("foo", ImmutableMap.of())).isEqualTo("src/bar");
   }
 
   @Test
   public void pathUnderExecRootUsesDotSlash() throws Exception {
     LocationFunction func =
         new LocationFunctionBuilder("//foo", false).add("//foo", "/exec/bar").build();
-    assertThat(func.apply("//foo")).isEqualTo("./bar");
+    assertThat(func.apply("//foo", ImmutableMap.of())).isEqualTo("./bar");
   }
 
   @Test
   public void noSuchLabel() throws Exception {
     LocationFunction func = new LocationFunctionBuilder("//foo", false).build();
     try {
-      func.apply("//bar");
+      func.apply("//bar", ImmutableMap.of());
       fail();
     } catch (IllegalStateException expected) {
       assertThat(expected).hasMessageThat()
@@ -72,7 +74,7 @@ public class LocationFunctionTest {
   public void emptyList() throws Exception {
     LocationFunction func = new LocationFunctionBuilder("//foo", false).add("//foo").build();
     try {
-      func.apply("//foo");
+      func.apply("//foo", ImmutableMap.of());
       fail();
     } catch (IllegalStateException expected) {
       assertThat(expected).hasMessageThat()
@@ -85,7 +87,7 @@ public class LocationFunctionTest {
     LocationFunction func =
         new LocationFunctionBuilder("//foo", false).add("//foo", "/exec/1", "/exec/2").build();
     try {
-      func.apply("//foo");
+      func.apply("//foo", ImmutableMap.of());
       fail();
     } catch (IllegalStateException expected) {
       assertThat(expected).hasMessageThat()
@@ -100,7 +102,7 @@ public class LocationFunctionTest {
   public void noSuchLabelMultiple() throws Exception {
     LocationFunction func = new LocationFunctionBuilder("//foo", true).build();
     try {
-      func.apply("//bar");
+      func.apply("//bar", ImmutableMap.of());
       fail();
     } catch (IllegalStateException expected) {
       assertThat(expected).hasMessageThat()
@@ -114,7 +116,7 @@ public class LocationFunctionTest {
   public void fileWithSpace() throws Exception {
     LocationFunction func =
         new LocationFunctionBuilder("//foo", false).add("//foo", "/exec/file/with space").build();
-    assertThat(func.apply("//foo")).isEqualTo("'file/with space'");
+    assertThat(func.apply("//foo", ImmutableMap.of())).isEqualTo("'file/with space'");
   }
 
   @Test
@@ -122,7 +124,7 @@ public class LocationFunctionTest {
     LocationFunction func = new LocationFunctionBuilder("//foo", true)
         .add("//foo", "/exec/foo/bar", "/exec/out/foo/foobar")
         .build();
-    assertThat(func.apply("//foo")).isEqualTo("foo/bar foo/foobar");
+    assertThat(func.apply("//foo", ImmutableMap.of())).isEqualTo("foo/bar foo/foobar");
   }
 
   @Test
@@ -130,7 +132,8 @@ public class LocationFunctionTest {
     LocationFunction func = new LocationFunctionBuilder("//foo", true)
         .add("//foo", "/exec/file/with space", "/exec/file/with spaces ")
         .build();
-    assertThat(func.apply("//foo")).isEqualTo("'file/with space' 'file/with spaces '");
+    assertThat(func.apply("//foo", ImmutableMap.of()))
+        .isEqualTo("'file/with space' 'file/with spaces '");
   }
 
   @Test
@@ -139,7 +142,27 @@ public class LocationFunctionTest {
         .setExecPaths(true)
         .add("//foo", "/exec/bar", "/exec/out/foobar")
         .build();
-    assertThat(func.apply("//foo")).isEqualTo("./bar out/foobar");
+    assertThat(func.apply("//foo", ImmutableMap.of())).isEqualTo("./bar out/foobar");
+  }
+
+  @Test
+  public void locationFunctionWithMappingReplace() throws Exception {
+    RepositoryName a = RepositoryName.create("@a");
+    RepositoryName b = RepositoryName.create("@b");
+    ImmutableMap<RepositoryName, RepositoryName> repositoryMapping = ImmutableMap.of(a, b);
+    LocationFunction func =
+        new LocationFunctionBuilder("//foo", false).add("@b//foo", "/exec/src/bar").build();
+    assertThat(func.apply("@a//foo", repositoryMapping)).isEqualTo("src/bar");
+  }
+
+  @Test
+  public void locationFunctionWithMappingIgnoreRepo() throws Exception {
+    RepositoryName a = RepositoryName.create("@a");
+    RepositoryName b = RepositoryName.create("@b");
+    ImmutableMap<RepositoryName, RepositoryName> repositoryMapping = ImmutableMap.of(a, b);
+    LocationFunction func =
+        new LocationFunctionBuilder("//foo", false).add("@potato//foo", "/exec/src/bar").build();
+    assertThat(func.apply("@potato//foo", repositoryMapping)).isEqualTo("src/bar");
   }
 }
 
diff --git a/src/test/shell/bazel/workspace_test.sh b/src/test/shell/bazel/workspace_test.sh
index 1250d69ee1..22b7a4741b 100755
--- a/src/test/shell/bazel/workspace_test.sh
+++ b/src/test/shell/bazel/workspace_test.sh
@@ -453,6 +453,43 @@ EOF
       || fail "Expected srcs to contain '@b//:x.txt'"
 }
 
+function test_repository_reassignment_location() {
+  # Repository a refers to @x
+  mkdir -p a
+  touch a/WORKSPACE
+  cat > a/BUILD<<EOF
+genrule(name = "a",
+        srcs = ["@x//:x.txt"],
+        outs = ["result.txt"],
+        cmd = "echo \$(location @x//:x.txt) > \$(location result.txt); \
+            cat \$(location @x//:x.txt)>> \$(location result.txt);"
+)
+EOF
+
+  # Repository b is a substitute for x
+  mkdir -p b
+  touch b/WORKSPACE
+  cat >b/BUILD <<EOF
+exports_files(srcs = ["x.txt"])
+EOF
+  echo "Hello from @b//:x.txt" > b/x.txt
+
+  # Main repo assigns @x to @b within @a
+  mkdir -p main
+  cat > main/WORKSPACE <<EOF
+workspace(name = "main")
+
+local_repository(name = "a", path="../a", repo_mapping = {"@x" : "@b"})
+local_repository(name = "b", path="../b")
+EOF
+  touch main/BUILD
+
+  cd main
+  bazel build --experimental_enable_repo_mapping @a//:a || fail "Expected build to succeed"
+  grep "external/b/x.txt" bazel-genfiles/external/a/result.txt \
+      || fail "expected external/b/x.txt in $(cat bazel-genfiles/external/a/result.txt)"
+}
+
 function test_workspace_addition_change_aspect() {
   mkdir -p repo_one
   mkdir -p repo_two