Make load() work in remote repositories too.

Previously, load() always looked up .bzl files in the main repository. Ideally, it would just take a label and then it would work by default, but for the time being, this quick fix will do.

I had to put in an evil hack to make load() statements work in the prelude, because we currently have no way to distinguish load() statements from the prelude and from the BUILD file. Again, a proper label-based load() would solve this.

--
MOS_MIGRATED_REVID=101677502

8 files changed, 112 insertions, 38 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/packages/PackageIdentifier.java b/src/main/java/com/google/devtools/build/lib/packages/PackageIdentifier.java
index 5a37107a29..f45ecc486a 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/PackageIdentifier.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/PackageIdentifier.java
@@ -21,6 +21,7 @@ import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
 import com.google.common.collect.ComparisonChain;
 import com.google.devtools.build.lib.cmdline.LabelValidator;
+import com.google.devtools.build.lib.syntax.Label;
 import com.google.devtools.build.lib.syntax.Label.SyntaxException;
 import com.google.devtools.build.lib.util.StringCanonicalizer;
 import com.google.devtools.build.lib.util.StringUtilities;
@@ -175,6 +176,15 @@ public final class PackageIdentifier implements Comparable<PackageIdentifier>, S
   }
 
   public static final String DEFAULT_REPOSITORY = "";
