diff --git a/src/generator.rs b/src/generator.rs index bf1729d..968a205 100644 --- a/src/generator.rs +++ b/src/generator.rs @@ -1,9 +1,10 @@ -use crate::google::protobuf::descriptor::{ - DescriptorProto, EnumDescriptorProto, FileDescriptorProto, FieldDescriptorProto, FileDescriptorSet -}; use crate::ProtoAccessor; -use std::str; +use crate::google::protobuf::descriptor::{ + DescriptorProto, EnumDescriptorProto, FieldDescriptorProto, FileDescriptorProto, + FileDescriptorSet, +}; use std::collections::{HashMap, HashSet}; +use std::str; pub fn to_pascal_case(s: &str) -> String { s.split('_') @@ -91,7 +92,8 @@ fn write_enum(enum_proto: &EnumDescriptorProto, output: &mut String) { let mut zero_variant_name = None; while let Some(val_res) = values.next() { let (val_data, _) = val_res.expect("Failed to iterate enum"); - let accessor = ProtoAccessor::new(val_data).expect("Failed to parse EnumValueDescriptorProto"); + let accessor = + ProtoAccessor::new(val_data).expect("Failed to parse EnumValueDescriptorProto"); let (name_bytes, _) = accessor.get_value(1).expect("Enum value name missing"); let name = str::from_utf8(name_bytes).expect("Enum value name invalid utf8"); let (num_bytes, _) = accessor.get_value(2).expect("Enum value number missing"); @@ -119,16 +121,26 @@ fn write_enum(enum_proto: &EnumDescriptorProto, output: &mut String) { let mut values = enum_proto.value(); while let Some(val_res) = values.next() { let (val_data, _) = val_res.expect("Failed to read enum value"); - let accessor = ProtoAccessor::new(val_data).expect("Failed to parse EnumValueDescriptorProto"); + let accessor = + ProtoAccessor::new(val_data).expect("Failed to parse EnumValueDescriptorProto"); let (name_bytes, _) = accessor.get_value(1).expect("Enum value name missing"); let name = str::from_utf8(name_bytes).expect("Enum value name invalid utf8"); let (num_bytes, _) = accessor.get_value(2).expect("Enum value number missing"); let (num, _) = crate::read_varint(num_bytes).expect("Enum value number invalid varint"); - output.push_str(&format!(" {} => {}::{},\n", num, enum_name, to_pascal_case(name))); + output.push_str(&format!( + " {} => {}::{},\n", + num, + enum_name, + to_pascal_case(name) + )); } - output.push_str(&format!(" _ => {}::{},\n", enum_name, zero_variant_name.as_ref().unwrap())); + output.push_str(&format!( + " _ => {}::{},\n", + enum_name, + zero_variant_name.as_ref().unwrap() + )); output.push_str(" }\n }\n}\n\n"); } @@ -138,7 +150,8 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) { let mut fields_info = Vec::new(); for field_res in msg_proto.field() { let (field_data, _) = field_res.expect("Failed to iterate field"); - let field_proto = FieldDescriptorProto::new(field_data).expect("Failed to parse FieldDescriptorProto"); + let field_proto = + FieldDescriptorProto::new(field_data).expect("Failed to parse FieldDescriptorProto"); let field_name = field_proto.name().unwrap(); let tag = field_proto.number().unwrap(); @@ -148,13 +161,8 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) { fields_info.push((field_name.to_string(), tag, f_type, f_label)); } - output.push_str(&format!( - "pub struct {}<'a> {{\n", - msg_name - )); - if !fields_info.is_empty() { - output.push_str(" accessor: crate::ProtoAccessor<'a>,\n"); - } + output.push_str(&format!("pub struct {}<'a> {{\n", msg_name)); + output.push_str(" accessor: crate::ProtoAccessor<'a>,\n"); for (field_name, _tag, _f_type, f_label) in &fields_info { if *f_label == 3 { @@ -168,9 +176,8 @@ 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]) -> crate::Result {\n"); + output.push_str(" let accessor = crate::ProtoAccessor::new(data)?;\n"); if !fields_info.is_empty() { - output.push_str(" let accessor = crate::ProtoAccessor::new(data)?;\n"); - for (name, _, _, label) in &fields_info { if *label == 3 { output.push_str(&format!(" let mut {}_start = None;\n", name)); @@ -186,22 +193,24 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) { 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!( + " 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(&format!( + " if tag.field_number == {} {{ {}_offset = Some(offset); }}\n", + tag, name + )); } } output.push_str(" }\n\n"); - } else { - output.push_str(" let _ = crate::ProtoAccessor::new(data)?;\n"); } output.push_str(" Ok(Self {\n"); - if !fields_info.is_empty() { - output.push_str(" accessor,\n"); - } + output.push_str(" accessor,\n"); for (name, _, _, label) in &fields_info { if *label == 3 { output.push_str(&format!("{}_start, {}_end,\n", name, name)); @@ -213,17 +222,36 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) { for (field_name, tag, f_type, f_label) in fields_info { let (rust_type, logic) = map_type_to_rust_accessor(f_type, f_label); - let safe_name = if field_name == "type" { format!("r#{}", field_name) } else { field_name.clone() }; + let safe_name = if field_name == "type" { + format!("r#{}", field_name) + } else { + field_name.clone() + }; 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!( + " 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)); + output.push_str(&format!( + " _ => self.accessor.iter_repeated({}),\n", + tag + )); output.push_str(" }\n }\n\n"); } else { - output.push_str(&format!(" pub fn {}(&self) -> crate::Result<{}> {{\n", safe_name, rust_type)); - output.push_str(&format!(" let offset = self.{}_offset.ok_or(crate::RotoError::FieldNotFound)?;\n", field_name)); + output.push_str(&format!( + " pub fn {}(&self) -> crate::Result<{}> {{\n", + safe_name, rust_type + )); + output.push_str(&format!( + " let offset = self.{}_offset.ok_or(crate::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"); @@ -243,9 +271,14 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) { for field_res in msg_proto.field() { let (field_data, _) = field_res.expect("Failed to iterate field"); - let field_proto = FieldDescriptorProto::new(field_data).expect("Failed to parse FieldDescriptorProto"); + let field_proto = + FieldDescriptorProto::new(field_data).expect("Failed to parse FieldDescriptorProto"); let field_name = field_proto.name().unwrap(); - let safe_name = if field_name == "type" { format!("r#{}", field_name) } else { field_name.to_string() }; + let safe_name = if field_name == "type" { + format!("r#{}", field_name) + } else { + field_name.to_string() + }; let tag = field_proto.number().unwrap(); let f_type = field_proto.r#type().unwrap() as i32; let (rust_type, method) = map_type_to_rust_builder(f_type); @@ -258,21 +291,32 @@ fn write_message(msg_proto: &DescriptorProto, output: &mut String) { let mut nested_enums = Vec::new(); for e_res in msg_proto.enum_type() { - if let Ok((e, _)) = e_res { nested_enums.push(e); } + if let Ok((e, _)) = e_res { + nested_enums.push(e); + } } let mut nested_msgs = Vec::new(); for m_res in msg_proto.nested_type() { - if let Ok((m, _)) = m_res { nested_msgs.push(m); } + if let Ok((m, _)) = m_res { + nested_msgs.push(m); + } } if !nested_enums.is_empty() || !nested_msgs.is_empty() { let mod_name = to_snake_case(msg_proto.name().unwrap()); output.push_str(&format!("pub mod {} {{\n", mod_name)); for e_data in nested_enums { - write_enum(&EnumDescriptorProto::new(e_data).expect("Failed to parse nested EnumDescriptorProto"), output); + write_enum( + &EnumDescriptorProto::new(e_data) + .expect("Failed to parse nested EnumDescriptorProto"), + output, + ); } for m_data in nested_msgs { - write_message(&DescriptorProto::new(m_data).expect("Failed to parse nested DescriptorProto"), output); + write_message( + &DescriptorProto::new(m_data).expect("Failed to parse nested DescriptorProto"), + output, + ); } output.push_str("}\n\n"); } @@ -299,7 +343,8 @@ pub fn generate_rust_code( for file_res in set.file() { let (file_data, _) = file_res.expect("Failed to iterate file"); - let file_proto = FileDescriptorProto::new(file_data).expect("Failed to parse FileDescriptorProto"); + let file_proto = + FileDescriptorProto::new(file_data).expect("Failed to parse FileDescriptorProto"); let proto_name = file_proto.name().expect("File proto name missing"); if let Some(filter) = files_to_generate { @@ -311,6 +356,7 @@ pub fn generate_rust_code( let rust_file_name = format!("{}.rs", proto_name.replace(".proto", "")); let mut output = String::new(); + output.push_str("// @generated by protoc-gen-roto — do not edit\n\n"); output.push_str("use crate::{ProtoAccessor, ProtoBuilder, Result, RotoError, read_varint, RepeatedFieldIterator};\n"); output.push_str("use std::str;\n\n"); @@ -325,13 +371,19 @@ pub fn generate_rust_code( // Enums for enum_res in file_proto.enum_type() { let (enum_data, _) = enum_res.expect("Failed to iterate enum"); - write_enum(&EnumDescriptorProto::new(enum_data).expect("Failed to parse EnumDescriptorProto"), &mut output); + write_enum( + &EnumDescriptorProto::new(enum_data).expect("Failed to parse EnumDescriptorProto"), + &mut output, + ); } // Messages for msg_res in file_proto.message_type() { let (msg_data, _) = msg_res.expect("Failed to iterate message"); - write_message(&DescriptorProto::new(msg_data).expect("Failed to parse DescriptorProto"), &mut output); + write_message( + &DescriptorProto::new(msg_data).expect("Failed to parse DescriptorProto"), + &mut output, + ); } generated_files.push((rust_file_name, output)); } @@ -365,6 +417,7 @@ pub fn generate_rust_code( } let mut root_mod_content = String::new(); + root_mod_content.push_str("// @generated by protoc-gen-roto — do not edit\n\n"); let mut sorted_root_mods: Vec<_> = root_mods.into_iter().collect(); sorted_root_mods.sort(); for m in sorted_root_mods { @@ -374,6 +427,7 @@ pub fn generate_rust_code( for (mod_path, sub_mods) in mod_files { let mut content = String::new(); + content.push_str("// @generated by protoc-gen-roto — do not edit\n\n"); let mut sorted_subs: Vec<_> = sub_mods.into_iter().collect(); sorted_subs.sort(); for sub in sorted_subs {