From dbf97258bd487e684081b9922d09a62018b55c6f Mon Sep 17 00:00:00 2001
From: Jakob Buchgraber <buchgr@google.com>
Date: Wed, 29 Mar 2023 06:10:35 -0700
Subject: [PATCH] Enable support for rust_upb_proto_library and
 rust_cc_proto_library within the same build

PiperOrigin-RevId: 520314976
---
 rust/BUILD                                     | 15 ++++++++++-----
 rust/aspects.bzl                               | 10 ++++++----
 src/google/protobuf/compiler/rust/generator.cc | 13 +++++++++++--
 3 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/rust/BUILD b/rust/BUILD
index fc086c9d60..ff47fccc5a 100644
--- a/rust/BUILD
+++ b/rust/BUILD
@@ -37,11 +37,16 @@ rust_test(
 
 # TODO(b/270125787): Move to the right location once rust_proto_library is no longer experimental.
 proto_lang_toolchain(
-    name = "proto_lang_toolchain",
-    command_line = "--rust_out=experimental-codegen=enabled,kernel=" + select({
-        ":use_upb_kernel": "upb",
-        "//conditions:default": "cpp",
-    }) + ":$(OUT)",
+    name = "proto_rust_upb_toolchain",
+    command_line = "--rust_out=experimental-codegen=enabled,kernel=upb:$(OUT)",
+    progress_message = "Generating Rust proto_library %{label}",
+    runtime = ":protobuf",
+    visibility = ["//visibility:public"],
+)
+
+proto_lang_toolchain(
+    name = "proto_rust_cpp_toolchain",
+    command_line = "--rust_out=experimental-codegen=enabled,kernel=cpp:$(OUT)",
     progress_message = "Generating Rust proto_library %{label}",
     runtime = ":protobuf",
     visibility = ["//visibility:public"],
diff --git a/rust/aspects.bzl b/rust/aspects.bzl
index dd12d09cd7..5285e73990 100644
--- a/rust/aspects.bzl
+++ b/rust/aspects.bzl
@@ -24,7 +24,8 @@ RustProtoInfo = provider(
 def _generate_rust_gencode(
         ctx,
         proto_info,
-        proto_lang_toolchain):
+        proto_lang_toolchain,
+        ext):
     """Generates Rust gencode
 
         This function uses proto_common APIs and a ProtoLangToolchain to register an action
@@ -34,6 +35,7 @@ def _generate_rust_gencode(
         proto_info (ProtoInfo): ProtoInfo of the proto_library target for which we are generating
                     gencode
         proto_lang_toolchain (ProtoLangToolchainInfo): proto lang toolchain for Rust
+        ext: the file extension to use for generated files
     Returns:
         rs_outputs ([File]): generated Rust source files
     """
@@ -41,7 +43,7 @@ def _generate_rust_gencode(
     rs_outputs = proto_common.declare_generated_files(
         actions = actions,
         proto_info = proto_info,
-        extension = ".pb.rs",
+        extension = ext,
     )
 
     proto_common.compile(
@@ -164,7 +166,7 @@ def _rust_proto_aspect_common(target, ctx, is_upb):
 
     proto_lang_toolchain = ctx.attr._proto_lang_toolchain[proto_common.ProtoLangToolchainInfo]
 
-    gencode = _generate_rust_gencode(ctx, target[ProtoInfo], proto_lang_toolchain)
+    gencode = _generate_rust_gencode(ctx, target[ProtoInfo], proto_lang_toolchain, (".u.pb.rs" if is_upb else ".c.pb.rs"))
 
     runtime = proto_lang_toolchain.runtime
     dep_variant_info_for_runtime = DepVariantInfo(
@@ -234,7 +236,7 @@ def _make_proto_library_aspect(is_upb):
                 cfg = "exec",
             ),
             "_proto_lang_toolchain": attr.label(
-                default = Label("//rust:proto_lang_toolchain"),
+                default = Label(("//rust:proto_rust_upb_toolchain" if is_upb else "//rust:proto_rust_cpp_toolchain")),
             ),
         },
         fragments = ["cpp"],
diff --git a/src/google/protobuf/compiler/rust/generator.cc b/src/google/protobuf/compiler/rust/generator.cc
index 09b66b8a92..da024f0897 100644
--- a/src/google/protobuf/compiler/rust/generator.cc
+++ b/src/google/protobuf/compiler/rust/generator.cc
@@ -88,6 +88,15 @@ std::string GetCrateName(const FileDescriptor* dependency) {
                                        });
 }
 
+std::string GetFileExtensionForKernel(Kernel kernel) {
+  switch (kernel) {
+    case Kernel::kUpb:
+      return ".u.pb.rs";
+    case Kernel::kCpp:
+      return ".c.pb.rs";
+  }
+}
+
 void GenerateForUpb(const FileDescriptor* file, google::protobuf::io::Printer& p) {
   for (int i = 0; i < file->message_type_count(); ++i) {
     // The prefix used by the UPB compiler to generate unique function
@@ -171,8 +180,8 @@ bool RustGenerator::Generate(const FileDescriptor* file,
   }
 
   auto basename = StripProto(file->name());
-  auto outfile = absl::WrapUnique(
-      generator_context->Open(absl::StrCat(basename, ".pb.rs")));
+  auto outfile = absl::WrapUnique(generator_context->Open(
+      absl::StrCat(basename, GetFileExtensionForKernel(*kernel))));
 
   google::protobuf::io::Printer p(outfile.get());
   p.Emit(R"rs(