From a6c3f2327df424e56674a6bd758566d2757afdc7 Mon Sep 17 00:00:00 2001
From: Zhongpeng Lin <zplin@uber.com>
Date: Tue, 25 Jan 2022 08:46:11 -0800
Subject: [PATCH] Adding Starlark dependencies to the package //external

This is an alternative approach to fix #14280. It adds transitive closure of Starlark dependencies to `//external` package when loading `WORKSPACE` file, so it can be processed in the same way as `BUILD` files during query execution. Comparing to the approach taken in #14497, this approach is less intrusive, but not able to distinguish the extension files needed by `//external:foo` and `//external:bar`, meaning `buildfiles(//external:foo)` returns the same result as `buildfiles(//external:*)`. However, this behavior is consistent with other packages. For example, `buildfiles(//foo:bar)` has the same result as `buildfiles(//foo:*)`.

Closes #14630.

PiperOrigin-RevId: 424092916
---
 .../devtools/build/lib/packages/Package.java  |  2 +-
 .../build/lib/packages/PackageFactory.java    |  2 +-
 .../lib/skyframe/WorkspaceFileFunction.java   |  9 ++
 src/test/py/bazel/query_test.py               | 89 +++++++++++++++++++
 4 files changed, 100 insertions(+), 2 deletions(-)

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 b5409285759896..598a05b59121c3 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
@@ -1370,7 +1370,7 @@ FailureDetail getFailureDetail() {
       return null;
     }
 
