Ensure that the rust_proto_library is in the same package as proto_library

PiperOrigin-RevId: 628308230

rust/BUILD

@@ -172,6 +172,11 @@ proto_lang_toolchain(
     visibility = ["//visibility:public"],
 )
 
+package_group(
+    name = "rust_proto_library_allowed_in_different_package",
+    packages = ["//rust/test"],  # for unittest proto_libraries
+)
+
 # This flag controls what kernel all Rust Protobufs are using in the current build.
 string_flag(
     name = "rust_proto_library_kernel",

rust/aspects.bzl

@@ -33,6 +33,12 @@ RustProtoInfo = provider(
     },
 )
 
+def proto_rust_toolchain_label(is_upb):
+    if is_upb:
+        return "//rust:proto_rust_upb_toolchain"
+    else:
+        return "//rust:proto_rust_cpp_toolchain"
+
 def _register_crate_mapping_write_action(name, actions, crate_mappings):
     """Registers an action that generates a crate mapping for a proto_library.
 
@@ -401,9 +407,7 @@ def _make_proto_library_aspect(is_upb):
                 cfg = "exec",
             ),
             "_proto_lang_toolchain": attr.label(
-                default = Label(
+                default = Label(proto_rust_toolchain_label(is_upb)),
-                    "//rust:proto_rust_upb_toolchain" if is_upb else "//rust:proto_rust_cpp_toolchain",
-                ),
             ),
         },
         fragments = ["cpp"],

rust/defs.bzl

@@ -3,9 +3,11 @@
 Disclaimer: This project is experimental, under heavy development, and should not
 be used yet."""
 
+load("@rules_proto//proto:defs.bzl", "ProtoInfo", "proto_common")
 load(
     "//rust:aspects.bzl",
     "RustProtoInfo",
+    "proto_rust_toolchain_label",
     "rust_cc_proto_library_aspect",
     "rust_upb_proto_library_aspect",
 )
@@ -65,6 +67,7 @@ def _rust_proto_library_impl(ctx):
 
     dep = deps[0]
     rust_proto_info = dep[RustProtoInfo]
+
     dep_variant_info = rust_proto_info.dep_variant_info
     return [
         dep_variant_info.crate_info,
@@ -82,6 +85,9 @@ def _make_rust_proto_library(is_upb):
                 providers = [ProtoInfo],
                 aspects = [rust_upb_proto_library_aspect if is_upb else rust_cc_proto_library_aspect],
             ),
+            "_proto_lang_toolchain": attr.label(
+                default = Label(proto_rust_toolchain_label(is_upb)),
+            ),
         },
     )