From 851c3304b88a3ea5352e387edb7adb7943f84fc8 Mon Sep 17 00:00:00 2001
From: Googler <messa@google.com>
Date: Tue, 24 Jan 2023 12:26:23 -0800
Subject: [PATCH] `same_pkg_direct_rdeps` only loads packages of its inputs

For graphless query mode, since `same_pkg_direct_rdeps` only tracks the same-package dependencies, there is no need to call `TargetProvider$getTarget()` for dependencies from different packages because this will [load unloaded packages](https://github.com/bazelbuild/bazel/blob/master/src/main/java/com/google/devtools/build/lib/pkgcache/TargetProvider.java#L30:L31) and their targets [will be ignored later](https://github.com/bazelbuild/bazel/blob/master/src/main/java/com/google/devtools/build/lib/query2/query/PathLabelVisitor.java#L257:L260).

PiperOrigin-RevId: 504345382
Change-Id: Idb76d3f2d1a49b2d5028c37a84a6bfbee9c22689
---
 .../lib/query2/query/PathLabelVisitor.java    |  6 ++++
 .../shell/integration/bazel_query_test.sh     | 34 +++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/src/main/java/com/google/devtools/build/lib/query2/query/PathLabelVisitor.java b/src/main/java/com/google/devtools/build/lib/query2/query/PathLabelVisitor.java
index df13235aeca6a3..5bbf8fdf74b84f 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/query/PathLabelVisitor.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/query/PathLabelVisitor.java
@@ -237,6 +237,12 @@ private void visitTargets(Iterable<Target> targets) throws InterruptedException
 
     private void enqueue(Target from, Attribute attribute, Label label)
         throws InterruptedException, NoSuchThingException {
+      if (mode == VisitorMode.SAME_PKG_DIRECT_RDEPS) {
+        // Only track same-package dependencies to avoid loading unneeded packages.
+        if (!label.getPackageIdentifier().equals(from.getLabel().getPackageIdentifier())) {
+          return;
+        }
+      }
       Target target = targetProvider.getTarget(eventHandler, label);
       enqueue(from, attribute, target);
     }
diff --git a/src/test/shell/integration/bazel_query_test.sh b/src/test/shell/integration/bazel_query_test.sh
index 5e2285691148cb..d97b5ceff4e3fd 100755
--- a/src/test/shell/integration/bazel_query_test.sh
+++ b/src/test/shell/integration/bazel_query_test.sh
@@ -1064,4 +1064,38 @@ EOF
   expect_log "//${package}:hint"
 }
 
+function test_same_pkg_direct_rdeps_loads_only_inputs_packages() {
+  mkdir -p "pkg1"
+  mkdir -p "pkg2"
+  mkdir -p "pkg3"
+
+  cat > "pkg1/BUILD" <<EOF
+sh_library(name = "t1", deps = [":t2", "//pkg2:t3"])
+sh_library(name = "t2")
+EOF
+
+  cat > "pkg2/BUILD" <<EOF
+sh_library(name = "t3")
+EOF
+
+  cat > "pkg3/BUILD" <<EOF
+sh_library(name = "t4", deps = [":t5"])
+sh_library(name = "t5")
+EOF
+
+  bazel query --experimental_ui_debug_all_events \
+     "same_pkg_direct_rdeps(//pkg1:t2+//pkg3:t5)"  >& $TEST_log \
+    || fail "Expected success"
+
+  expect_log "Loading package: pkg1"
+  expect_log "Loading package: pkg3"
+  # For graphless query mode, pkg2 should not be loaded because
+  # same_pkg_direct_rdeps only cares about the targets in the same package
+  # as its inputs.
+  expect_not_log "Loading package: pkg2"
+  # the result of "same_pkg_direct_rdeps(//pkg1:t2+//pkg3:t5)"
+  expect_log "//pkg1:t1"
+  expect_log "//pkg3:t4"
+}
+
 run_suite "${PRODUCT_NAME} query tests"