From ac53ab2bb9737c4308df7d21efe2547a23904a43 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?X=C3=B9d=C5=8Dng=20Y=C3=A1ng?= <wyverald@gmail.com>
Date: Fri, 16 Feb 2024 17:16:20 -0500
Subject: [PATCH] Fix rule definition environment for repo rules

See https://github.com/bazelbuild/bazel/pull/21131#discussion_r1471924084
---
 .../starlark/StarlarkRepositoryModule.java    | 18 +++++++---
 .../shell/bazel/starlark_repository_test.sh   | 33 +++++++++++++++++++
 2 files changed, 46 insertions(+), 5 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryModule.java b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryModule.java
index 9637172fcf0238..6b90a7c8543d3f 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryModule.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryModule.java
@@ -107,11 +107,19 @@ public StarlarkCallable repositoryRule(
       }
     }
     builder.setConfiguredTargetFunction(implementation);
-    // TODO(b/291752414): If we care about the digest of repository rules, we should be using the
-    // transitive bzl digest of the module of the outermost stack frame, not the innermost.
-    BazelModuleContext moduleContext = BazelModuleContext.ofInnermostBzlOrThrow(thread);
-    builder.setRuleDefinitionEnvironmentLabelAndDigest(
-        moduleContext.label(), moduleContext.bzlTransitiveDigest());
+    var bzlInitContext = (BzlInitThreadContext) thread.getThreadLocal(BazelStarlarkContext.class);
+    if (bzlInitContext != null) {
+      builder.setRuleDefinitionEnvironmentLabelAndDigest(
+          bzlInitContext.getBzlFile(), bzlInitContext.getTransitiveDigest());
+    } else {
+      // TODO: this branch is wrong, but cannot be removed until we deprecate WORKSPACE because
+      //   WORKSPACE can currently call unexported repo rules (so there's potentially no
+      //   BzlInitThreadContext. See
+      //   https://github.com/bazelbuild/bazel/pull/21131#discussion_r1471924084 for more details.
+      BazelModuleContext moduleContext = BazelModuleContext.ofInnermostBzlOrThrow(thread);
+      builder.setRuleDefinitionEnvironmentLabelAndDigest(
+          moduleContext.label(), moduleContext.bzlTransitiveDigest());
+    }
     Label.RepoMappingRecorder repoMappingRecorder =
         thread.getThreadLocal(Label.RepoMappingRecorder.class);
     if (repoMappingRecorder != null) {
diff --git a/src/test/shell/bazel/starlark_repository_test.sh b/src/test/shell/bazel/starlark_repository_test.sh
index 06e70124d12120..54da676930a190 100755
--- a/src/test/shell/bazel/starlark_repository_test.sh
+++ b/src/test/shell/bazel/starlark_repository_test.sh
@@ -784,6 +784,39 @@ function test_starlark_repository_bzl_invalidation_batch() {
   bzl_invalidation_test_template --batch
 }
 
+function test_starlark_repo_bzl_invalidation_wrong_digest() {
+  # regression test for https://github.com/bazelbuild/bazel/pull/21131#discussion_r1471924084
+  create_new_workspace
+  cat > MODULE.bazel <<EOF
+ext = use_extension("//:r.bzl", "ext")
+use_repo(ext, "r")
+EOF
+  touch BUILD
+  cat > r.bzl <<EOF
+load(":make_repo_rule.bzl", "make_repo_rule")
+def _r(rctx):
+  print("I'm here")
+  rctx.file("BUILD", "filegroup(name='r')")
+r=make_repo_rule(_r)
+ext=module_extension(lambda mctx: r(name='r'))
+EOF
+  cat > make_repo_rule.bzl << EOF
+def make_repo_rule(impl):
+  return repository_rule(impl)
+EOF
+
+  bazel build @r >& $TEST_log || fail "Failed to build"
+  expect_log "I'm here"
+
+  cat <<EOF >>r.bzl
+
+# Just add a comment
+EOF
+  # the above SHOULD trigger a refetch.
+  bazel build @r >& $TEST_log || fail "failed to build"
+  expect_log "I'm here"
+}
+
 # Test invalidation based on change to the bzl files
 function file_invalidation_test_template() {
   local startup_flag="${1-}"