Add EnumDescriptorProto and clap dependency

src/bin/generator.rs

@@ -0,0 +1,228 @@
+use clap::Parser;
+use roto::proto_gen::google::protobuf::descriptor::{
+    DescriptorProto, EnumDescriptorProto, FieldDescriptorProto, FileDescriptorProto, FileDescriptorSet,
+};
+use std::fs;
+use std::path::PathBuf;
+
+#[derive(Parser)]
+#[command(author, version, about = "Generates Rust accessor and builder code from a protobuf descriptor set")]
+struct Args {
+    /// Path to the descriptor set file (.desc)
+    #[arg(short, long)]
+    input: PathBuf,
+
+    /// Path to the output Rust file (.rs)
+    #[arg(short, long)]
+    output: PathBuf,
+}
+
+fn to_pascal_case(s: &str) -> String {
+    s.split('_')
+        .map(|word| {
+            let mut chars = word.chars();
+            match chars.next() {
+                None => String::new(),
+                Some(f) => f.to_uppercase().collect::<String>() + chars.as_str(),
+            }
+        })
+        .collect()
+}
+
+fn map_type_to_rust_accessor(field_type: i32, label: i32) -> (String, String) {
+    if label == 3 {
+        // LABEL_REPEATED
+        return (
+            "crate::RepeatedFieldIterator<'a>".to_string(),
+            "self.0.iter_repeated(%d)".to_string(),
+        );
+    }
+
+    match field_type {
+        9 => (
+            "&'a str".to_string(),
+            "str::from_utf8(bytes).map_err(|_| RotoError::WireFormatViolation)".to_string(),
+        ), // TYPE_STRING
+        1 => (
+            "f64".to_string(),
+            "f64::from_le_bytes(bytes.try_into().map_err(|_| RotoError::WireFormatViolation)?)".to_string(),
+        ), // TYPE_DOUBLE
+        2 => (
+            "f32".to_string(),
+            "f32::from_le_bytes(bytes.try_into().map_err(|_| RotoError::WireFormatViolation)?)".to_string(),
+        ), // TYPE_FLOAT
+        3 | 5 | 15 | 17 => (
+            "i32".to_string(),
+            "crate::read_varint(bytes).map(|(v, _)| v as i32).map_err(|_| RotoError::WireFormatViolation)".to_string(),
+        ), // INT/SINT/SFIXED 32
+        4 | 6 | 13 => (
+            "u32".to_string(),
+            "crate::read_varint(bytes).map(|(v, _)| v as u32).map_err(|_| RotoError::WireFormatViolation)".to_string(),
+        ), // UINT/FIXED 32
+        16 | 18 => (
+            "i64".to_string(),
+            "crate::read_varint(bytes).map(|(v, _)| v as i64).map_err(|_| RotoError::WireFormatViolation)".to_string(),
+        ), // SINT/SFIXED 64
+        7 | 14 => (
+            "u64".to_string(),
+            "crate::read_varint(bytes).map(|(v, _)| v as u64).map_err(|_| RotoError::WireFormatViolation)".to_string(),
+        ), // UINT/FIXED 64
+        8 => (
+            "bool".to_string(),
+            "crate::read_varint(bytes).map(|(v, _)| v != 0).map_err(|_| RotoError::WireFormatViolation)".to_string(),
+        ), // TYPE_BOOL
+        11 | 12 => ("&'a [u8]".to_string(), "Ok(bytes)".to_string()), // MESSAGE/BYTES
+        _ => ("&'a [u8]".to_string(), "Ok(bytes)".to_string()),
+    }
+}
+
+fn map_type_to_rust_builder(field_type: i32) -> (String, String) {
+    match field_type {
+        9 => ("&str".to_string(), "write_string".to_string()),
+        5 | 17 => ("i32".to_string(), "write_int32".to_string()),
+        3 | 4 | 8 | 13 | 14 | 18 => ("i64".to_string(), "write_varint".to_string()),
+        7 | 15 => ("u32".to_string(), "write_fixed32".to_string()),
+        6 | 16 => ("u64".to_string(), "write_fixed64".to_string()),
+        11 | 12 => ("&[u8]".to_string(), "write_bytes".to_string()),
+        _ => ("&[u8]".to_string(), "write_bytes".to_string()),
+    }
+}
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let args = Args::parse();
+    let data = fs::read(&args.input)?;
+    let set = FileDescriptorSet::new(&data).expect("Failed to parse FileDescriptorSet");
+
+    let mut output = String::new();
+    output.push_str("use crate::{ProtoAccessor, ProtoBuilder, Result, RotoError};\n");
+    output.push_str("use std::str;\n\n");
+
+    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");
+
+        // Enums
+        for enum_res in file_proto.enum_type() {
+            let (enum_data, _) = enum_res.expect("Failed to iterate enum");
+            let enum_proto = EnumDescriptorProto::new(enum_data).expect("Failed to parse EnumDescriptorProto");
+            let enum_name = to_pascal_case(enum_proto.name().unwrap());
+
+            output.push_str(&format!(
+                "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\n#[repr(i32)]\npub enum {} {{\n",
+                enum_name
+            ));
+
+            let mut values = enum_proto.enum_value();
+            let mut variant_count = 0;
+            while let Some(val_res) = values.next() {
+                let (val_data, _) = val_res.expect("Failed to read enum value");
+                let accessor = roto::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, _) = roto::read_varint(num_bytes).expect("Enum value number invalid varint");
+
+                output.push_str(&format!("    {}, = {},\n", to_pascal_case(name), num));
+                variant_count += 1;
+            }
+
+            if variant_count == 0 {
+                output.push_str("    Unknown = 0,\n");
+            }
+
+            output.push_str("}\n\n");
+
+            output.push_str(&format!(
+                "impl {} {{\n    pub fn from_i32(value: i32) -> Self {{\n        match value {{\n",
+                enum_name
+            ));
+
+            let mut values = enum_proto.enum_value();
+            while let Some(val_res) = values.next() {
+                let (val_data, _) = val_res.expect("Failed to read enum value");
+                let accessor = roto::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, _) = roto::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!("            _ => {}::Unknown,\n", enum_name));
+            output.push_str("        }\n    }}\n}\n\n");
+        }
+
+        // Messages
+        for msg_res in file_proto.message_type() {
+            let (msg_data, _) = msg_res.expect("Failed to iterate message");
+            let msg_proto = DescriptorProto::new(msg_data).expect("Failed to parse DescriptorProto");
+            let msg_name = to_pascal_case(msg_proto.name().unwrap());
+
+            // Accessor
+            output.push_str(&format!(
+                "pub struct {}<'a>(ProtoAccessor<'a>);\n\nimpl<'a> {}<'a> {{\n",
+                msg_name, msg_name
+            ));
+            output.push_str(&format!(
+                "    pub fn new(data: &'a [u8]) -> Result<Self> {{\n        Ok(Self(ProtoAccessor::new(data)?))\n    }}\n\n"
+            ));
+
+            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_name = field_proto.name().unwrap();
+                let tag = field_proto.number().unwrap();
+                let f_type = field_proto.field_type().unwrap() as i32;
+                let f_label = field_proto.label().unwrap() as i32;
+
+                let (rust_type, logic) = map_type_to_rust_accessor(f_type, f_label);
+
+                if f_label == 3 {
+                    output.push_str(&format!(
+                        "    pub fn {}(&self) -> {} {{\n        {}\n    }}\n\n",
+                        field_name, rust_type, logic.replace("%d", &tag.to_string())
+                    ));
+                } else {
+                    output.push_str(&format!(
+                        "    pub fn {}(&self) -> Result<{}> {{\n        let (bytes, _) = self.0.get_value({})?;\n        {}\n    }}\n\n",
+                        field_name, rust_type, tag, logic
+                    ));
+                }
+            }
+            output.push_str("}\n\n");
+
+            // Builder
+            output.push_str(&format!(
+                "pub struct {}Builder<'b> {{\n    builder: ProtoBuilder<'b>,\n}}\n\nimpl<'b> {}Builder<'b> {{\n",
+                msg_name, msg_name
+            ));
+            output.push_str(&format!(
+                "    pub fn builder(buf: &mut [u8]) -> {}Builder<'_> {{\n        {}Builder {{\n            builder: ProtoBuilder::new(buf),\n        }}\n    }}\n\n",
+                msg_name, msg_name
+            ));
+
+            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_name = field_proto.name().unwrap();
+                let tag = field_proto.number().unwrap();
+                let f_type = field_proto.field_type().unwrap() as i32;
+
+                let (rust_type, method) = map_type_to_rust_builder(f_type);
+
+                output.push_str(&format!(
+                    "    pub fn {}(mut self, value: {}) -> Result<Self> {{\n        self.builder.{}({}, value)?;\n        Ok(self)\n    }}\n\n",
+                    field_name, rust_type, method, tag
+                ));
+            }
+
+            output.push_str(&format!(
+                "    pub fn finish(self) -> Result<&'b mut [u8]> {{\n        self.builder.finish()\n    }}\n}}\n\n"
+            ));
+        }
+    }
+
+    fs::write(&args.output, output)?;
+    Ok(())
+}

