Fix merge conflicts and generated code logic in generator.rs

This commit is contained in:
2026-05-07 20:53:08 -07:00
6 changed files with 201 additions and 54 deletions
+9
View File
@@ -0,0 +1,9 @@
«
codegen/data/test_map.proto roto.test"y
MapTest4
my_map ( 2.roto.test.MapTest.MyMapEntryRmyMap8
MyMapEntry
key ( Rkey
value (Rvalue:8bproto3
+7
View File
@@ -0,0 +1,7 @@
syntax = "proto3";
package roto.test;
message MapTest {
map<string, int32> my_map = 1;
}
+71 -53
View File
@@ -1,6 +1,6 @@
use crate::google::protobuf::descriptor::{
DescriptorProto, EnumDescriptorProto, FieldDescriptorProto, FileDescriptorProto,
FileDescriptorSet, OneofDescriptorProto,
FileDescriptorSet, MessageOptions, OneofDescriptorProto,
};
use roto_runtime::ProtoAccessor;
use std::collections::{HashMap, HashSet};
@@ -33,11 +33,16 @@ pub fn to_snake_case(s: &str) -> String {
result
}
fn map_type_to_rust_accessor(field_type: i32, label: i32) -> (String, String) {
fn map_type_to_rust_accessor(field_type: i32, label: i32, is_map: bool) -> (String, String) {
if label == 3 {
// LABEL_REPEATED
let iterator_type = if is_map {
"roto_runtime::MapFieldIterator<'a>"
} else {
"roto_runtime::RepeatedFieldIterator<'a>"
};
return (
"roto_runtime::RepeatedFieldIterator<'a>".to_string(),
iterator_type.to_string(),
"".to_string(), // Not used for repeated fields in the same way
);
}
@@ -160,8 +165,24 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) {
let f_type = field_proto.r#type().unwrap() as i32;
let f_label = field_proto.label().unwrap() as i32;
let oneof_index = field_proto.oneof_index().ok();
let is_map = field_proto
.options()
.map(|opt| {
MessageOptions::new(opt)
.unwrap()
.map_entry()
.unwrap_or(false)
})
.unwrap_or(false);
fields_info.push((field_name.to_string(), tag, f_type, f_label, oneof_index));
fields_info.push((
field_name.to_string(),
tag,
f_type,
f_label,
oneof_index,
is_map,
));
}
let mut oneofs = Vec::new();
@@ -173,7 +194,7 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) {
output.push_str(&format!("pub struct {}<'a> {{\n", msg_name));
output.push_str(" accessor: roto_runtime::ProtoAccessor<'a>,\n");
for (field_name, _tag, _f_type, f_label, _oneof_index) in &fields_info {
for (field_name, _tag, _f_type, f_label, _oneof_index, _is_map) in &fields_info {
if *f_label == 3 {
output.push_str(&format!(" {}_start: Option<usize>,\n", field_name));
output.push_str(&format!(" {}_end: Option<usize>,\n", field_name));
@@ -186,41 +207,38 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) {
output.push_str(&format!("impl<'a> {}<'a> {{\n", msg_name));
output.push_str(" pub fn new(data: &'a [u8]) -> roto_runtime::Result<Self> {\n");
output.push_str(" let accessor = roto_runtime::ProtoAccessor::new(data)?;\n");
if !fields_info.is_empty() {
for (name, _, _, label, _) in &fields_info {
if *label == 3 {
output.push_str(&format!(" let mut {}_start = None;\n", name));
output.push_str(&format!(" let mut {}_end = None;\n", name));
} else {
output.push_str(&format!(" let mut {}_offset = None;\n", name));
}
for (name, _, _, label, _oneof_index, _is_map) in &fields_info {
if *label == 3 {
output.push_str(&format!(" let mut {}_start = None;\n", name));
output.push_str(&format!(" let mut {}_end = None;\n", name));
} else {
output.push_str(&format!(" let mut {}_offset = None;\n", name));
}
output.push_str(" for item in accessor.fields() {\n");
output.push_str(" let (offset, tag, _) = item?;\n");
for (name, tag, _, label, _) in &fields_info {
if *label == 3 {
output.push_str(&format!(" if tag.field_number == {} {{\n", tag));
output.push_str(&format!(
" if {}_start.is_none() {{ {}_start = Some(offset); }}\n",
name, name
));
output.push_str(&format!(" {}_end = Some(offset);\n", name));
output.push_str(" }\n");
} else {
output.push_str(&format!(
" if tag.field_number == {} {{ {}_offset = Some(offset); }}\n",
tag, name
));
}
}
output.push_str(" }\n\n");
}
output.push_str(" for item in accessor.fields() {\n");
output.push_str(" let (offset, tag, _) = item?;\n");
for (name, tag, _, label, _oneof_index, _is_map) in &fields_info {
if *label == 3 {
output.push_str(&format!(" if tag.field_number == {} {{\n", tag));
output.push_str(&format!(
" if {}_start.is_none() {{ {}_start = Some(offset); }}\n",
name, name
));
output.push_str(&format!(" {}_end = Some(offset);\n", name));
output.push_str(" }\n");
} else {
output.push_str(&format!(
" if tag.field_number == {} {{ {}_offset = Some(offset); }}\n",
tag, name
));
}
}
output.push_str(" }\n\n");
output.push_str(" Ok(Self {\n");
output.push_str(" accessor,\n");
for (name, _, _, label, _) in &fields_info {
for (name, _, _, label, _oneof_index, _is_map) in &fields_info {
if *label == 3 {
output.push_str(&format!("{}_start, {}_end,\n", name, name));
} else {
@@ -229,8 +247,8 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) {
}
output.push_str(" })\n }\n\n");
for (field_name, tag, f_type, f_label, _oneof_index) in &fields_info {
let (rust_type, logic) = map_type_to_rust_accessor(*f_type, *f_label);
for (field_name, tag, f_type, f_label, _oneof_index, is_map) in &fields_info {
let (rust_type, logic) = map_type_to_rust_accessor(*f_type, *f_label, *is_map);
let safe_name = if field_name == "type" {
format!("r#{}", field_name)
} else {
@@ -238,19 +256,19 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) {
};
if *f_label == 3 {
output.push_str(&format!(
" pub fn {}(&self) -> {} {{\n",
safe_name, rust_type
));
output.push_str(&format!(
" match (self.{}_start, self.{}_end) {{\n",
field_name, field_name
));
output.push_str(&format!(" (Some(start), Some(end)) => self.accessor.iter_repeated_range({}, start, end),\n", tag));
output.push_str(&format!(
" _ => self.accessor.iter_repeated({}),\n",
tag
));
if *is_map {
output.push_str(&format!(" (Some(start), Some(end)) => roto_runtime::MapFieldIterator::new(self.accessor.iter_repeated_range({}, start, end)),\n", tag));
output.push_str(&format!(
" _ => roto_runtime::MapFieldIterator::new(self.accessor.iter_repeated({})),\n",
tag
));
} else {
output.push_str(&format!(" (Some(start), Some(end)) => self.accessor.iter_repeated_range({}, start, end),\n", tag));
output.push_str(&format!(
" _ => self.accessor.iter_repeated({}),\n",
tag
));
}
output.push_str(" }\n }\n\n");
} else {
output.push_str(&format!(
@@ -282,7 +300,7 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) {
" pub fn which_{}(&self) -> roto_runtime::Result<Option<{}::{}<'a>>> {{\n",
snake_oneof_name, msg_name, pascal_oneof_name
));
for (field_name, _tag, _f_type, _f_label, f_oneof_index) in &fields_info {
for (field_name, _tag, _f_type, _f_label, f_oneof_index, _is_map) in &fields_info {
if *f_oneof_index == Some(oneof_index as i32) {
let safe_field_name = if field_name == "type" {
format!("r#{}", field_name)
@@ -416,9 +434,9 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) {
let oneof_name = oneof_desc.name().unwrap();
let pascal_oneof_name = to_pascal_case(oneof_name);
output.push_str(&format!("pub enum {}<'a> {{\n", pascal_oneof_name));
for (field_name, _tag, f_type, f_label, f_oneof_index) in &fields_info {
for (field_name, _tag, f_type, f_label, f_oneof_index, _is_map) in &fields_info {
if *f_oneof_index == Some(oneof_index as i32) {
let (rust_type, _) = map_type_to_rust_accessor(*f_type, *f_label);
let (rust_type, _) = map_type_to_rust_accessor(*f_type, *f_label, *_is_map);
let safe_field_name = if field_name == "type" {
format!("r#{}", field_name)
} else {
+67
View File
@@ -0,0 +1,67 @@
use roto_codegen::google::protobuf::descriptor::FileDescriptorSet;
use std::fs;
use std::process::Command;
#[test]
fn test_map_generated_code_builds() {
// 1. Load FileDescriptorSet from data/test_map.desc
let desc_path = "data/test_map.desc";
let data = fs::read(desc_path).expect("Failed to read test_map.desc");
let set = FileDescriptorSet::new(&data)
.expect("Failed to create FileDescriptorSet from test_map.desc");
let generated_files = roto_codegen::generator::generate_rust_code(&set, None, false);
assert!(
!generated_files.is_empty(),
"Generated code should not be empty"
);
// 2. Setup a temporary Cargo project to verify the code builds
let root = std::env::current_dir().expect("Failed to get current directory");
let temp_project_dir = root.join("test_map_gen_project");
// Clean up previous runs
if temp_project_dir.exists() {
fs::remove_dir_all(&temp_project_dir).expect("Failed to clean up temp project directory");
}
// Create new library project
let status = Command::new("cargo")
.args(["new", "--lib", "test_map_gen_project"])
.current_dir(&root)
.status()
.expect("Failed to run cargo new");
assert!(status.success(), "cargo new failed");
// 3. Configure the project to depend on the current roto crate
let cargo_toml_path = temp_project_dir.join("Cargo.toml");
let cargo_toml_content =
fs::read_to_string(&cargo_toml_path).expect("Failed to read Cargo.toml");
let updated_cargo_toml = format!(
"{}\n\nroto-codegen = {{ path = \"..\" }}\nroto-runtime = {{ path = \"../../runtime\" }}\n\n[workspace]\n",
cargo_toml_content
);
fs::write(cargo_toml_path, updated_cargo_toml).expect("Failed to write Cargo.toml");
// 4. Write the generated code to src/lib.rs
let mut all_code = String::new();
for (_, content) in generated_files {
all_code.push_str(&content);
all_code.push_str("\n");
}
let final_code = all_code.replace("use crate::", "use roto::");
let lib_path = temp_project_dir.join("src/lib.rs");
fs::write(lib_path, final_code).expect("Failed to write generated code to src/lib.rs");
// 5. Attempt to build the project
let build_status = Command::new("cargo")
.args(["--offline", "build"])
.current_dir(&temp_project_dir)
.status()
.expect("Failed to run cargo build");
assert!(
build_status.success(),
"The generated Rust code for test_map.proto failed to build in a standalone project!"
);
}