charles/roto
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add support for protobuf map fields

Browse Source 
Update the generator to detect map fields and use MapFieldIterator.
Implement MapFieldIterator in the runtime to handle key-value pair
extraction and add write_map_entry to ProtoBuilder. Add tests to
verify that map-bearing messages generate and compile correctly.
This commit is contained in:
charles
2026-05-07 20:15:08 -07:00
parent 8395195ac1
commit a20eed7223
6 changed files with 167 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codegen/src/generator.rs
+37 -15
View File
@@ -1,6 +1,6 @@
use crate::google::protobuf::descriptor::{
DescriptorProto, EnumDescriptorProto, FieldDescriptorProto, FileDescriptorProto,
FileDescriptorSet,
FileDescriptorSet, MessageOptions,
};
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
);
}
@@ -159,14 +164,23 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) {
let tag = field_proto.number().unwrap();
let f_type = field_proto.r#type().unwrap() as i32;
let f_label = field_proto.label().unwrap() as i32;
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));
fields_info.push((field_name.to_string(), tag, f_type, f_label, is_map));
}
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) in &fields_info {
for (field_name, _tag, _f_type, f_label, _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));
@@ -180,7 +194,7 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) {
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 {
for (name, _, _, label, _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));
@@ -192,7 +206,7 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) {
output.push_str(" for item in accessor.fields() {\n");
output.push_str(" let (offset, tag, _) = item?;\n");
for (name, tag, _, label) in &fields_info {
for (name, tag, _, label, _is_map) in &fields_info {
if *label == 3 {
output.push_str(&format!(" if tag.field_number == {} {{\n", tag));
output.push_str(&format!(
@@ -213,7 +227,7 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) {
output.push_str(" Ok(Self {\n");
output.push_str(" accessor,\n");
for (name, _, _, label) in &fields_info {
for (name, _, _, label, _is_map) in &fields_info {
if *label == 3 {
output.push_str(&format!("{}_start, {}_end,\n", name, name));
} else {
@@ -222,8 +236,8 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) {
}
output.push_str(" })\n }\n\n");
for (field_name, tag, f_type, f_label) in fields_info {
let (rust_type, logic) = map_type_to_rust_accessor(f_type, f_label);
for (field_name, tag, f_type, f_label, 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 {
@@ -239,11 +253,19 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) {
" 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!(