+  public static final RepositoryName DEFAULT_REPOSITORY_NAME;
+
+  static {
+    try {
+      DEFAULT_REPOSITORY_NAME = RepositoryName.create(DEFAULT_REPOSITORY);
+    } catch (Label.SyntaxException e) {
+      throw new IllegalStateException(e);
+    }
+  }
 
   /**
    * Helper for serializing PackageIdentifiers.
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ASTFileLookupFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/ASTFileLookupFunction.java
index fbf3870f09..6d51828a93 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ASTFileLookupFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ASTFileLookupFunction.java
@@ -14,7 +14,9 @@
 
 package com.google.devtools.build.lib.skyframe;
 
+import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.packages.CachingPackageLocator;
+import com.google.devtools.build.lib.packages.PackageIdentifier;
 import com.google.devtools.build.lib.packages.RuleClassProvider;
 import com.google.devtools.build.lib.pkgcache.PathPackageLocator;
 import com.google.devtools.build.lib.syntax.BuildFileAST;
@@ -28,6 +30,7 @@ import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 
 import java.io.IOException;
+import java.util.List;
 import java.util.concurrent.atomic.AtomicReference;
 
 import javax.annotation.Nullable;
@@ -84,7 +87,7 @@ public class ASTFileLookupFunction implements SkyFunction {
 
       @Override
       public RootedPath rootedPath() {
-        throw new IllegalStateException("unsucessful lookup");
+        throw new IllegalStateException("unsuccessful lookup");
       }
     }
   }
@@ -104,8 +107,9 @@ public class ASTFileLookupFunction implements SkyFunction {
   @Override
   public SkyValue compute(SkyKey skyKey, Environment env) throws SkyFunctionException,
       InterruptedException {
-    PathFragment astFilePathFragment = (PathFragment) skyKey.argument();
-    FileLookupResult lookupResult = getASTFile(env, astFilePathFragment);
+    PackageIdentifier key = (PackageIdentifier) skyKey.argument();
+    PathFragment astFilePathFragment = key.getPackageFragment();
+    FileLookupResult lookupResult = getASTFile(env, key);
     if (lookupResult == null) {
       return null;
     }
@@ -132,15 +136,27 @@ public class ASTFileLookupFunction implements SkyFunction {
     return ASTFileLookupValue.withFile(ast);
   }
 
-  private FileLookupResult getASTFile(Environment env, PathFragment astFilePathFragment)
+  private FileLookupResult getASTFile(Environment env, PackageIdentifier key)
       throws ASTLookupFunctionException {
-    for (Path packagePathEntry : pkgLocator.get().getPathEntries()) {
-      RootedPath rootedPath = RootedPath.toRootedPath(packagePathEntry, astFilePathFragment);
-      SkyKey fileSkyKey = FileValue.key(rootedPath);
-      FileValue fileValue = null;
+    List<Path> candidateRoots;
+    if (!key.getRepository().isDefault()) {
+      RepositoryValue repository =
+          (RepositoryValue) env.getValue(RepositoryValue.key(key.getRepository()));
+      if (repository == null) {
+        return null;
+      }
+
+      candidateRoots = ImmutableList.of(repository.getPath());
+    } else {
+      candidateRoots = pkgLocator.get().getPathEntries();
+    }
+
+    for (Path root : candidateRoots) {
+      RootedPath rootedPath = RootedPath.toRootedPath(root, key.getPackageFragment());
+      FileValue fileValue;
       try {
-        fileValue = (FileValue) env.getValueOrThrow(fileSkyKey, IOException.class,
-            FileSymlinkException.class, InconsistentFilesystemException.class);
+        fileValue = (FileValue) env.getValueOrThrow(FileValue.key(rootedPath),
+            IOException.class, FileSymlinkException.class, InconsistentFilesystemException.class);
       } catch (IOException | FileSymlinkException e) {
         throw new ASTLookupFunctionException(new ErrorReadingSkylarkExtensionException(
             e.getMessage()), Transience.PERSISTENT);
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ASTFileLookupValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/ASTFileLookupValue.java
index c9b7a88773..65a20f47a5 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ASTFileLookupValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ASTFileLookupValue.java
@@ -14,8 +14,8 @@
 
 package com.google.devtools.build.lib.skyframe;
 
+import com.google.devtools.build.lib.packages.PackageIdentifier;
 import com.google.devtools.build.lib.syntax.BuildFileAST;
-import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 
@@ -49,16 +49,8 @@ public class ASTFileLookupValue implements SkyValue {
     return ast;
   }
 
-  static void checkInputArgument(PathFragment astFilePathFragment) throws ASTLookupInputException {
-    if (astFilePathFragment.isAbsolute()) {
-      throw new ASTLookupInputException(String.format(
-          "Input file '%s' cannot be an absolute path.", astFilePathFragment));
-    }
-  }
-
-  static SkyKey key(PathFragment astFilePathFragment) throws ASTLookupInputException {
-    checkInputArgument(astFilePathFragment);
-    return new SkyKey(SkyFunctions.AST_FILE_LOOKUP, astFilePathFragment);
+  static SkyKey key(PackageIdentifier astFileIdentifier) throws ASTLookupInputException {
+    return new SkyKey(SkyFunctions.AST_FILE_LOOKUP, astFileIdentifier);
   }
 
   static final class ASTLookupInputException extends Exception {
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/PackageFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/PackageFunction.java
index af32b56959..398f973d75 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/PackageFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/PackageFunction.java
@@ -15,7 +15,6 @@ package com.google.devtools.build.lib.skyframe;
 
 import com.google.common.base.Preconditions;
 import com.google.common.cache.Cache;
-import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -38,6 +37,7 @@ import com.google.devtools.build.lib.packages.Package;
 import com.google.devtools.build.lib.packages.PackageFactory;
 import com.google.devtools.build.lib.packages.PackageFactory.Globber;
 import com.google.devtools.build.lib.packages.PackageIdentifier;
+import com.google.devtools.build.lib.packages.PackageIdentifier.RepositoryName;
 import com.google.devtools.build.lib.packages.Preprocessor;
 import com.google.devtools.build.lib.packages.RuleVisibility;
 import com.google.devtools.build.lib.packages.Target;
@@ -89,7 +89,7 @@ public class PackageFunction implements SkyFunction {
   private final AtomicInteger numPackagesLoaded;
   private final Profiler profiler = Profiler.instance();
 
-  private static final PathFragment PRELUDE_FILE_FRAGMENT =
+  static final PathFragment PRELUDE_FILE_FRAGMENT =
       new PathFragment(Constants.PRELUDE_FILE_DEPOT_RELATIVE_PATH);
 
   static final String DEFAULTS_PACKAGE_NAME = "tools/defaults";
@@ -422,7 +422,8 @@ public class PackageFunction implements SkyFunction {
     ASTFileLookupValue astLookupValue = null;
     SkyKey astLookupKey = null;
     try {
-      astLookupKey = ASTFileLookupValue.key(PRELUDE_FILE_FRAGMENT);
+      astLookupKey = ASTFileLookupValue.key(
+          PackageIdentifier.createInDefaultRepo(PRELUDE_FILE_FRAGMENT));
     } catch (ASTLookupInputException e) {
       // There's a static check ensuring that PRELUDE_FILE_FRAGMENT is relative.
       throw new IllegalStateException(e);
@@ -584,13 +585,22 @@ public class PackageFunction implements SkyFunction {
       BuildFileAST buildFileAST,
       Environment env)
       throws PackageFunctionException {
-    ImmutableCollection<PathFragment> imports = buildFileAST.getImports();
+    ImmutableMap<Location, PathFragment> imports = buildFileAST.getImports();
     Map<PathFragment, SkylarkEnvironment> importMap = new HashMap<>();
     ImmutableList.Builder<SkylarkFileDependency> fileDependencies = ImmutableList.builder();
     try {
-      for (PathFragment importFile : imports) {
-        SkyKey importsLookupKey =
-            SkylarkImportLookupValue.key(packageId.getRepository(), buildFileFragment, importFile);
+      for (Map.Entry<Location, PathFragment> entry : imports.entrySet()) {
+        PathFragment importFile = entry.getValue();
+        // HACK: The prelude sometimes contains load() statements, which need to be resolved
+        // relative to the prelude file. However, we don't have a good way to tell "this should come
+        // from the main repository" in a load() statement, and we don't have a good way to tell if
+        // a load() statement comes from the prelude, since we just prepend those statements before
+        // the actual BUILD file. So we use this evil .endsWith() statement to figure it out.
+        RepositoryName repository =
+            entry.getKey().getPath().endsWith(PRELUDE_FILE_FRAGMENT)
+                ? PackageIdentifier.DEFAULT_REPOSITORY_NAME : packageId.getRepository();
+        SkyKey importsLookupKey = SkylarkImportLookupValue.key(
+            repository, buildFileFragment, importFile);
         SkylarkImportLookupValue importLookupValue = (SkylarkImportLookupValue)
             env.getValueOrThrow(importsLookupKey, SkylarkImportFailedException.class,
                 InconsistentFilesystemException.class, ASTLookupInputException.class,
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkylarkImportLookupFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkylarkImportLookupFunction.java
index 724b6b2ec6..203f259815 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SkylarkImportLookupFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkylarkImportLookupFunction.java
@@ -15,6 +15,7 @@ package com.google.devtools.build.lib.skyframe;
 
 import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.events.Event;
+import com.google.devtools.build.lib.events.Location;
 import com.google.devtools.build.lib.events.StoredEventHandler;
 import com.google.devtools.build.lib.packages.BuildFileNotFoundException;
 import com.google.devtools.build.lib.packages.PackageFactory;
@@ -60,7 +61,7 @@ public class SkylarkImportLookupFunction implements SkyFunction {
     PathFragment file = arg.getPackageFragment();
     ASTFileLookupValue astLookupValue = null;
     try {
-      SkyKey astLookupKey = ASTFileLookupValue.key(file);
+      SkyKey astLookupKey = ASTFileLookupValue.key(arg);
       astLookupValue = (ASTFileLookupValue) env.getValueOrThrow(astLookupKey,
           ErrorReadingSkylarkExtensionException.class, InconsistentFilesystemException.class);
     } catch (ErrorReadingSkylarkExtensionException e) {
@@ -83,10 +84,18 @@ public class SkylarkImportLookupFunction implements SkyFunction {
     ImmutableList.Builder<SkylarkFileDependency> fileDependencies = ImmutableList.builder();
     BuildFileAST ast = astLookupValue.getAST();
     // TODO(bazel-team): Refactor this code and PackageFunction to reduce code duplications.
-    for (PathFragment importFile : ast.getImports()) {
+    for (Map.Entry<Location, PathFragment> entry : ast.getImports().entrySet()) {
       try {
-        SkyKey importsLookupKey =
-            SkylarkImportLookupValue.key(arg.getRepository(), file, importFile);
+        PathFragment importFile = entry.getValue();
+        // HACK: The prelude sometimes contains load() statements, which need to be resolved
+        // relative to the prelude file. However, we don't have a good way to tell "this should come
+        // from the main repository" in a load() statement, and we don't have a good way to tell if
+        // a load() statement comes from the prelude, since we just prepend those statements before
+        // the actual BUILD file. So we use this evil .endsWith() statement to figure it out.
+        RepositoryName repository =
+            entry.getKey().getPath().endsWith(PackageFunction.PRELUDE_FILE_FRAGMENT)
+                ? PackageIdentifier.DEFAULT_REPOSITORY_NAME : arg.getRepository();
+        SkyKey importsLookupKey = SkylarkImportLookupValue.key(repository, file, importFile);
         SkylarkImportLookupValue importsLookupValue;
         importsLookupValue = (SkylarkImportLookupValue) env.getValueOrThrow(
             importsLookupKey, ASTLookupInputException.class);
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkylarkImportLookupValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkylarkImportLookupValue.java
index 20579cce0d..9b307be7b8 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SkylarkImportLookupValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkylarkImportLookupValue.java
@@ -58,6 +58,14 @@ public class SkylarkImportLookupValue implements SkyValue {
     return dependency;
   }
 
+  private static void checkInputArgument(PathFragment astFilePathFragment)
+      throws ASTLookupInputException {
+    if (astFilePathFragment.isAbsolute()) {
+      throw new ASTLookupInputException(String.format(
+          "Input file '%s' cannot be an absolute path.", astFilePathFragment));
+    }
+  }
+
   @VisibleForTesting
   static SkyKey key(PackageIdentifier pkgIdentifier) throws ASTLookupInputException {
     return key(pkgIdentifier.getRepository(), pkgIdentifier.getPackageFragment());
@@ -79,7 +87,7 @@ public class SkylarkImportLookupValue implements SkyValue {
   private static SkyKey key(RepositoryName repo, PathFragment fileToImport)
       throws ASTLookupInputException {
     // Skylark import lookup keys need to be valid AST file lookup keys.
-    ASTFileLookupValue.checkInputArgument(fileToImport);
+    checkInputArgument(fileToImport);
     return new SkyKey(
         SkyFunctions.SKYLARK_IMPORTS_LOOKUP,
         new PackageIdentifier(repo, fileToImport));
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/BuildFileAST.java b/src/main/java/com/google/devtools/build/lib/syntax/BuildFileAST.java
index e2a333ab29..8dac01edcf 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/BuildFileAST.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/BuildFileAST.java
@@ -14,6 +14,7 @@
 package com.google.devtools.build.lib.syntax;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.hash.HashCode;
 import com.google.devtools.build.lib.events.Event;
@@ -37,7 +38,7 @@ public class BuildFileAST extends ASTNode {
 
   private final ImmutableList<Comment> comments;
 
-  private ImmutableSet<PathFragment> loads;
+  private ImmutableMap<Location, PathFragment> loads;
 
   private ImmutableSet<Label> includes;
 
@@ -105,12 +106,12 @@ public class BuildFileAST extends ASTNode {
   }
 
   /** Collects paths from all load statements */