-    Builder setStarlarkFileDependencies(ImmutableList<Label> starlarkFileDependencies) {
+    public Builder setStarlarkFileDependencies(ImmutableList<Label> starlarkFileDependencies) {
       this.starlarkFileDependencies = starlarkFileDependencies;
       return this;
     }
diff --git a/src/main/java/com/google/devtools/build/lib/packages/PackageFactory.java b/src/main/java/com/google/devtools/build/lib/packages/PackageFactory.java
index b0e956743c53c2..9cd3424641c9ee 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/PackageFactory.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/PackageFactory.java
@@ -706,7 +706,7 @@ private static ImmutableList<Label> transitiveClosureOfLabels(
     return ImmutableList.copyOf(set);
   }
 
-  private static void transitiveClosureOfLabelsRec(
+  public static void transitiveClosureOfLabelsRec(
       Set<Label> set, ImmutableMap<String, Module> loads) {
     for (Module m : loads.values()) {
       BazelModuleContext ctx = BazelModuleContext.of(m);
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/WorkspaceFileFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/WorkspaceFileFunction.java
index 9021c952a319e8..7da227dfba4c21 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/WorkspaceFileFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/WorkspaceFileFunction.java
@@ -24,6 +24,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSortedSet;
+import com.google.common.collect.Sets;
 import com.google.devtools.build.lib.actions.FileValue;
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
 import com.google.devtools.build.lib.cmdline.Label;
@@ -58,6 +59,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import net.starlark.java.eval.Module;
 import net.starlark.java.eval.Mutability;
 import net.starlark.java.eval.Starlark;
@@ -325,14 +327,21 @@ public SkyValue compute(SkyKey skyKey, Environment env)
               directories.getWorkspace(),
               directories.getLocalJavabase(),
               starlarkSemantics);
+      Set<Label> starlarkFileDependencies;
       if (prevValue != null) {
+        starlarkFileDependencies =
+            Sets.newLinkedHashSet(prevValue.getPackage().getStarlarkFileDependencies());
         try {
           parser.setParent(
               prevValue.getPackage(), prevValue.getLoadedModules(), prevValue.getBindings());
         } catch (NameConflictException e) {
           throw new WorkspaceFileFunctionException(e, Transience.PERSISTENT);
         }
+      } else {
+        starlarkFileDependencies = Sets.newLinkedHashSet();
       }
+      PackageFactory.transitiveClosureOfLabelsRec(starlarkFileDependencies, loadedModules);
+      builder.setStarlarkFileDependencies(ImmutableList.copyOf(starlarkFileDependencies));
       // Execute the partial files that comprise this chunk.
       for (StarlarkFile partialFile : chunk) {
         parser.execute(partialFile, loadedModules, key);
diff --git a/src/test/py/bazel/query_test.py b/src/test/py/bazel/query_test.py
index d5b37fa8fa5781..cec898e57ad577 100644
--- a/src/test/py/bazel/query_test.py
+++ b/src/test/py/bazel/query_test.py
@@ -38,6 +38,87 @@ def testSimpleQuery(self):
     self._AssertQueryOutput('deps(//foo:top-rule, 1)', '//foo:top-rule',
                             '//foo:dep-rule')
 
+  def testBuildFilesForExternalRepos_Simple(self):
+    self.ScratchFile('WORKSPACE', [
+        'load("//:deps.bzl", "repos")',
+        'repos()',
+    ])
+    self.ScratchFile('BUILD.bazel')
+    self.ScratchFile('deps.bzl', [
+        'def repos():',
+        '    native.new_local_repository(',
+        '        name = "io_bazel_rules_go",',
+        '        path = ".",',
+        """        build_file_content = "exports_files(glob(['*.go']))",""",
+        '    )',
+    ])
+    self._AssertQueryOutputContains('buildfiles(//external:io_bazel_rules_go)',
+                                    '//external:WORKSPACE', '//:deps.bzl',
+                                    '//:BUILD.bazel')
+
+  def testBuildFilesForExternalRepos_IndirectLoads(self):
+    self.ScratchFile('WORKSPACE', [
+        'load("//:deps.bzl", "repos")',
+        'repos()',
+    ])
+    self.ScratchFile('BUILD.bazel')
+    self.ScratchFile('deps.bzl', [
+        'load("//:private_deps.bzl", "other_repos")',
+        'def repos():',
+        '    native.new_local_repository(',
+        '        name = "io_bazel_rules_go",',
+        '        path = ".",',
+        """        build_file_content = "exports_files(glob(['*.go']))",""",
+        '    )',
+        '    other_repos()',
+        '',
+    ])
+    self.ScratchFile('private_deps.bzl', [
+        'def other_repos():',
+        '    native.new_local_repository(',
+        '        name = "io_bazel_rules_python",',
+        '        path = ".",',
+        """        build_file_content = "exports_files(glob(['*.py']))",""",
+        '    )',
+    ])
+
+    self._AssertQueryOutputContains(
+        'buildfiles(//external:io_bazel_rules_python)', '//external:WORKSPACE',
+        '//:deps.bzl', '//:private_deps.bzl', '//:BUILD.bazel')
+
+  def testBuildFilesForExternalRepos_NoDuplicates(self):
+    self.ScratchFile('WORKSPACE', [
+        'load("//:deps.bzl", "repos")',
+        'repos()',
+    ])
+    self.ScratchFile('BUILD.bazel')
+    self.ScratchFile('deps.bzl', [
+        'def repos():',
+        '    native.new_local_repository(',
+        '        name = "io_bazel_rules_go",',
+        '        path = ".",',
+        """        build_file_content = "exports_files(glob(['*.go']))",""",
+        '    )',
+        '    other_repos()',
+        '',
+        'def other_repos():',
+        '    native.new_local_repository(',
+        '        name = "io_bazel_rules_python",',
+        '        path = ".",',
+        """        build_file_content = "exports_files(glob(['*.py']))",""",
+        '    )',
+    ])
+
+    exit_code, stdout, stderr = self.RunBazel(
+        ['query', 'buildfiles(//external:io_bazel_rules_python)'])
+    self.AssertExitCode(exit_code, 0, stderr)
+    result = set()
+    for item in stdout:
+      if not item:
+        continue
+      self.assertNotIn(item, result)
+      result.add(item)
+
   def _AssertQueryOutput(self, query_expr, *expected_results):
     exit_code, stdout, stderr = self.RunBazel(['query', query_expr])
     self.AssertExitCode(exit_code, 0, stderr)
@@ -46,6 +127,14 @@ def _AssertQueryOutput(self, query_expr, *expected_results):
     self.assertEqual(len(stdout), len(expected_results))
     self.assertListEqual(stdout, sorted(expected_results))
 
+  def _AssertQueryOutputContains(self, query_expr, *expected_content):
+    exit_code, stdout, stderr = self.RunBazel(['query', query_expr])
+    self.AssertExitCode(exit_code, 0, stderr)
+
+    stdout = {x for x in stdout if x}
+    for item in expected_content:
+      self.assertIn(item, stdout)
+
 
 if __name__ == '__main__':
   unittest.main()