codas 0.7.1

//! Rust code generators.
//!
//! ## What's Here
//!
//! For a given coda, the following code will be
//! generated:
//!
//! - An `enum` containing variants for each data
//!   type documented by the coda.
//! - A `struct` for each data type.
//! - Codecs for the `enum` and every `struct`.
//!
//! The code generated by this implementation assumes
//! that the `codas` crate (with _at least_ `default`
//! features) is present wherever the code is compiled.
use alloc::{format, string::String, vec};

use crate::{
    codec::WritesEncodable,
    stream::{StreamError, Writes},
    types::{Coda, Text, Type, Unspecified},
};

/// Generates the rust types for `coda`,
/// writing them to `stream`.
///
/// Iff `with_serde`, the generated types
/// will be convertable to and from serde-equivalent types.
pub fn generate_types(
    coda: &Coda,
    stream: &mut impl Writes,
    with_serde: bool,
) -> Result<(), StreamError> {
    // Extract coda metadata.
    let coda_type_name = format!("{}Data", coda.local_name.trim());
    let coda_type_docs = match &coda.docs {
        Some(docs) => docs.as_ref(),
        None => "Undocumented Coda. How could you? ;~;",
    };

    // Escape double-quotes in docs.
    let coda_type_docs = coda_type_docs.replace('"', "\\\"");

    // Track the generated struct and variant fragments.
    let mut enum_variants = vec![];
    let mut enum_variant_ordinals_raw = vec![];
    let mut enum_variant_ordinals = vec![];
    let mut enum_variant_encoders = vec![];
    let mut enum_variant_header_encoders = vec![];
    let mut enum_variant_decoders = vec![];
    let mut enum_variant_converters = vec![];
    let mut type_structs = vec![];

    // Extract data types, which implicitly include
    // the [`Unspecified`] type.
    for (expected_ordinal, typing) in [Unspecified::DATA_TYPE]
        .iter()
        .chain(coda.iter())
        .enumerate()
    {
        // Validate ordinal.
        let type_ordinal = typing.format().as_data_format().ordinal;
        assert_eq!(expected_ordinal, type_ordinal as usize);
        enum_variant_ordinals_raw.push(type_ordinal);

        // Extract type metadata.
        let type_name = &typing.name;
        let type_docs = match &typing.docs {
            Some(docs) => docs.as_ref(),
            None => "Undocumented Type. How could you? ;~;",
        };

        // Escape double-quotes in docs.
        let type_docs = type_docs.replace('"', "\\\"");

        // The first data type in a coda is Unspecified;
        // this type is defined in the codas crate so
        // we don't re-generate it for every coda.
        let (struct_fqn, struct_name) = if type_ordinal == 0 {
            (
                "codas::types::Unspecified".into(),
                Text::from("Unspecified"),
            )
        } else {
            (format!("self::{type_name}"), type_name.clone())
        };

        // Extract field name-type pairings.
        let mut type_fields = vec![];
        for field in typing.iter() {
            let mut field_type = get_rust_type(&field.typing);
            if field.optional {
                field_type = format!("Option<{field_type}>").into();
            }

            type_fields.push((
                field.name.clone(),
                field_type,
                field.docs.clone(),
                field.flattened,
            ));
        }

        // Generate enum variant.
        let enum_variant = format!(
            r#"
            #[doc = "{type_docs}"]
            {struct_name}({struct_fqn})
        "#
        );
        enum_variants.push(enum_variant);

        // Generate enum variant ordinal.
        let enum_ordinal = format!("Self::{type_name}(..) => {type_ordinal}");
        enum_variant_ordinals.push(enum_ordinal);

        // Generate enum variant encoder.
        let enum_encoder = format!(
            r#"
            Self::{struct_name}(data) => {{
                data.encode(writer)
            }}
        "#
        );
        enum_variant_encoders.push(enum_encoder);

        // Generate enum variant header encoder.
        let enum_header_encoder = format!(
            r#"
            Self::{struct_name}(data) => {{
                data.encode_header(writer)
            }}
        "#
        );
        enum_variant_header_encoders.push(enum_header_encoder);

        // Generate enum variant decoder.
        let enum_decoder = format!(
            r#"
            {type_ordinal} => {{
                let mut data = {struct_fqn}::default();
                data.decode(reader, Some(header))?;
                *self = Self::{struct_name}(data);
                Ok(())
            }}
        "#
        );
        enum_variant_decoders.push(enum_decoder);

        // Generate enum variant converters.
        let enum_converter = format!(
            r#"
            impl From<{struct_fqn}> for {coda_type_name} {{
                fn from(data: {struct_fqn}) -> {coda_type_name} {{
                    {coda_type_name}::{struct_name}(data)
                }}
            }}

            impl codas::types::TryAsFormat<{struct_fqn}> for {coda_type_name} {{
                type Error = u8;

                fn try_as_format(&self) -> Result<&{struct_fqn}, Self::Error> {{
                    match self {{
                        {coda_type_name}::{struct_name}(data) => Ok(&data),
                        _ => Err(self.ordinal()),
                    }}
                }} 
            }}
        "#
        );
        enum_variant_converters.push(enum_converter);

        // If this type is Unspecified, don't generate
        // any struct or codec; it already exists.
        if type_ordinal == 0 {
            continue;
        }

        // Generate struct and codec.
        let mut type_struct = String::default();
        type_struct += &format!("#[doc = \"{type_docs}\"]\n");
        if with_serde {
            type_struct += "#[derive(serde::Serialize, serde::Deserialize)]\n";
        }
        type_struct += "#[derive(Default, Clone, Debug, PartialEq)]\n";
        type_struct += &format!("pub struct {struct_name} {{\n");
        for (name, typing, docs, flattened) in &type_fields {
            if let Some(docs) = docs {
                // Escape double-quotes in docs.
                let docs = docs.replace('"', "\\\"");
                type_struct += &format!("#[doc = \"{docs}\"]\n");
            }

            if *flattened && with_serde {
                type_struct += "#[serde(flatten)]\n";
            }

            type_struct += &format!("pub {name}: {typing},\n");
        }
        type_struct += "}";

        // Encoder impl.
        type_struct += &format!("impl codas::codec::Encodable for {struct_name} {{\n");

        // `FORMAT`
        type_struct += &format!(
            "const FORMAT: codas::codec::Format = codas::codec::Format::data({type_ordinal})"
        );
        for (_, typing, _, _) in &type_fields {
            type_struct += &format!("\n.with(<{typing} as codas::codec::Encodable>::FORMAT)");
        }
        type_struct += ";\n";

        // `fn encode`
        type_struct +=
            "fn encode(&self, writer: &mut (impl codas::codec::WritesEncodable + ?Sized),)\n";
        type_struct += "-> core::result::Result<(), codas::codec::CodecError> {\n";
        for (name, _, _, _) in &type_fields {
            type_struct += &format!("writer.write_data(&self.{name})?;\n");
        }
        type_struct += "Ok(())\n";
        type_struct += "}\n";
        type_struct += "}\n";

        // Decoder impl.
        type_struct += &format!("impl codas::codec::Decodable for {struct_name} {{\n");

        // `fn decode`
        type_struct += "fn decode(\n";
        type_struct += "&mut self,\n";
        type_struct += "reader: &mut (impl codas::codec::ReadsDecodable + ?Sized),\n";
        type_struct += "header: Option<codas::codec::DataHeader>,\n";
        type_struct += ") -> core::result::Result<(), codas::codec::CodecError> {\n";
        type_struct += &format!("let _ = Self::ensure_header(header, &[{type_ordinal}])?;\n");
        for (name, _, _, _) in &type_fields {
            type_struct += &format!("reader.read_data_into(&mut self.{name})?;\n");
        }
        type_struct += "Ok(())\n";
        type_struct += "}\n";
        type_struct += "}\n";

        type_structs.push(type_struct);
    }

    // Generate coda enum.
    let mut coda_enum = String::default();
    coda_enum += &format!("#[doc = \"{coda_type_docs}\"]\n");
    if with_serde {
        coda_enum += "#[derive(serde::Serialize, serde::Deserialize)]\n";
    }
    coda_enum += "#[derive(Clone, Debug, PartialEq)]\n";
    coda_enum += &format!("pub enum {coda_type_name} {{\n");

    // Enum variants.
    for variant in enum_variants {
        coda_enum += &variant;
        coda_enum += ",\n";
    }
    coda_enum += "}\n";

    // Enum impls.
    coda_enum += &format!("impl {coda_type_name} {{\n");

    // `CODA_BYTES`
    let mut coda_bytes = vec![];
    coda_bytes.write_data(coda).expect("codas are encodable");
    coda_enum += "/// The full Coda doc, in Coda-encoded bytes.\n";
    coda_enum += &format!(
        "const CODA_BYTES: &'static [u8; {}] = &{:?};\n",
        coda_bytes.len(),
        coda_bytes,
    );

    // `fn ordinal`
    coda_enum += "/// This variant's ordinal in the coda.\n";
    coda_enum += "pub fn ordinal(&self) -> u8 {\n";
    coda_enum += "match self {\n";
    for variant in enum_variant_ordinals {
        coda_enum += &variant;
        coda_enum += ",\n";
    }
    coda_enum += "}\n";
    coda_enum += "}\n";

    coda_enum += "}\n";

    // Enum encoder.
    coda_enum += &format!("impl codas::codec::Encodable for {coda_type_name} {{\n");

    // `FORMAT`
    coda_enum += "const FORMAT: codas::codec::Format = codas::codec::Format::Fluid;\n";

    // `fn encode`
    coda_enum += "fn encode(&self, writer: &mut (impl codas::codec::WritesEncodable + ?Sized),)\n";
    coda_enum += "-> core::result::Result<(), codas::codec::CodecError> {\n";
    coda_enum += "match self {\n";
    for variant in enum_variant_encoders {
        coda_enum += &variant;
        coda_enum += ",\n";
    }
    coda_enum += "}\n";
    coda_enum += "}\n";

    // `fn encoded_header`

    coda_enum +=
        "fn encode_header(&self, writer: &mut (impl codas::codec::WritesEncodable + ?Sized),)\n";
    coda_enum += "-> core::result::Result<(), codas::codec::CodecError> {\n";
    coda_enum += "match self {\n";
    for variant in enum_variant_header_encoders {
        coda_enum += &variant;
        coda_enum += ",\n";
    }
    coda_enum += "}\n";
    coda_enum += "}\n";

    coda_enum += "}\n";

    // Enum decoder.
    coda_enum += &format!("impl codas::codec::Decodable for {coda_type_name} {{\n");

    // `fn decode`
    coda_enum += "fn decode(\n";
    coda_enum += "&mut self,\n";
    coda_enum += "reader: &mut (impl codas::codec::ReadsDecodable + ?Sized),\n";
    coda_enum += "header: Option<codas::codec::DataHeader>,\n";
    coda_enum += ") -> core::result::Result<(), codas::codec::CodecError> {\n";

    // Ensure header.
    coda_enum += "let header = Self::ensure_header(header, &[\n";
    for ordinal in enum_variant_ordinals_raw {
        coda_enum += &format!("{ordinal},");
    }
    coda_enum += "])?;\n";

    // Decode variants.
    coda_enum += "match header.format.ordinal {\n";
    for variant in enum_variant_decoders {
        coda_enum += &variant;
        coda_enum += ",\n";
    }
    coda_enum += "_ => unreachable!(),\n";
    coda_enum += "}\n";
    coda_enum += "}\n";

    coda_enum += "}\n";

    // Enum default.
    coda_enum += &format!("impl core::default::Default for {coda_type_name} {{\n");
    coda_enum += &format!("fn default() -> {coda_type_name} {{\n");
    coda_enum += "Self::Unspecified(codas::types::Unspecified::default())\n";
    coda_enum += "}\n";
    coda_enum += "}\n";

    // Generate final code from the enum.
    let mut codegen = coda_enum;

    // Add type structs to the output.
    for variant in type_structs {
        codegen += &variant;
        codegen += "\n";
    }

    // Add struct-to-enum converters.
    for variant in enum_variant_converters {
        codegen += &variant;
        codegen += "\n";
    }

    stream.write_all(codegen.as_bytes())
}

