Read crate mapping in protoc

This enables us to get the correct crate names for Rust gencode.

PiperOrigin-RevId: 597850176

pkg/BUILD.bazel

@@ -243,5 +243,6 @@ cc_dist_library(
         "//src/google/protobuf/compiler:command_line_interface_tester",
         "//src/google/protobuf/compiler:mock_code_generator",
         "//src/google/protobuf/testing",
+        "//src/google/protobuf/testing:file",
     ],
 )

pkg/cc_dist_library.bzl

@@ -293,7 +293,9 @@ def _subtract_files(a, b):
     cc_file_list_args = {}
     file_list_fields = ["srcs", "hdrs", "internal_hdrs", "textual_hdrs"]
     for field in file_list_fields:
-        to_remove = {e: None for e in getattr(b.cc_file_list, field).to_list()}
+        # only a subset of file.cc is used from protoc, to get all its symbols for tests we need to
+        # also build & link it to tests.
+        to_remove = {e: None for e in getattr(b.cc_file_list, field).to_list() if "src/google/protobuf/testing/file.cc" not in e.path}
         cc_file_list_args[field] = depset(
             [e for e in getattr(a.cc_file_list, field).to_list() if not e in to_remove],
         )

src/google/protobuf/compiler/rust/BUILD.bazel