-  private ImmutableSet<PathFragment> fetchLoads(List<Statement> stmts) {
-    ImmutableSet.Builder<PathFragment> loads = new ImmutableSet.Builder<>();
+  private ImmutableMap<Location, PathFragment> fetchLoads(List<Statement> stmts) {
+    ImmutableMap.Builder<Location, PathFragment> loads = ImmutableMap.builder();
     for (Statement stmt : stmts) {
       if (stmt instanceof LoadStatement) {
         LoadStatement imp = (LoadStatement) stmt;
-        loads.add(imp.getImportPath());
+        loads.put(imp.getLocation(), imp.getImportPath());
       }
     }
     return loads.build();
@@ -142,7 +143,7 @@ public class BuildFileAST extends ASTNode {
   /**
    * Returns a set of loads in this BUILD file.
    */
-  public synchronized ImmutableSet<PathFragment> getImports() {
+  public synchronized ImmutableMap<Location, PathFragment> getImports() {
     if (loads == null) {
       loads = fetchLoads(stmts);
     }
diff --git a/src/test/shell/bazel/local_repository_test.sh b/src/test/shell/bazel/local_repository_test.sh
index 4629ae337e..bb797724bd 100755
--- a/src/test/shell/bazel/local_repository_test.sh
+++ b/src/test/shell/bazel/local_repository_test.sh
@@ -763,4 +763,32 @@ EOF
   expect_log "Target '//:private' is not visible from target '@r//:private'"
 }
 
+function test_load_in_remote_repository() {
+  local r=$TEST_TMPDIR/r
+  rm -fr $r
+  mkdir -p $r
+  cat > $r/BUILD <<EOF
+package(default_visibility=["//visibility:public"])
+load("r", "r_filegroup")
+r_filegroup(name="rfg", srcs=["rfgf"])
+EOF
+
+  cat > $r/r.bzl <<EOF
+def r_filegroup(name, srcs):
+    native.filegroup(name=name, srcs=srcs)
+EOF
+
+  touch $r/rfgf
+
+  cat > WORKSPACE <<EOF
+local_repository(name="r", path="$r")
+EOF
+
+  cat > BUILD <<EOF
+filegroup(name="fg", srcs=["@r//:rfg"])
+EOF
+
+  bazel build //:fg || fail "failed to build target"
+}
+
 run_suite "local repository tests"