Allow PackageRootResolver to take in directories for findPackageRoots().

--
MOS_MIGRATED_REVID=117512506

4 files changed, 61 insertions, 7 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/actions/PackageRootResolver.java b/src/main/java/com/google/devtools/build/lib/actions/PackageRootResolver.java
index 6f058a5728..0db0075ddc 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/PackageRootResolver.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/PackageRootResolver.java
@@ -39,4 +39,21 @@ public interface PackageRootResolver {
   @Nullable
   Map<PathFragment, Root> findPackageRootsForFiles(Iterable<PathFragment> execPaths)
       throws PackageRootResolutionException;
+
+  /**
+   * Returns mapping from execPath to Root. Root will be null if the path has no containing
+   * package. Unlike {@link #findPackageRootsForFiles(Iterable)}, this function allows directories
+   * in the list of exec paths.
+   *
+   * @param execPaths the paths to find {@link Root}s for. The search for a containing package will
+   *    start with the path's parent directory, since the path is assumed to be a file.
+   * @return mappings from {@code execPath} to {@link Root}, or null if for some reason we
+   *    cannot determine the result at this time (such as when used within a SkyFunction)
+   * @throws PackageRootResolutionException if unable to determine package roots or lack thereof,
+   *    typically caused by exceptions encountered while attempting to locate BUILD files
+   */
+  // TODO(bazel-team): Remove this once we don't need to find package roots for directories.
+  @Nullable
+  Map<PathFragment, Root> findPackageRoots(Iterable<PathFragment> execPaths)
+      throws PackageRootResolutionException;
 }
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/SkyframePackageRootResolver.java b/src/main/java/com/google/devtools/build/lib/analysis/SkyframePackageRootResolver.java
index 8fed1a7993..2b4a61e6b7 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/SkyframePackageRootResolver.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/SkyframePackageRootResolver.java
@@ -23,6 +23,8 @@ import com.google.devtools.build.lib.vfs.PathFragment;
 
 import java.util.Map;
 
+import javax.annotation.Nullable;
+
 /**
  * Skyframe implementation of {@link PackageRootResolver}.
  *
@@ -40,6 +42,13 @@ public final class SkyframePackageRootResolver implements PackageRootResolver {
   @Override
   public Map<PathFragment, Root> findPackageRootsForFiles(Iterable<PathFragment> execPaths)
       throws PackageRootResolutionException {
+    return executor.getArtifactRootsForFiles(eventHandler, execPaths);
+  }
+  
+  @Override
+  @Nullable
+  public Map<PathFragment, Root> findPackageRoots(Iterable<PathFragment> execPaths)
+      throws PackageRootResolutionException {
     return executor.getArtifactRoots(eventHandler, execPaths);
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java
index 8b37d9f933..ba4c70aeb3 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java
@@ -312,6 +312,15 @@ public class ActionExecutionFunction implements SkyFunction, CompletionReceiver
       // If some values are missing, return null.
       return env.valuesMissing() ? null : result;
     }
+
+    @Override
+    @Nullable
+    public Map<PathFragment, Root> findPackageRoots(Iterable<PathFragment> execPaths)
+        throws PackageRootResolutionException {
+      // call sites for this implementation of PackageRootResolver shouldn't be passing in
+      // directories.
+      return findPackageRootsForFiles(execPaths);
+    }
   }
 
   private ActionExecutionValue checkCacheAndExecuteIfNeeded(
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
index b948490519..d37c582440 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
@@ -702,15 +702,34 @@ public abstract class SkyframeExecutor implements WalkableGraphFactory {
   }
 
   // TODO(bazel-team): Make this take a PackageIdentifier.
+  public Map<PathFragment, Root> getArtifactRootsForFiles(final EventHandler eventHandler,
+      Iterable<PathFragment> execPaths) throws PackageRootResolutionException {
+    return getArtifactRoots(eventHandler, execPaths, true);
+  }
+
   public Map<PathFragment, Root> getArtifactRoots(final EventHandler eventHandler,
       Iterable<PathFragment> execPaths) throws PackageRootResolutionException {
+    return getArtifactRoots(eventHandler, execPaths, false);
+  }
+
+  private Map<PathFragment, Root> getArtifactRoots(final EventHandler eventHandler,
+      Iterable<PathFragment> execPaths, boolean forFiles) throws PackageRootResolutionException {
+
     final List<SkyKey> packageKeys = new ArrayList<>();
-    for (PathFragment execPath : execPaths) {
-      PathFragment parent = Preconditions.checkNotNull(
-          execPath.getParentDirectory(), "Must pass in files, not root directory");
-      Preconditions.checkArgument(!parent.isAbsolute(), execPath);
-      packageKeys.add(ContainingPackageLookupValue.key(
-          PackageIdentifier.createInMainRepo(parent)));
+    if (forFiles) {
+      for (PathFragment execPath : execPaths) {
+        PathFragment parent = Preconditions.checkNotNull(
+            execPath.getParentDirectory(), "Must pass in files, not root directory");
+        Preconditions.checkArgument(!parent.isAbsolute(), execPath);
+        packageKeys.add(ContainingPackageLookupValue.key(
+            PackageIdentifier.createInMainRepo(parent)));
+      }
+    } else {
+      for (PathFragment execPath : execPaths) {
+        Preconditions.checkArgument(!execPath.isAbsolute(), execPath);
+        packageKeys.add(ContainingPackageLookupValue.key(
+            PackageIdentifier.createInMainRepo(execPath)));
+      }
     }
 
     EvaluationResult<ContainingPackageLookupValue> result;
@@ -736,7 +755,7 @@ public abstract class SkyframeExecutor implements WalkableGraphFactory {
     Map<PathFragment, Root> roots = new HashMap<>();
     for (PathFragment execPath : execPaths) {
       ContainingPackageLookupValue value = result.get(ContainingPackageLookupValue.key(
-          PackageIdentifier.createInMainRepo(execPath.getParentDirectory())));
+          PackageIdentifier.createInMainRepo(forFiles ? execPath.getParentDirectory() : execPath)));
       if (value.hasContainingPackage()) {
         roots.put(execPath, Root.asSourceRoot(value.getContainingPackageRoot()));
       } else {