Fix rule definition environment for repo rules

See #21131 (comment)

Closes #21390.

PiperOrigin-RevId: 607821270
Change-Id: I031222ee80c4db4ca50372c549e2b522ac606cb7

src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryModule.java

@@ -107,11 +107,20 @@ 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 threadContext = BazelStarlarkContext.fromOrFail(thread);
+    if (threadContext instanceof BzlInitThreadContext) {
+      var bzlInitContext = (BzlInitThreadContext) threadContext;
+      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) {

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-}"