Provide an option to opt-out of compile_well_known_types

extensions/prost/private/prost.bzl

@@ -75,13 +75,15 @@ def _compile_proto(
     additional_args = ctx.actions.args()
 
     # Prost process wrapper specific args
+    compile_well_known_types = prost_toolchain.compile_well_known_types
     additional_args.add("--protoc={}".format(proto_compiler.executable.path))
     additional_args.add("--label={}".format(ctx.label))
     additional_args.add("--out_librs={}".format(lib_rs.path))
     additional_args.add("--package_info_output={}".format("{}={}".format(crate_name, package_info_file.path)))
     additional_args.add("--deps_info={}".format(deps_info_file.path))
     additional_args.add("--direct_dep_crate_names={}".format(",".join(direct_crate_names)))
-    additional_args.add("--prost_opt=compile_well_known_types")
+    if compile_well_known_types:
+        additional_args.add("--compile_well_known_types")
     additional_args.add("--descriptor_set={}".format(proto_info.direct_descriptor_set.path))
     additional_args.add("--additional_srcs={}".format(",".join([f.path for f in all_additional_srcs.to_list()])))
     additional_args.add_all(prost_toolchain.prost_opts + prost_opts, format_each = "--prost_opt=%s")
@@ -90,7 +92,6 @@ def _compile_proto(
         tonic_plugin = prost_toolchain.tonic_plugin[DefaultInfo].files_to_run
         additional_args.add(prost_toolchain.tonic_plugin_flag % tonic_plugin.executable.path)
         additional_args.add("--tonic_opt=no_include")
-        additional_args.add("--tonic_opt=compile_well_known_types")
         additional_args.add("--is_tonic")
 
         additional_args.add_all(prost_toolchain.tonic_opts + tonic_opts, format_each = "--tonic_opt=%s")
@@ -232,7 +233,10 @@ def _rust_prost_aspect_impl(target, ctx):
 
     rustfmt_toolchain = ctx.toolchains["@rules_rust//rust/rustfmt:toolchain_type"]
     prost_toolchain = ctx.toolchains[TOOLCHAIN_TYPE]
-    for prost_runtime in [prost_toolchain.prost_runtime, prost_toolchain.tonic_runtime]:
+    runtimes = [prost_toolchain.prost_runtime, prost_toolchain.tonic_runtime]
+    if not prost_toolchain.compile_well_known_types:
+        runtimes += [prost_toolchain.prost_types]
+    for prost_runtime in runtimes:
         if not prost_runtime:
             continue
         if rust_common.crate_group_info in prost_runtime:
@@ -442,13 +446,18 @@ def _rust_prost_toolchain_impl(ctx):
         tonic_plugin_flag = ctx.attr.tonic_plugin_flag,
         tonic_runtime = ctx.attr.tonic_runtime,
         include_transitive_deps = ctx.attr.include_transitive_deps,
+        compile_well_known_types = ctx.attr.compile_well_known_types,
     )]
 
 rust_prost_toolchain = rule(
     implementation = _rust_prost_toolchain_impl,
     doc = "Rust Prost toolchain rule.",
     fragments = ["proto"],
     attrs = dict({
+        "compile_well_known_types": attr.bool(
+            doc = "Corresponds to prost_build's `compile_well_known_types` option. If set to False, well-known-types will not be compiled by prost, and instead rely on the provided Prost types crate.",
+            default = True,
+        ),
         "include_transitive_deps": attr.bool(
             doc = "Whether to include transitive dependencies. If set to True, all transitive dependencies will directly accessible by the dependent crate.",
             default = False,

extensions/prost/private/protoc_wrapper.rs

@@ -314,11 +314,16 @@ impl From<&str> for RustModulePath {
 fn get_extern_paths(
     descriptor_set: &FileDescriptorSet,
     crate_name: &str,
+    compile_well_known_types: bool,
 ) -> Result<BTreeMap<ProtoPath, RustModulePath>, String> {
     let mut extern_paths = BTreeMap::new();
     let rust_path = RustModulePath(escape_keyword(crate_name.to_string()));
 
     for file in descriptor_set.file.iter() {
+        if !compile_well_known_types && file.package.as_deref() == Some("google.protobuf") {
+            continue;
+        }
+
         descriptor_set_file_to_extern_paths(&mut extern_paths, &rust_path, file);
     }
 
@@ -451,6 +456,9 @@ struct Args {
     /// Direct dependency crate names.
     direct_dep_crate_names: Vec<String>,
 
+    /// Whether to compile well known types.
+    compile_well_known_types: bool,
+
     /// The path to the rustfmt binary.
     rustfmt: Option<PathBuf>,
 
@@ -479,6 +487,7 @@ impl Args {
         let mut tonic_or_prost_opts = Vec::new();
         let mut direct_dep_crate_names = Vec::new();
         let mut is_tonic = false;
+        let mut compile_well_known_types = false;
 
         let mut extra_args = Vec::new();
 
@@ -502,6 +511,11 @@ impl Args {
                 return;
             }
 
+            if arg == "--compile_well_known_types" {
+                compile_well_known_types = true;
+                return;
+            }
+
             if !arg.contains('=') {
                 extra_args.push(arg);
                 return;
@@ -646,6 +660,7 @@ impl Args {
             is_tonic,
             label: label.unwrap(),
             extra_args,
+            compile_well_known_types,
         })
     }
 }
@@ -749,6 +764,7 @@ fn main() {
         direct_dep_crate_names,
         is_tonic,
         extra_args,
+        compile_well_known_types,
     } = Args::parse().expect("Failed to parse args");
 
     let out_dir = get_and_create_output_dir(&out_dir, &label);
@@ -794,6 +810,14 @@ fn main() {
     if is_tonic {
         args.push(format!("--tonic_out={}", out_dir.display()));
     }
+
+    if compile_well_known_types {
+        args.push("--prost_opt=compile_well_known_types".to_owned());
+        if is_tonic {
+            args.push("--tonic_opt=compile_well_known_types".to_owned());
+        }
+    }
+
     args.extend(extra_args);
     args.extend(
         proto_paths
@@ -906,7 +930,7 @@ fn main() {
         }
     }
 
-    let extern_paths = get_extern_paths(&descriptor_set, &crate_name)
+    let extern_paths = get_extern_paths(&descriptor_set, &crate_name, compile_well_known_types)
         .expect("Failed to compute proto package info");
 
     // Write outputs