From 52df4b6908636f647ec846d461387caf5ded8781 Mon Sep 17 00:00:00 2001 From: charles Date: Thu, 25 Jun 2026 01:57:08 -0700 Subject: [PATCH] Resolve type paths and add typed iterators Build a type registry across all files to correctly handle nested and cross-referenced proto messages. Qualify generated Rust paths based on module depth. Introduce dedicated iterator structs for repeated message fields to replace raw byte accessors. --- codegen/src/generator/mod.rs | 206 +++++++++++++++++++++++++++++++++-- 1 file changed, 197 insertions(+), 9 deletions(-) diff --git a/codegen/src/generator/mod.rs b/codegen/src/generator/mod.rs index ea5de5b..920ba12 100644 --- a/codegen/src/generator/mod.rs +++ b/codegen/src/generator/mod.rs @@ -191,7 +191,52 @@ fn write_enum(enum_proto: &EnumDescriptorProto, output: &mut String) { output.push_str(" }\n }\n}\n\n"); } -fn write_message(msg_proto: &DescriptorProto, output: &mut String) { +/// Recursively build a mapping from proto type short name to generated Rust path. +/// Handles nested messages like DescriptorProto.ExtensionRange -> descriptor_proto::ExtensionRange +fn collect_type_paths( + msg_proto: &DescriptorProto, + parent_mod: &str, + registry: &mut HashMap, +) { + let msg_name = to_pascal_case(msg_proto.name().unwrap()); + let mod_name = to_snake_case(msg_proto.name().unwrap()); + + let rust_path = if parent_mod.is_empty() { + msg_name.clone() + } else { + format!("{}::{}", parent_mod, msg_name) + }; + + // Register this message under its short name + registry.insert(msg_name.clone(), rust_path); + + let child_mod = if parent_mod.is_empty() { + mod_name.clone() + } else { + format!("{}::{}", parent_mod, mod_name) + }; + + for m_res in msg_proto.nested_type() { + let (m_data, _) = m_res.expect("Failed to iterate nested message"); + let nested = DescriptorProto::new(m_data).expect("Failed to parse nested message"); + collect_type_paths(&nested, &child_mod, registry); + } +} + +pub fn write_message( + msg_proto: &DescriptorProto, + output: &mut String, + type_registry: &HashMap, + module_depth: usize, +) { + fn qualify_type(type_path: &str, depth: usize) -> String { + if depth == 0 { + type_path.to_string() + } else { + let prefix = (0..depth).map(|_| "super::").collect::(); + format!("{}{}", prefix, type_path) + } + } let msg_name = to_pascal_case(msg_proto.name().unwrap()); let mod_name = to_snake_case(msg_proto.name().unwrap()); @@ -216,6 +261,22 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) { }) .unwrap_or(false); + let type_name = if f_type == 11 { + field_proto.type_name().ok().and_then(|s| { + // Extract the short type name (last component) + let short = s + .split('.') + .filter(|p| !p.is_empty()) + .last() + .unwrap_or("") + .to_string(); + // Look up the full Rust path in the registry + type_registry.get(&short).cloned() + }) + } else { + None + }; + fields_info.push(( field_name.to_string(), tag, @@ -223,6 +284,7 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) { f_label, oneof_index, is_map, + type_name, )); } @@ -235,7 +297,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, _is_map) in &fields_info { + for (field_name, _tag, _f_type, f_label, _oneof_index, _is_map, _type_name) in &fields_info { if *f_label == 3 { output.push_str(&format!(" {}_start: Option,\n", field_name)); output.push_str(&format!(" {}_end: Option,\n", field_name)); @@ -248,7 +310,7 @@ 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 {\n"); output.push_str(" let accessor = roto_runtime::ProtoAccessor::new(data)?;\n"); - for (name, _, _, label, _oneof_index, _is_map) in &fields_info { + for (name, _, _, label, _oneof_index, _is_map, _type_name) 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)); @@ -259,7 +321,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, _oneof_index, _is_map) in &fields_info { + for (name, tag, _, label, _oneof_index, _is_map, _type_name) in &fields_info { if *label == 3 { output.push_str(&format!(" if tag.field_number == {} {{\n", tag)); output.push_str(&format!( @@ -279,7 +341,7 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) { output.push_str(" Ok(Self {\n"); output.push_str(" accessor,\n"); - for (name, _, _, label, _oneof_index, _is_map) in &fields_info { + for (name, _, _, label, _oneof_index, _is_map, _type_name) in &fields_info { if *label == 3 { output.push_str(&format!("{}_start, {}_end,\n", name, name)); } else { @@ -288,8 +350,10 @@ 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, is_map) in &fields_info { - let (rust_type, logic, default_val) = map_type_to_rust_accessor(*f_type, *f_label, *is_map); + // Collect typed iterator info for repeated message fields (generated after the impl block) + let mut typed_iters: Vec<(String, String, String)> = Vec::new(); + + for (field_name, tag, f_type, f_label, _oneof_index, is_map, type_name) in &fields_info { let safe_name = if field_name == "type" { format!("r#{}", field_name) } else { @@ -297,6 +361,8 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) { }; if *f_label == 3 { + // Repeated field — use the original accessor for the raw iterator + let (rust_type, _, _) = map_type_to_rust_accessor(*f_type, *f_label, *is_map); output.push_str(&format!( " pub fn {}(&self) -> {} {{\n", safe_name, rust_type @@ -319,7 +385,71 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) { )); } output.push_str(" }\n }\n\n"); + + // Generate typed iterator for repeated message fields + if *f_type == 11 { + if let Some(inner_type) = type_name { + let qualified = qualify_type(inner_type, module_depth); + // Use {ParentMsg}{Field}Iter to guarantee uniqueness + let pascal_field = to_pascal_case(field_name); + let iter_name = format!("{}{}Iter", msg_name, pascal_field); + // Schedule iterator struct for generation after impl block + typed_iters.push((iter_name, qualified, safe_name.clone())); + } + } + } else if *f_type == 11 { + // Typed singular message field + if let Some(inner_type) = type_name { + let qualified = qualify_type(inner_type, module_depth); + output.push_str(&format!( + " pub fn {}(&self) -> roto_runtime::Result<{}<'a>> {{\n", + safe_name, qualified + )); + output.push_str(&format!( + " let offset = self.{}_offset.ok_or(roto_runtime::RotoError::FieldNotFound)?;\n", + field_name + )); + output.push_str(" let (bytes, _) = self.accessor.get_value_at(offset)?;\n"); + output.push_str(&format!(" {}::new(bytes)\n", qualified)); + output.push_str(" }\n\n"); + + output.push_str(&format!( + " pub fn has_{}(&self) -> bool {{ self.{}_offset.is_some() }}\n\n", + field_name, field_name + )); + } else { + // Fallback — no type name, treat as raw bytes + let (rust_type, logic, default_val) = + map_type_to_rust_accessor(*f_type, *f_label, *is_map); + output.push_str(&format!( + " pub fn {}(&self) -> roto_runtime::Result<{}> {{\n", + safe_name, rust_type + )); + output.push_str(&format!( + " let offset = self.{}_offset.ok_or(roto_runtime::RotoError::FieldNotFound)?;\n", + field_name + )); + output.push_str(" let (bytes, _) = self.accessor.get_value_at(offset)?;\n"); + output.push_str(&format!(" {}\n", logic)); + output.push_str(" }\n\n"); + output.push_str(&format!( + " pub fn {}_or_default(&self) -> roto_runtime::Result<{}> {{\n", + safe_name, rust_type + )); + output.push_str(&format!( + " self.{}().or(Ok({}))\n", + safe_name, default_val + )); + output.push_str(" }\n\n"); + output.push_str(&format!( + " pub fn has_{}(&self) -> bool {{ self.{}_offset.is_some() }}\n\n", + field_name, field_name + )); + } } else { + // Non-message singular field — use the original accessor + let (rust_type, logic, default_val) = + map_type_to_rust_accessor(*f_type, *f_label, *is_map); output.push_str(&format!( " pub fn {}(&self) -> roto_runtime::Result<{}> {{\n", safe_name, rust_type @@ -362,7 +492,9 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) { snake_oneof_name, return_type ); output.push_str(&signature); - for (field_name, _tag, _f_type, _f_label, f_oneof_index, _is_map) in &fields_info { + for (field_name, _tag, _f_type, _f_label, f_oneof_index, _is_map, _type_name) in + &fields_info + { if *f_oneof_index == Some(oneof_index as i32) { let safe_field_name = if field_name == "type" { format!("r#{}", field_name) @@ -383,12 +515,47 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) { output.push_str(" Ok(None)\n }\n\n"); } + // Typed iterator accessors (for repeated message fields) + for (iter_name, _inner_type, safe_field) in &typed_iters { + output.push_str(&format!( + " pub fn {}_iter(&self) -> {}<'a> {{\n", + safe_field, iter_name + )); + output.push_str(&format!( + " {}::new(self.{}())\n", + iter_name, safe_field + )); + output.push_str(" }\n\n"); + } + // raw_fields() convenience on the message struct (before closing the impl) output.push_str(" pub fn raw_fields(&self) -> roto_runtime::RawFieldIterator<'a> {\n"); output.push_str(" self.accessor.raw_fields()\n"); output.push_str(" }\n\n"); output.push_str("}\n\n"); + // Generate typed iterator structs (outside the impl block) + for (iter_name, inner_type, _safe_field) in &typed_iters { + output.push_str(&format!("pub struct {}<'a> {{\n", iter_name)); + output.push_str(" inner: roto_runtime::RepeatedFieldIterator<'a>,\n"); + output.push_str("}\n\n"); + output.push_str(&format!("impl<'a> {}<'a> {{\n", iter_name)); + output.push_str(&format!( + " pub fn new(iter: roto_runtime::RepeatedFieldIterator<'a>) -> Self {{\n Self {{ inner: iter }}\n }}\n\n" + )); + output.push_str("}\n\n"); + output.push_str(&format!("impl<'a> Iterator for {}<'a> {{\n", iter_name)); + output.push_str(&format!( + " type Item = roto_runtime::Result<{}<'a>>;\n", + inner_type + )); + output.push_str(" fn next(&mut self) -> Option {\n"); + output.push_str(" self.inner.next().map(|item| {\n"); + output.push_str(" let (bytes, _) = item?;\n"); + output.push_str(&format!(" {}::new(bytes)\n", inner_type)); + output.push_str(" })\n }\n}\n\n"); + } + // Collect builder field info so we can use it multiple times below. // Tuple: (field_name, safe_name, tag, rust_type, write_method) let mut builder_fields: Vec<(String, String, u32, String, String)> = Vec::new(); @@ -569,6 +736,8 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) { write_message( &DescriptorProto::new(m_data).expect("Failed to parse nested DescriptorProto"), output, + type_registry, + module_depth + 1, ); } @@ -578,7 +747,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, _is_map) in &fields_info { + for (field_name, _tag, f_type, f_label, f_oneof_index, _is_map, _type_name) in + &fields_info + { if *f_oneof_index == Some(oneof_index as i32) { let (rust_type, _, _) = map_type_to_rust_accessor(*f_type, *f_label, *_is_map); let safe_field_name = if field_name == "type" { @@ -728,6 +899,21 @@ pub fn generate_protobuf_code( generate_mod_files, DATA_IMPORTS, |file_proto, output| { + // Build a global type registry from ALL files in the descriptor set, + // so cross-file type references (e.g., FeatureSet defined in one file, + // used in another) resolve correctly. + let mut type_registry: HashMap = HashMap::new(); + for file_res in set.file() { + let (file_data, _) = file_res.expect("Failed to iterate file"); + let fp = FileDescriptorProto::new(file_data).expect("Failed to parse file"); + for msg_res in fp.message_type() { + let (msg_data, _) = msg_res.expect("Failed to iterate message"); + let msg_proto = + DescriptorProto::new(msg_data).expect("Failed to parse message"); + collect_type_paths(&msg_proto, "", &mut type_registry); + } + } + // Enums for enum_res in file_proto.enum_type() { let (enum_data, _) = enum_res.expect("Failed to iterate enum"); @@ -744,6 +930,8 @@ pub fn generate_protobuf_code( write_message( &DescriptorProto::new(msg_data).expect("Failed to parse DescriptorProto"), output, + &type_registry, + 0, ); } },