@@ -17,6 +17,7 @@ cc_library(
     ],
     deps = [
         ":context",
+        ":crate_mapping",
         ":enum",
         ":message",
         ":naming",
@@ -29,9 +30,11 @@ cc_library(
         "//src/google/protobuf/io:printer",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:btree",
+        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:optional",
@@ -39,6 +42,39 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "crate_mapping",
+    srcs = ["crate_mapping.cc"],
+    hdrs = ["crate_mapping.h"],
+    strip_include_prefix = "/src",
+    deps = [
+        ":context",
+        "//src/google/protobuf",
+        "//src/google/protobuf/testing:file",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_test(
+    name = "crate_mapping_unittest",
+    srcs = ["crate_mapping_unittest.cc"],
+    deps = [
+        ":context",
+        ":crate_mapping",
+        "//src/google/protobuf/testing",
+        "//src/google/protobuf/testing:file",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "message",
     srcs = ["message.cc"],
@@ -104,6 +140,8 @@ cc_library(
         "//src/google/protobuf/compiler:code_generator",
         "//src/google/protobuf/io:printer",
         "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/log:absl_log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",

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

@@ -65,6 +65,15 @@ absl::StatusOr<Options> Options::Parse(absl::string_view param) {
                          kernel_arg->second));
   }
 
+  auto mapping_arg = absl::c_find_if(
+      args, [](auto& arg) { return arg.first == "bazel_crate_mapping"; });
+  if (mapping_arg == args.end()) {
+    return absl::InvalidArgumentError(
+        "Mandatory option `mapping` missing, please specify mapping file "
+        "path.");
+  }
+  opts.mapping_file_path = mapping_arg->second;
+
   return opts;
 }
 

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

@@ -9,8 +9,11 @@
 #define GOOGLE_PROTOBUF_COMPILER_RUST_CONTEXT_H__
 
 #include <algorithm>
+#include <string>
 #include <vector>
 
+#include "absl/container/flat_hash_map.h"
+#include "absl/log/absl_check.h"
 #include "absl/log/absl_log.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
@@ -43,6 +46,7 @@ inline absl::string_view KernelRsName(Kernel kernel) {
 // Global options for a codegen invocation.
 struct Options {
   Kernel kernel;
+  std::string mapping_file_path;
 
   static absl::StatusOr<Options> Parse(absl::string_view param);
 };
@@ -50,8 +54,11 @@ struct Options {
 class RustGeneratorContext {
  public:
   explicit RustGeneratorContext(
-      const std::vector<const FileDescriptor*>* files_in_current_crate)
-      : files_in_current_crate_(*files_in_current_crate) {}
+      const std::vector<const FileDescriptor*>* files_in_current_crate,
+      const absl::flat_hash_map<std::string, std::string>*
+          import_path_to_crate_name)
+      : files_in_current_crate_(*files_in_current_crate),
+        import_path_to_crate_name_(*import_path_to_crate_name) {}
 
   const FileDescriptor& primary_file() const {
     return *files_in_current_crate_.front();
@@ -63,8 +70,14 @@ class RustGeneratorContext {
                      &f) != files_in_current_crate_.end();
   }
 
+  absl::string_view ImportPathToCrateName(absl::string_view import_path) const {
+    return import_path_to_crate_name_.at(import_path);
+  }
+
  private:
   const std::vector<const FileDescriptor*>& files_in_current_crate_;
+  const absl::flat_hash_map<std::string, std::string>&
+      import_path_to_crate_name_;
 };
 
 // A context for generating a particular kind of definition.

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

@@ -0,0 +1,59 @@
+#include "google/protobuf/compiler/rust/crate_mapping.h"
+
+#include <fcntl.h>
+
+#include <cstddef>
+#include <cstdio>
+#include <string>
+#include <vector>
+
+#include "google/protobuf/testing/file.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/numbers.h"
+#include "absl/strings/str_split.h"
+#include "absl/strings/string_view.h"
+#include "google/protobuf/compiler/rust/context.h"
+
+namespace google {
+namespace protobuf {
+namespace compiler {
+namespace rust {
+
+absl::StatusOr<absl::flat_hash_map<std::string, std::string>>
+GetImportPathToCrateNameMap(const Options* opts) {
+  if (opts->mapping_file_path.empty()) {
+    return absl::InvalidArgumentError("Mapping file path is not specified");
+  }
+  std::string mapping_contents;
+  absl::Status status =
+      File::ReadFileToString(opts->mapping_file_path, &mapping_contents, true);
+  if (!status.ok()) {
+    return status;
+  }
+
+  absl::flat_hash_map<std::string, std::string> mapping;
+  std::vector<absl::string_view> lines =
+      absl::StrSplit(mapping_contents, '\n', absl::SkipEmpty());
+  size_t len = lines.size();
+
+  size_t idx = 0;
+  while (idx < len) {
+    absl::string_view crate_name = lines[idx++];
+    size_t files_cnt;
+    if (!absl::SimpleAtoi(lines[idx++], &files_cnt)) {
+      return absl::InvalidArgumentError(
+          "Couldn't parse number of import paths in mapping file");
+    }
+    for (size_t i = 0; i < files_cnt; ++i) {
+      mapping.insert({std::string(lines[idx++]), std::string(crate_name)});
+    }
+  }
+  return mapping;
+}
+
+}  // namespace rust
+}  // namespace compiler
+}  // namespace protobuf
+}  // namespace google

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

@@ -0,0 +1,42 @@
+#ifndef GOOGLE_PROTOBUF_COMPILER_RUST_MAPPING_FILE_H__
+#define GOOGLE_PROTOBUF_COMPILER_RUST_MAPPING_FILE_H__
+
+#include <string>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/statusor.h"
+#include "google/protobuf/compiler/rust/context.h"
+
+namespace google {
+namespace protobuf {
+namespace compiler {
+namespace rust {
+
+// Returns a map from import path of a .proto file to the name of the crate
+// (proto_library) covering that file.
+//
+// This function parses a .rust_crate_mapping file generated by Bazel. The file
+// contains:
+//
+//    <crate_name>\n
+//    <number of .proto files covered by the proto_library with that name>\n
+//    <import path of the first .proto file of the proto_library>\n
+//    ...
+//    <import path of the last .proto file of the proto_library>\n
+//
+// repeated for each proto_library transitively reachable from the current
+// proto_library.
+//
+//  Note that the logic of translating the proto_library label to a crate name
+//  is handled by the build system completely, protoc is only given the end
+//  results. See `_render_text_crate_mapping` in
+//  //third_party/protobuf/rust:aspects.bzl for how Bazel does this.
+absl::StatusOr<absl::flat_hash_map<std::string, std::string>>
+GetImportPathToCrateNameMap(const Options* opts);
+
+}  // namespace rust
+}  // namespace compiler
+}  // namespace protobuf
+}  // namespace google
+
+#endif  // GOOGLE_PROTOBUF_COMPILER_RUST_MAPPING_FILE_H__

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

@@ -0,0 +1,106 @@
+#include "google/protobuf/compiler/rust/crate_mapping.h"
+
+#include <string>
+
+#include "google/protobuf/testing/file.h"
+#include "google/protobuf/testing/googletest.h"
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/strings/ascii.h"
+#include "absl/strings/str_split.h"
+#include "absl/strings/string_view.h"
+#include "google/protobuf/compiler/rust/context.h"
+
+using google::protobuf::compiler::rust::Options;
+using testing::Eq;
+using testing::UnorderedElementsAreArray;
+
+namespace google {
+namespace protobuf {
+namespace compiler {
+namespace rust {
+
+namespace {
+
+std::string WriteStringToTempFile(absl::string_view text) {
+  std::string path = absl::StrCat(TestTempDir(), "crate_mapping");
+  File::WriteStringToFileOrDie(text, path);
+  return path;
+}
+
+std::string SkipLeadingWhitespace(absl::string_view text) {
+  std::string result;
+  for (absl::string_view line : absl::StrSplit(text, '\n', absl::SkipEmpty())) {
+    absl::string_view stripped_line = absl::StripLeadingAsciiWhitespace(line);
+    // Deal with old libc++ on OSS Protobuf CI workers
+    result.append(stripped_line.data(), stripped_line.size());
+    result.append("\n");
+  }
+  return result;
+}
+
+TEST(RustGenerator, GetImportPathToCrateNameMapEmpty) {
+  std::string mapping_file_path = WriteStringToTempFile("");
+  const Options opts = {Kernel::kUpb, mapping_file_path};
+  absl::flat_hash_map<std::string, std::string> expected = {};
+  auto result = GetImportPathToCrateNameMap(&opts);
+  EXPECT_TRUE(result.ok());
+  EXPECT_THAT(result.value(), Eq(expected));
+}
+
+TEST(RustGenerator, GetImportPathToCrateNameMapSimple) {
+  std::string mapping_file_path =
+      WriteStringToTempFile(SkipLeadingWhitespace(R"mapping(
+    crate_name
+    3
+    proto1.proto
+    google.protobuf.proto
+    proto3.proto
+  )mapping"));
+  const Options opts = {Kernel::kUpb, mapping_file_path};
+  absl::flat_hash_map<std::string, std::string> expected = {
+      {"proto1.proto", "crate_name"},
+      {"google.protobuf.proto", "crate_name"},
+      {"proto3.proto", "crate_name"},
+  };
+  auto result = GetImportPathToCrateNameMap(&opts);
+  EXPECT_TRUE(result.ok());
+  EXPECT_THAT(result.value(), UnorderedElementsAreArray(expected));
+}
+
+TEST(RustGenerator,
+     GetImportPathToCrateNameMapErrorsOnNotSpecifiedMappingFile) {
+  const Options opts = {Kernel::kUpb};
+  auto result = GetImportPathToCrateNameMap(&opts);
+  EXPECT_FALSE(result.ok());
+  EXPECT_THAT(result.status().code(), Eq(absl::StatusCode::kInvalidArgument));
+  EXPECT_THAT(result.status().message(),
+              Eq("Mapping file path is not specified"));
+}
+
+TEST(RustGenerator, GetImportPathToCrateNameMapErrorsOnMissingFile) {
+  const Options opts = {Kernel::kUpb, "missing_file_path"};
+  auto result = GetImportPathToCrateNameMap(&opts);
+  EXPECT_FALSE(result.ok());
+  EXPECT_THAT(result.status().code(), Eq(absl::StatusCode::kNotFound));
+  EXPECT_THAT(result.status().message(), Eq("Could not open file"));
+}
+
+TEST(RustGenerator, GetImportPathToCrateNameMapInvalidFormat) {
+  std::string mapping_file_path =
+      WriteStringToTempFile("crate_name\nnot_a_number");
+  const Options opts = {Kernel::kUpb, mapping_file_path};
+  absl::flat_hash_map<std::string, std::string> expected = {};
+  auto result = GetImportPathToCrateNameMap(&opts);
+  EXPECT_FALSE(result.ok());
+  EXPECT_THAT(result.status().code(), Eq(absl::StatusCode::kInvalidArgument));
+  EXPECT_THAT(result.status().message(),
+              Eq("Couldn't parse number of import paths in mapping file"));
+}
+}  // namespace
+}  // namespace rust
+}  // namespace compiler
+}  // namespace protobuf
+}  // namespace google

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

@@ -7,15 +7,16 @@
 
 #include "google/protobuf/compiler/rust/generator.h"
 
-#include <iterator>
 #include <memory>
 #include <string>
 #include <vector>
 
 #include "absl/algorithm/container.h"
 #include "absl/container/btree_map.h"
+#include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/memory/memory.h"
+#include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_split.h"
 #include "absl/strings/string_view.h"
@@ -23,6 +24,7 @@
 #include "google/protobuf/compiler/code_generator.h"
 #include "google/protobuf/compiler/cpp/names.h"
 #include "google/protobuf/compiler/rust/context.h"
+#include "google/protobuf/compiler/rust/crate_mapping.h"
 #include "google/protobuf/compiler/rust/enum.h"
 #include "google/protobuf/compiler/rust/message.h"
 #include "google/protobuf/compiler/rust/naming.h"
@@ -30,7 +32,6 @@
 #include "google/protobuf/descriptor.h"
 #include "google/protobuf/descriptor.pb.h"
 #include "google/protobuf/io/printer.h"
-#include "google/protobuf/io/zero_copy_stream.h"
 
 namespace google {
 namespace protobuf {
@@ -156,8 +157,8 @@ void EmitPublicImports(Context& ctx, const FileDescriptor& primary_file) {
   }
 }
 
-// Emits submodule declarations so `rustc` can find non primary sources from the
-// primary file.
+// 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) {
@@ -179,8 +180,9 @@ void DeclareSubmodulesForNonPrimarySrcs(
   }
 }
 
-// Emits `pub use <...>::Msg` for all messages in non primary sources into their
-// corresponding packages (each source file can declare a different package).
+// Emits `pub use <...>::Msg` for all messages in non primary sources into
+// their corresponding packages (each source file can declare a different
+// package).
 //
 // Returns the non-primary sources that should be reexported from the package of
 // the primary file.
@@ -222,7 +224,15 @@ bool RustGenerator::Generate(const FileDescriptor* file,
   std::vector<const FileDescriptor*> files_in_current_crate;
   generator_context->ListParsedFiles(&files_in_current_crate);
 
-  RustGeneratorContext rust_generator_context(&files_in_current_crate);
+  absl::StatusOr<absl::flat_hash_map<std::string, std::string>>
+      import_path_to_crate_name = GetImportPathToCrateNameMap(&*opts);
+  if (!import_path_to_crate_name.ok()) {
+    *error = std::string(import_path_to_crate_name.status().message());
+    return false;
+  }
+
+  RustGeneratorContext rust_generator_context(&files_in_current_crate,
+                                              &*import_path_to_crate_name);
 
   Context ctx_without_printer(&*opts, &rust_generator_context, nullptr);
 

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

@@ -41,8 +41,7 @@ std::string GetUnderscoreDelimitedFullName(Context& ctx,
 
 std::string GetCrateName(Context& ctx, const FileDescriptor& dep) {
   absl::string_view path = dep.name();
-  auto basename = path.substr(path.rfind('/') + 1);
-  return absl::StrReplaceAll(basename, {{".", "_"}, {"-", "_"}});
+  return std::string(ctx.generator_context().ImportPathToCrateName(path));
 }
 
 std::string GetRsFile(Context& ctx, const FileDescriptor& file) {

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

@@ -4,6 +4,7 @@
 #include <vector>
 
 #include <gtest/gtest.h>
+#include "absl/container/flat_hash_map.h"
 #include "google/protobuf/compiler/rust/context.h"
 #include "google/protobuf/descriptor.h"
 #include "google/protobuf/io/zero_copy_stream_impl_lite.h"
@@ -27,7 +28,8 @@ TEST(RustProtoNaming, RustInternalModuleName) {
 
   const Options opts = {Kernel::kUpb};
   std::vector<const google::protobuf::FileDescriptor*> files{fd};
-  const RustGeneratorContext rust_generator_context(&files);
+  absl::flat_hash_map<std::string, std::string> mapping;
+  const RustGeneratorContext rust_generator_context(&files, &mapping);
   std::string output;
   StringOutputStream stream{&output};
   Printer printer(&stream);