Drop the notion of 'primary' file in all rs code generation.

Before this change, the way it works is that we emit one .rs file per input .proto file, and a multi-src proto_library are handled by considering the first file as specially the 'primary' one which specially re-exports everything defined in other files.

After this change, we instead emit the .rs file per .proto file equally, and then we additionally emit a generated.rs file which re-exports all of them.

PiperOrigin-RevId: 707569226

rust/aspects.bzl

@@ -106,6 +106,12 @@ def _generate_rust_gencode(
         proto_info = proto_info,
         extension = ".{}.pb.rs".format("u" if is_upb else "c"),
     )
+
+    entry_point_rs_output = actions.declare_file(
+        "{}.generated.{}.rs".format(ctx.label.name, "u" if is_upb else "c"),
+        sibling = proto_info.direct_sources[0],
+    )
+
     if is_upb:
         cc_outputs = []
     else:
@@ -117,9 +123,10 @@ 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={},bazel_crate_mapping={},generated_entry_point_rs_file_name={}".format(
             "upb" if is_upb else "cpp",
             crate_mapping.path,
+            entry_point_rs_output.basename,
         ),
     )
 
@@ -128,10 +135,10 @@ def _generate_rust_gencode(
         proto_info = proto_info,
         additional_inputs = depset(direct = [crate_mapping]),
         additional_args = additional_args,
-        generated_files = rs_outputs + cc_outputs,
+        generated_files = [entry_point_rs_output] + rs_outputs + cc_outputs,
         proto_lang_toolchain_info = proto_lang_toolchain,
     )
-    return (rs_outputs, cc_outputs)
+    return (entry_point_rs_output, rs_outputs, cc_outputs)
 
 def _get_crate_info(providers):
     for provider in providers:
@@ -316,7 +323,7 @@ def _rust_proto_aspect_common(target, ctx, is_upb):
         mapping_for_current_target,
     )
 
-    (gencode, thunks) = _generate_rust_gencode(
+    (entry_point_rs_output, rs_gencode, cc_thunks_gencode) = _generate_rust_gencode(
         ctx,
         target[ProtoInfo],
         proto_lang_toolchain,
@@ -338,7 +345,7 @@ def _rust_proto_aspect_common(target, ctx, is_upb):
             attr = attr,
             cc_toolchain = cc_toolchain,
             cc_infos = [target[CcInfo]] + [dep[CcInfo] for dep in ctx.attr._cpp_thunks_deps] + dep_cc_infos,
-        ) for thunk in thunks])
+        ) for thunk in cc_thunks_gencode])
 
     runtime = proto_lang_toolchain.runtime
     dep_variant_info_for_runtime = DepVariantInfo(
@@ -357,8 +364,8 @@ def _rust_proto_aspect_common(target, ctx, is_upb):
         dep_variant_info = _compile_rust(
             ctx = ctx,
             attr = ctx.rule.attr,
-            src = gencode[0],
-            extra_srcs = gencode[1:],
+            src = entry_point_rs_output,
+            extra_srcs = rs_gencode,
             deps = [dep_variant_info_for_runtime, dep_variant_info_for_native_gencode] + dep_variant_infos,
             runtime = runtime,
         )

rust/release_crates/protobuf_example/src/main.rs

@@ -1,7 +1,7 @@
 use protobuf::proto;
 
 mod protos {
-    include!(concat!(env!("OUT_DIR"), "/protobuf_generated/proto_example/foo.u.pb.rs"));
+    include!(concat!(env!("OUT_DIR"), "/protobuf_generated/proto_example/generated.rs"));
 }
 use protos::Foo;
 

rust/test/rust_proto_library_unit_test/rust_proto_library_unit_test.bzl

@@ -173,12 +173,20 @@ def _rust_outputs_test_impl(ctx):
     env = analysistest.begin(ctx)
     target_under_test = analysistest.target_under_test(env)
 
-    label_to_file = {
-        "child_cpp_rust_proto": "child.c.pb.rs",
-        "child_upb_rust_proto": "child.u.pb.rs",
+    label_to_files = {
+        "child_cpp_rust_proto": ["generated.c.rs", "child.c.pb.rs"],
+        "child_upb_rust_proto": ["generated.u.rs", "child.u.pb.rs"],
     }
-    expected_output = label_to_file[target_under_test.label.name]
-    asserts.true(env, target_under_test.files.to_list()[0].path.endswith(expected_output))
+    expected_outputs = label_to_files[target_under_test.label.name]
+    actual_outputs = target_under_test.files.to_list()
+    asserts.equals(env, len(expected_outputs), len(actual_outputs))
+    for expected in expected_outputs:
+        found = False
+        for actual in actual_outputs:
+            if actual.path.endswith(expected):
+                found = True
+                break
+        asserts.true(env, found)
 
     return analysistest.end(env)
 

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

@@ -78,6 +78,15 @@ absl::StatusOr<Options> Options::Parse(absl::string_view param) {
     opts.strip_nonfunctional_codegen = true;
   }
 
+  auto generated_entry_point_rs_file_name_arg =
+      absl::c_find_if(args, [](auto& arg) {
+        return arg.first == "generated_entry_point_rs_file_name";
+      });
+  if (generated_entry_point_rs_file_name_arg != args.end()) {
+    opts.generated_entry_point_rs_file_name =
+        generated_entry_point_rs_file_name_arg->second;
+  }
+
   return opts;
 }
 

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

@@ -10,6 +10,7 @@
 
 #include <algorithm>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "absl/container/flat_hash_map.h"
@@ -48,6 +49,9 @@ struct Options {
   std::string mapping_file_path;
   bool strip_nonfunctional_codegen = false;
 
+  // The name to use for the generated entry point rs file.
+  std::string generated_entry_point_rs_file_name = "generated.rs";
+
   static absl::StatusOr<Options> Parse(absl::string_view param);
 };
 
@@ -128,11 +132,16 @@ class Context {
     auto it =
         rust_generator_context_->import_path_to_crate_name_.find(import_path);
     if (it == rust_generator_context_->import_path_to_crate_name_.end()) {
-      ABSL_LOG(FATAL)
+      ABSL_LOG(ERROR)
           << "Path " << import_path
-          << " not found in crate mapping. Crate mapping has "
+          << " not found in crate mapping. Crate mapping contains "
           << rust_generator_context_->import_path_to_crate_name_.size()
-          << " entries";
+          << " entries:";
+      for (const auto& entry :
+           rust_generator_context_->import_path_to_crate_name_) {
+        ABSL_LOG(ERROR) << "  " << entry.first << " : " << entry.second << "\n";
+      }
+      ABSL_LOG(FATAL) << "Cannot continue with missing crate mapping.";
     }
     return it->second;
   }

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

@@ -93,27 +93,39 @@ void EmitPublicImports(const RustGeneratorContext& rust_generator_context,
   }
 }
 
