roto/src/generator.rs

use crate::proto_gen::google::protobuf::descriptor::{
    DescriptorProto, EnumDescriptorProto, FieldDescriptorProto, FileDescriptorProto, FileDescriptorSet,
};
use crate::{ProtoAccessor, Result, RotoError};
use std::str;

pub 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 => ("u64".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()),
    }
}

pub fn generate_rust_code(set: &FileDescriptorSet) -> String {
    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;
            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 (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");

                let pascal_name = to_pascal_case(name);
                if num == 0 {
                    zero_variant_name = Some(pascal_name.clone());
                }
                output.push_str(&format!("    {} = {},\n", pascal_name, num));
                variant_count += 1;
            }

            if zero_variant_name.is_none() {
                output.push_str("    Unknown = 0,\n");
                zero_variant_name = Some("Unknown".to_string());
            }

            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 = 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", enum_name, zero_variant_name.as_ref().unwrap()));
            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 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.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",
                        safe_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",
                        safe_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 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.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",
                    safe_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"
            ));
        }
    }

    output
}