From dbdddb722af9bcbe4459f3b3c4eae4a2b4b10c7c Mon Sep 17 00:00:00 2001 From: QP Hou Date: Sun, 6 Oct 2024 14:58:58 -0700 Subject: [PATCH] [Bugfix] protos with the same package name prefix fail to compile When proto files sharing the same package name prefixes reference types from each other, protoc_wrapper will generate code with incorrect module scope resulting in the following compilation error: ``` error[E0412]: cannot find type `CMessage` in module `super` --> bazel-out/darwin_arm64-fastbuild/bin/proto/prost/private/tests/shared_package_prefix/shared_pkg_prefix_proto.lib.rs:21:46 | 21 | pub c: ::core::option::Option, | ^^^^^^^^ not found in `super` | help: consider importing this struct | 15 + use crate::pkg::CMessage; | help: if you import `CMessage`, refer to it directly | 21 - pub c: ::core::option::Option, 21 + pub c: ::core::option::Option, ``` This Patch refactors the module tree walk implementation to ensure generated rust code are wrapped with the correct rust module hierarchy. --- proto/prost/private/protoc_wrapper.rs | 81 ++++++++----------- .../tests/shared_package_prefix/BUILD.bazel | 26 ++++++ .../tests/shared_package_prefix/pkg.c.proto | 6 ++ .../shared_package_prefix/pkg.foo.a.proto | 9 +++ .../shared_package_prefix/pkg.foo.b.proto | 6 ++ .../shared_pkg_prefix_test.rs | 22 +++++ 6 files changed, 102 insertions(+), 48 deletions(-) create mode 100644 proto/prost/private/tests/shared_package_prefix/BUILD.bazel create mode 100644 proto/prost/private/tests/shared_package_prefix/pkg.c.proto create mode 100644 proto/prost/private/tests/shared_package_prefix/pkg.foo.a.proto create mode 100644 proto/prost/private/tests/shared_package_prefix/pkg.foo.b.proto create mode 100644 proto/prost/private/tests/shared_package_prefix/shared_pkg_prefix_test.rs diff --git a/proto/prost/private/protoc_wrapper.rs b/proto/prost/private/protoc_wrapper.rs index 9c41892ccf..21371acd76 100644 --- a/proto/prost/private/protoc_wrapper.rs +++ b/proto/prost/private/protoc_wrapper.rs @@ -68,7 +68,7 @@ struct Module { contents: String, /// The names of any other modules which are submodules of this module. - submodules: BTreeSet, + submodules: BTreeMap, } /// Generate a lib.rs file with all prost/tonic outputs embeeded in modules which @@ -118,7 +118,12 @@ struct Module { /// } /// ``` fn generate_lib_rs(prost_outputs: &BTreeSet, is_tonic: bool) -> String { - let mut module_info = BTreeMap::new(); + // dummy root module to host all discovered modules in its submodules map + let mut root_module = Module { + name: "".to_string(), + contents: "".to_string(), + submodules: BTreeMap::new(), + }; for path in prost_outputs.iter() { let mut package = path @@ -154,43 +159,31 @@ fn generate_lib_rs(prost_outputs: &BTreeSet, is_tonic: bool) -> String // Avoid a stack overflow by skipping a known bad package name let module_name = snake_cased_package_name(&package); - - module_info.insert( - module_name.clone(), - Module { - name, - contents: fs::read_to_string(path).expect("Failed to read file"), - submodules: BTreeSet::new(), - }, - ); - + let contents = fs::read_to_string(path).expect("Failed to read file"); let module_parts = module_name.split('.').collect::>(); - for parent_module_index in 0..module_parts.len() { - let child_module_index = parent_module_index + 1; - if child_module_index >= module_parts.len() { - break; - } - let full_parent_module_name = module_parts[0..parent_module_index + 1].join("."); - let parent_module_name = module_parts[parent_module_index]; - let child_module_name = module_parts[child_module_index]; - - module_info - .entry(full_parent_module_name.clone()) - .and_modify(|parent_module| { - parent_module - .submodules - .insert(child_module_name.to_string()); - }) - .or_insert(Module { - name: parent_module_name.to_string(), + + let mut curr_module = &mut root_module; + // Iterate package path to create leaf node + for part in module_parts.iter() { + curr_module = curr_module + .submodules + .entry(part.to_string()) + .or_insert_with(|| Module { + name: part.to_string(), contents: "".to_string(), - submodules: [child_module_name.to_string()].iter().cloned().collect(), + submodules: BTreeMap::new(), }); } + // Assign generated rust code at the leaf node + curr_module.contents = contents; + curr_module.name = name; } let mut content = "// @generated\n\n".to_string(); - write_module(&mut content, &module_info, "", 0); + let module_info = &root_module.submodules; + for submodule_name in module_info.keys() { + write_module(&mut content, module_info, submodule_name, 1); + } content } @@ -201,40 +194,32 @@ fn write_module( module_name: &str, depth: usize, ) { - if module_name.is_empty() { - for submodule_name in module_info.keys() { - write_module(content, module_info, submodule_name, depth + 1); - } - return; - } let module = module_info.get(module_name).expect("Failed to get module"); let indent = " ".repeat(depth); let is_rust_module = module.name != "_"; if is_rust_module { - let rust_module_name = escape_keyword(module.name.clone()); - content - .write_str(&format!("{}pub mod {} {{\n", indent, rust_module_name)) - .expect("Failed to write string"); + let rust_module_name = escape_keyword(module_name.to_string()); + let opening = format!("{}pub mod {} {{\n", indent, rust_module_name); + content.write_str(&opening).expect("Failed to write string"); } content .write_str(&module.contents) .expect("Failed to write string"); - for submodule_name in module.submodules.iter() { + for submodule_name in module.submodules.keys() { write_module( content, - module_info, - [module_name, submodule_name].join(".").as_str(), + &module.submodules, + submodule_name.as_str(), depth + 1, ); } if is_rust_module { - content - .write_str(&format!("{}}}\n", indent)) - .expect("Failed to write string"); + let closing = format!("{}}}\n", indent); + content.write_str(&closing).expect("Failed to write string"); } } diff --git a/proto/prost/private/tests/shared_package_prefix/BUILD.bazel b/proto/prost/private/tests/shared_package_prefix/BUILD.bazel new file mode 100644 index 0000000000..6c5ddb1f6b --- /dev/null +++ b/proto/prost/private/tests/shared_package_prefix/BUILD.bazel @@ -0,0 +1,26 @@ +load("@rules_proto//proto:defs.bzl", "proto_library") +load("@rules_rust//rust:defs.bzl", "rust_test") +load("//proto/prost:defs.bzl", "rust_prost_library") + +proto_library( + name = "shared_pkg_prefix_proto", + srcs = [ + "pkg.foo.a.proto", + "pkg.foo.b.proto", + "pkg.c.proto", + ], +) + +rust_prost_library( + name = "shared_pkg_prefix_rs_proto", + proto = ":shared_pkg_prefix_proto", +) + +rust_test( + name = "shared_pkg_prefix_test", + srcs = ["shared_pkg_prefix_test.rs"], + edition = "2021", + deps = [ + ":shared_pkg_prefix_rs_proto", + ], +) diff --git a/proto/prost/private/tests/shared_package_prefix/pkg.c.proto b/proto/prost/private/tests/shared_package_prefix/pkg.c.proto new file mode 100644 index 0000000000..8872412415 --- /dev/null +++ b/proto/prost/private/tests/shared_package_prefix/pkg.c.proto @@ -0,0 +1,6 @@ +syntax = "proto3"; +package pkg; + +message CMessage { + string name = 1; +} diff --git a/proto/prost/private/tests/shared_package_prefix/pkg.foo.a.proto b/proto/prost/private/tests/shared_package_prefix/pkg.foo.a.proto new file mode 100644 index 0000000000..97a15bd792 --- /dev/null +++ b/proto/prost/private/tests/shared_package_prefix/pkg.foo.a.proto @@ -0,0 +1,9 @@ +syntax = "proto3"; +package pkg.foo; + +import "proto/prost/private/tests/shared_package_prefix/pkg.c.proto"; + +message AMessage { + string name = 1; + CMessage c = 2; +} diff --git a/proto/prost/private/tests/shared_package_prefix/pkg.foo.b.proto b/proto/prost/private/tests/shared_package_prefix/pkg.foo.b.proto new file mode 100644 index 0000000000..b7496301b8 --- /dev/null +++ b/proto/prost/private/tests/shared_package_prefix/pkg.foo.b.proto @@ -0,0 +1,6 @@ +syntax = "proto3"; +package pkg.foo; + +message BMessage { + string name = 1; +} diff --git a/proto/prost/private/tests/shared_package_prefix/shared_pkg_prefix_test.rs b/proto/prost/private/tests/shared_package_prefix/shared_pkg_prefix_test.rs new file mode 100644 index 0000000000..8501642a65 --- /dev/null +++ b/proto/prost/private/tests/shared_package_prefix/shared_pkg_prefix_test.rs @@ -0,0 +1,22 @@ +use shared_pkg_prefix_proto::pkg::foo::AMessage; +use shared_pkg_prefix_proto::pkg::foo::BMessage; +use shared_pkg_prefix_proto::pkg::CMessage; + +#[test] +fn test_packages() { + let pkg_b = BMessage { + name: "pkg_b".to_string(), + }; + let pkg_c = CMessage { + name: "pkg_c".to_string(), + }; + let pkg_a = AMessage { + name: "pkg_a".to_string(), + c: Some(pkg_c.clone()), + }; + + assert_eq!(pkg_a.name, "pkg_a"); + assert_eq!(pkg_a.c.unwrap().name, "pkg_c"); + assert_eq!(pkg_b.name, "pkg_b"); + assert_eq!(pkg_c.name, "pkg_c"); +}