fix: tests

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!(
+    output.push_str(&format!("pub struct {}<'a> {{\n", msg_name));
-        "pub struct {}<'a> {{\n",
+    output.push_str("    accessor: crate::ProtoAccessor<'a>,\n");
-        msg_name
-    ));
-    if !fields_info.is_empty() {
-        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<Self> {\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!(
-            output.push_str(&format!("        match (self.{}_start, self.{}_end) {{\n", field_name, field_name));
+                "    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!(
-            output.push_str(&format!("        let offset = self.{}_offset.ok_or(crate::RotoError::FieldNotFound)?;\n", field_name));
+                "    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 {