src/proto_gen/google/protobuf/descriptor.rs

@@ -254,6 +254,67 @@ impl<'b> DescriptorProtoBuilder<'b> {
     }
 }
 
+pub struct EnumDescriptorProto<'a>(ProtoAccessor<'a>);
+
+impl<'a> EnumDescriptorProto<'a> {
+    pub fn new(data: &'a [u8]) -> Result<Self> {
+        Ok(Self(ProtoAccessor::new(data)?))
+    }
+
+    pub fn name(&self) -> Result<&'a str> {
+        let (bytes, _) = self.0.get_value(1)?;
+        str::from_utf8(bytes).map_err(|_| RotoError::WireFormatViolation)
+    }
+
+    pub fn enum_value(&self) -> crate::RepeatedFieldIterator<'a> {
+        self.0.iter_repeated(2)
+    }
+
+    pub fn reserved_value(&self) -> crate::RepeatedFieldIterator<'a> {
+        self.0.iter_repeated(3)
+    }
+
+    pub fn options(&self) -> crate::RepeatedFieldIterator<'a> {
+        self.0.iter_repeated(4)
+    }
+
+    pub fn builder(buf: &mut [u8]) -> EnumDescriptorProtoBuilder<'_> {
+        EnumDescriptorProtoBuilder {
+            builder: ProtoBuilder::new(buf),
+        }
+    }
+}
+
+pub struct EnumDescriptorProtoBuilder<'b> {
+    builder: ProtoBuilder<'b>,
+}
+
+impl<'b> EnumDescriptorProtoBuilder<'b> {
+    pub fn name(mut self, value: &str) -> Result<Self> {
+        self.builder.write_string(1, value)?;
+        Ok(self)
+    }
+
+    pub fn add_enum_value(mut self, data: &[u8]) -> Result<Self> {
+        self.builder.write_bytes(2, data)?;
+        Ok(self)
+    }
+
+    pub fn add_reserved_value(mut self, value: i32) -> Result<Self> {
+        self.builder.write_int32(3, value)?;
+        Ok(self)
+    }
+
+    pub fn add_options(mut self, data: &[u8]) -> Result<Self> {
+        self.builder.write_bytes(4, data)?;
+        Ok(self)
+    }
+
+    pub fn finish(self) -> Result<&'b mut [u8]> {
+        self.builder.finish()
+    }
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum FieldType {
     Double = 1,