-// Emits submodule declarations so `rustc` can find non primary sources from
-// the primary file.
-void DeclareSubmodulesForNonPrimarySrcs(
-    Context& ctx, const FileDescriptor& primary_file,
-    absl::Span<const FileDescriptor* const> non_primary_srcs) {
-  std::string primary_file_path = GetRsFile(ctx, primary_file);
-  RelativePath primary_relpath(primary_file_path);
-  for (const FileDescriptor* non_primary_src : non_primary_srcs) {
-    std::string non_primary_file_path = GetRsFile(ctx, *non_primary_src);
+void EmitEntryPointRsFile(GeneratorContext* generator_context,
+                          Context& ctx_without_printer,
+                          const std::vector<const FileDescriptor*>& files) {
+  // Besides the one .rs file per .proto file, we additional emit one
+  // entry_point rs file here which re-exports all of the types generated by
+  // this same proto_library.
+  std::string entry_point_rs_file_path =
+      GetEntryPointRsFilePath(ctx_without_printer, *files.front());
+  auto outfile =
+      absl::WrapUnique(generator_context->Open(entry_point_rs_file_path));
+  io::Printer printer(outfile.get());
+  Context ctx = ctx_without_printer.WithPrinter(&printer);
+
+  // Declare the submodules for all of the the generated code and pub re-export
+  // all of them into a flat namespace.
+  RelativePath primary_relpath(entry_point_rs_file_path);
+  for (const FileDescriptor* file : files) {
+    std::string non_primary_file_path = GetRsFile(ctx, *file);
     std::string relative_mod_path =
         primary_relpath.Relative(RelativePath(non_primary_file_path));
+    // Temporarily emit these re-exported mods as pub to avoid issues with
+    // Crubit. In a future change we should change these back to be private
+    // mods.
     ctx.Emit({{"file_path", relative_mod_path},
-              {"mod_name", RustInternalModuleName(*non_primary_src)}},
+              {"mod_name", RustInternalModuleName(*file)}},
              R"rs(
-                        #[path="$file_path$"]
-                        #[allow(non_snake_case)]
-                        mod $mod_name$;
+              #[path="$file_path$"]
+              #[allow(non_snake_case)]
+              pub mod internal_do_not_use_$mod_name$;
 
-                        #[allow(unused_imports)]
-                        pub use $mod_name$::*;
-                      )rs");
+              #[allow(unused_imports)]
+              pub use internal_do_not_use_$mod_name$::*;
+            )rs");
   }
 }
 