/// Returns the native Rust identifier of `type`.
///
/// If `type` is a [`codas::spec::Type::Data`], the
/// data's name will be interpereted as a
/// native Rust identifier.
fn get_rust_type(typing: &Type) -> Text {
    match typing {
        Type::Unspecified => Text::Static("codas::types::Unspecified"),
        Type::U8 => Text::Static("u8"),
        Type::U16 => Text::Static("u16"),
        Type::U32 => Text::Static("u32"),
        Type::U64 => Text::Static("u64"),
        Type::I8 => Text::Static("i8"),
        Type::I16 => Text::Static("i16"),
        Type::I32 => Text::Static("i32"),
        Type::I64 => Text::Static("i64"),
        Type::F32 => Text::Static("f32"),
        Type::F64 => Text::Static("f64"),
        Type::Bool => Text::Static("bool"),
        Type::Text => Text::Static("codas::types::Text"),
        Type::Data(typing) => typing.name.clone(),
        Type::List(typing) => {
            let typing = get_rust_type(typing.as_ref());
            format!("alloc::vec::Vec<{typing}>").into()
        }
        Type::Map(typing) => {
            let key_typing = get_rust_type(&typing.as_ref().0);
            let value_typing = get_rust_type(&typing.as_ref().1);
            format!("alloc::collections::BTreeMap<{key_typing}, {value_typing}>").into()
        }
    }
}