Rust: bundle the well-known types in the protobuf crate

This is a hacky first attempt.

PiperOrigin-RevId: 703645757

rust/aspects.bzl

@@ -117,7 +117,7 @@ def _generate_rust_gencode(
     additional_args = ctx.actions.args()
 
     additional_args.add(
-        "--rust_opt=experimental-codegen=enabled,kernel={},bazel_crate_mapping={}".format(
+        "--rust_opt=experimental-codegen=enabled,kernel={},build_system=bazel,bazel_crate_mapping={}".format(
             "upb" if is_upb else "cpp",
             crate_mapping.path,
         ),

rust/release_crates/BUILD

@@ -8,6 +8,7 @@ sh_binary(
         "//rust/release_crates/protobuf:protobuf_crate",
         "//rust/release_crates/protobuf_codegen:protobuf_codegen_crate",
         "//rust/release_crates/protobuf_example:protobuf_example_crate",
+        "//src/google/protobuf/compiler:protoc",
     ],
     tags = ["manual"],
     deps = ["@bazel_tools//tools/bash/runfiles"],

rust/release_crates/cargo_test.sh

@@ -56,7 +56,10 @@ mkdir $EXAMPLE_ROOT
 EXAMPLE_TAR=$(rlocation com_google_protobuf/rust/release_crates/protobuf_example/protobuf_example_crate.tar)
 
 echo "Expanding protobuf_example crate tar"
-tar -xvf $EXAMPLE_TAR -C $EXAMPLE_ROOT 
+tar -xvf $EXAMPLE_TAR -C $EXAMPLE_ROOT
+
+# Put the Bazel-built protoc at the beginning of $PATH
+PATH=$(dirname $(rlocation com_google_protobuf/protoc)):$PATH
 
 cd $CRATE_ROOT
 CARGO_HOME=$CARGO_HOME cargo test

rust/release_crates/protobuf/BUILD

@@ -6,6 +6,7 @@ pkg_tar(
     name = "protobuf_crate",
     srcs = [
         ":crate_root_files",
+        ":well_known_types",
         "//:LICENSE",
         "//rust:rust_protobuf_libupb_src",
         "//rust:rust_protobuf_src_dir",
@@ -35,3 +36,9 @@ filegroup(
     ],
     visibility = ["//rust:__subpackages__"],
 )
+
+pkg_files(
+    name = "well_known_types",
+    srcs = ["//src/google/protobuf:well_known_type_protos"],
+    prefix = "google/protobuf",
+)

rust/release_crates/protobuf/build.rs

@@ -8,6 +8,54 @@
 const VERSION: &str = env!("CARGO_PKG_VERSION");
 
 fn main() {
+    // Generate code for the well-known types
+    let well_known_types = &[
+        "google/protobuf/any.proto",
+        "google/protobuf/api.proto",
+        "google/protobuf/duration.proto",
+        "google/protobuf/empty.proto",
+        "google/protobuf/field_mask.proto",
+        "google/protobuf/source_context.proto",
+        "google/protobuf/struct.proto",
+        "google/protobuf/timestamp.proto",
+        "google/protobuf/type.proto",
+        "google/protobuf/wrappers.proto",
+    ];
+    for f in well_known_types {
+        let mut cmd = std::process::Command::new("protoc");
+        cmd.arg(format!("--rust_out={}", std::env::var("OUT_DIR").unwrap()))
+            .arg("--rust_opt=experimental-codegen=enabled,kernel=upb,build_system=bazel")
+            .arg(f);
+        let output = cmd.output().map_err(|e| format!("failed to run protoc: {}", e)).unwrap();
+        println!("{}", std::str::from_utf8(&output.stdout).unwrap());
+        eprintln!("{}", std::str::from_utf8(&output.stderr).unwrap());
+        assert!(output.status.success());
+    }
+
+    let generated_files = &[
+        "google/protobuf/any.u.pb.rs",
+        "google/protobuf/api.u.pb.rs",
+        "google/protobuf/duration.u.pb.rs",
+        "google/protobuf/empty.u.pb.rs",
+        "google/protobuf/field_mask.u.pb.rs",
+        "google/protobuf/source_context.u.pb.rs",
+        "google/protobuf/struct.u.pb.rs",
+        "google/protobuf/timestamp.u.pb.rs",
+        "google/protobuf/type.u.pb.rs",
+        "google/protobuf/wrappers.u.pb.rs",
+    ];
+    for f in generated_files {
+        let mut cmd = std::process::Command::new("sed");
+        cmd.arg("-i");
+        cmd.arg("-e");
+        cmd.arg("s/::protobuf/crate/g");
+        cmd.arg(format!("{}/{}", std::env::var("OUT_DIR").unwrap(), f));
+        let output = cmd.output().map_err(|e| format!("failed to run sed: {}", e)).unwrap();
+        println!("{}", std::str::from_utf8(&output.stdout).unwrap());
+        eprintln!("{}", std::str::from_utf8(&output.stderr).unwrap());
+        assert!(output.status.success());
+    }
+
     cc::Build::new()
         .flag("-std=c99")
         // TODO: Come up with a way to enable lto

rust/release_crates/protobuf_codegen/src/lib.rs

@@ -98,7 +98,7 @@ impl CodeGen {
         }
 
         cmd.arg(format!("--rust_out={}", self.output_dir.display()))
-            .arg("--rust_opt=experimental-codegen=enabled,kernel=upb")
+            .arg("--rust_opt=experimental-codegen=enabled,kernel=upb,build_system=cargo")
             .arg(format!(
                 "--plugin=protoc-gen-upb_minitable={}",
                 protoc_gen_upb_minitable_path.display()

rust/shared.rs

@@ -104,3 +104,29 @@ impl fmt::Display for SerializeError {
         write!(f, "Couldn't serialize proto into bytes (depth too deep or missing required fields)")
     }
 }
+
+#[cfg(not(bzl))]
+pub mod well_known_types {
+    include!(concat!(env!("OUT_DIR"), "/google/protobuf/any.u.pb.rs"));
+
+    include!(concat!(env!("OUT_DIR"), "/google/protobuf/api.u.pb.rs"));
+
+    include!(concat!(env!("OUT_DIR"), "/google/protobuf/duration.u.pb.rs"));
+
+    include!(concat!(env!("OUT_DIR"), "/google/protobuf/empty.u.pb.rs"));
+
+    include!(concat!(env!("OUT_DIR"),
+"/google/protobuf/field_mask.u.pb.rs"));
+
+    include!(concat!(env!("OUT_DIR"),
+"/google/protobuf/source_context.u.pb.rs"));
+
+    include!(concat!(env!(" OUT_DIR"), "/google/protobuf/struct.u.pb.rs"));
+
+    include!(concat!(env!(" OUT_DIR"),
+"/google/protobuf/timestamp.u.pb.rs"));
+
+    include!(concat!(env! ("OUT_DIR"), "/google/protobuf/type.u.pb.rs"));
+
+    include!(concat!(env!(" OUT_DIR"), "/google/protobuf/wrappers.u.pb.rs"));
+}

src/google/protobuf/compiler/rust/context.cc

@@ -65,6 +65,22 @@ absl::StatusOr<Options> Options::Parse(absl::string_view param) {
                          kernel_arg->second));
   }
 
+  auto build_system_arg = absl::c_find_if(
+      args, [](auto& arg) { return arg.first == "build_system"; });
+  if (build_system_arg == args.end()) {
+    return absl::InvalidArgumentError(
+        "Mandatory option `build_system` missing, please specify `bazel` or "
+        "`cargo`.");
+  } else if (build_system_arg->second == "bazel") {
+    opts.build_system = BuildSystem::kBazel;
+  } else if (build_system_arg->second == "cargo") {
+    opts.build_system = BuildSystem::kCargo;
+  } else {
+    return absl::InvalidArgumentError(absl::Substitute(
+        "Unknown build system `$0`, please specify `bazel` or `cargo`.",
+        build_system_arg->second));
+  }
+
   auto mapping_arg = absl::c_find_if(
       args, [](auto& arg) { return arg.first == "bazel_crate_mapping"; });
   if (mapping_arg != args.end()) {

src/google/protobuf/compiler/rust/context.h

@@ -30,6 +30,11 @@ enum class Kernel {
   kCpp,
 };
 
+enum class BuildSystem {
+  kBazel,
+  kCargo,
+};
+
 inline absl::string_view KernelRsName(Kernel kernel) {
   switch (kernel) {
     case Kernel::kUpb:
@@ -45,6 +50,7 @@ inline absl::string_view KernelRsName(Kernel kernel) {
 // Global options for a codegen invocation.
 struct Options {
   Kernel kernel;
+  BuildSystem build_system;
   std::string mapping_file_path;
   bool strip_nonfunctional_codegen = false;
 

src/google/protobuf/compiler/rust/naming.cc

@@ -40,6 +40,10 @@ namespace compiler {
 namespace rust {
 
 std::string GetCrateName(Context& ctx, const FileDescriptor& dep) {
+  if (ctx.opts().build_system == BuildSystem::kCargo &&
+      cpp::IsWellKnownMessage(&dep)) {
+    return absl::StrCat("::protobuf::well_known_types");
+  }
   return absl::StrCat("::", RsSafeName(ctx.ImportPathToCrateName(dep.name())));
 }