@@ -143,10 +155,7 @@ bool RustGenerator::Generate(const FileDescriptor* file,
                                               &*import_path_to_crate_name);
 
   std::vector<std::string> modules;
-  if (file != files_in_current_crate[0]) {
-    // This is not the primary file, so its generated code will be in a module.
-    modules.emplace_back(RustInternalModuleName(*file));
-  }
+  modules.emplace_back(RustInternalModuleName(*file));
   Context ctx_without_printer(&*opts, &rust_generator_context, nullptr,
                               std::move(modules));
 
@@ -178,10 +187,11 @@ bool RustGenerator::Generate(const FileDescriptor* file,
   std::vector<const FileDescriptor*> file_contexts(
       files_in_current_crate.begin(), files_in_current_crate.end());
 
-  // Generating the primary file?
-  if (file == &rust_generator_context.primary_file()) {
-    auto non_primary_srcs = absl::MakeConstSpan(file_contexts).subspan(1);
-    DeclareSubmodulesForNonPrimarySrcs(ctx, *file, non_primary_srcs);
+  // When the generator is called for the 'first' file we also want to emit the
+  // 'entry point' rs file. This is the file that will simply pub re-export all
+  // everything from all of the other generated .rs files.
+  if (file == files_in_current_crate.front()) {
+    EmitEntryPointRsFile(generator_context, ctx_without_printer, file_contexts);
   }
 
   std::unique_ptr<io::ZeroCopyOutputStream> thunks_cc;

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

@@ -22,7 +22,6 @@
 #include "absl/strings/str_replace.h"
 #include "absl/strings/string_view.h"
 #include "absl/strings/strip.h"
-#include "absl/strings/substitute.h"
 #include "google/protobuf/compiler/code_generator.h"
 #include "google/protobuf/compiler/cpp/helpers.h"
 #include "google/protobuf/compiler/rust/context.h"
@@ -43,6 +42,14 @@ std::string GetCrateName(Context& ctx, const FileDescriptor& dep) {
   return absl::StrCat("::", RsSafeName(ctx.ImportPathToCrateName(dep.name())));
 }
 
+std::string GetEntryPointRsFilePath(Context& ctx, const FileDescriptor& file) {
+  size_t last_slash = file.name().find_last_of('/');
+  std::string dir = last_slash == std::string::npos
+                        ? ""
+                        : file.name().substr(0, last_slash + 1);
+  return absl::StrCat(dir, ctx.opts().generated_entry_point_rs_file_name);
+}
+
 std::string GetRsFile(Context& ctx, const FileDescriptor& file) {
   auto basename = StripProto(file.name());
   switch (auto k = ctx.opts().kernel) {
@@ -232,7 +239,11 @@ std::string RustModule(Context& ctx, const OneofDescriptor& oneof) {
 
 std::string RustInternalModuleName(const FileDescriptor& file) {
   return RsSafeName(
-      absl::StrReplaceAll(StripProto(file.name()), {{"_", "__"}, {"/", "_s"}}));
+      absl::StrReplaceAll(StripProto(file.name()), {
+                                                       {"_", "__"},
+                                                       {"/", "_s"},
+                                                       {"-", "__"},
+                                                   }));
 }
 
 std::string FieldInfoComment(Context& ctx, const FieldDescriptor& field) {

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

@@ -22,6 +22,12 @@ namespace compiler {
 namespace rust {
 std::string GetCrateName(Context& ctx, const FileDescriptor& dep);
 
+// Gets the file name for the entry point rs file. This path will be in the same
+// directory as the provided file. This will be the path provided by command
+// line flag, or a default path relative to the provided `file` (which should
+// be the first .proto src proto file).
+std::string GetEntryPointRsFilePath(Context& ctx, const FileDescriptor& file);
+
 std::string GetRsFile(Context& ctx, const FileDescriptor& file);
 std::string GetThunkCcFile(Context& ctx, const FileDescriptor& file);
 std::string GetHeaderFile(Context& ctx, const FileDescriptor& file);