acadia/yunq/rust/src/codegen.rs

use crate::parser::{Decl, Interface, Message, Method, Type};
use crate::parser::{Field, FieldType};
use convert_case::{Case, Casing};
use proc_macro2::Ident;
use proc_macro2::Span;
use proc_macro2::TokenStream;
use quote::quote;

fn ident(s: &str) -> Ident {
    Ident::new(s, Span::call_site())
}

fn to_path(field_type: &FieldType) -> TokenStream {
    let rust_type = ident(&field_type.inner_type.rust_type());
    if field_type.repeated {
        quote! {
            Vec::<#rust_type>
        }
    } else {
        quote! {
            #rust_type
        }
    }
}

fn serialize_field(field: &Field) -> proc_macro2::TokenStream {
    let ind = field.number as usize;
    let name = ident(&field.name);
    if field.field_type.repeated {
        match &field.field_type.inner_type {
            Type::String => unimplemented!(),
            Type::Capability => unimplemented!(),
            Type::U64 => quote! {
                {
                    let rep_offset = next_extension;
                    let rep_len = self.#name.len() as u32;
                    next_extension = yunq::message::serialize_repeated(buf, offset + next_extension as usize, &self.#name)? as u32;

                    buf.write_at(yunq::message::field_offset(offset, #ind), rep_offset)?;
                    buf.write_at(yunq::message::field_offset(offset, #ind) + 4, rep_len)?;
                }
            },
            Type::I64 => unimplemented!(),
            Type::Bytes => unimplemented!(),
            Type::Message(_) => quote! {
                {
                    let rep_offset = next_extension;
                    let rep_len = self.#name.len() as u32;
                    next_extension = yunq::message::serialize_repeated_message(buf, offset + next_extension as usize, &self.#name, caps)? as u32;

                    buf.write_at(yunq::message::field_offset(offset, #ind), rep_offset)?;
                    buf.write_at(yunq::message::field_offset(offset, #ind) + 4, rep_len)?;
                }
            },
        }
    } else {
        match &field.field_type.inner_type {
            Type::String => quote! {
                {
                    let str_offset = next_extension;
                    let str_length = self.#name.len() as u32;

                    buf.write_str_at(offset + next_extension as usize, &self.#name)?;
                    next_extension += str_length;

                    buf.write_at(yunq::message::field_offset(offset, #ind), str_offset)?;
                    buf.write_at(yunq::message::field_offset(offset, #ind) + 4, str_length)?;
                }
            },
            Type::Capability => quote! {
                {
                    let cap_ind = caps.len();
                    caps.push(self.#name);

                    buf.write_at(yunq::message::field_offset(offset, #ind), cap_ind as u64)?;
                }
            },
            Type::U64 => quote! {
                {
                    buf.write_at(yunq::message::field_offset(offset, #ind), self.#name)?;
                }
            },
            Type::I64 => quote! {
                {
                    buf.write_at(yunq::message::field_offset(offset, #ind), self.#name)?;
                }
            },
            Type::Message(_) => quote! {
                {
                    let msg_offset = next_extension;
                    let msg_length = self.#name.serialize(buf, offset + next_extension as usize, caps)? as u32;
                    next_extension += msg_length;

                    buf.write_at(yunq::message::field_offset(offset, #ind), msg_offset)?;
                    buf.write_at(yunq::message::field_offset(offset, #ind) + 4, msg_length)?;
                }
            },
            Type::Bytes => unimplemented!(),
        }
    }
}

fn parse_field(field: &Field) -> TokenStream {
    let ind = field.number as usize;
    let name = ident(&field.name);
    if field.field_type.repeated {
        match &field.field_type.inner_type {
            Type::String => unimplemented!(),
            Type::Capability => unimplemented!(),
            Type::U64 => quote! {
                let #name = {
                    let rep_offset = buf.at::<u32>(yunq::message::field_offset(offset, #ind))?;
                    let rep_len = buf.at::<u32>(yunq::message::field_offset(offset, #ind) + 4)?;

                    yunq::message::parse_repeated(buf, offset + rep_offset as usize, rep_len as usize)?
                };
            },
            Type::I64 => unimplemented!(),
            Type::Bytes => unimplemented!(),
            Type::Message(_s) => quote! {
                let #name = {
                    let rep_offset = buf.at::<u32>(yunq::message::field_offset(offset, #ind))?;
                    let rep_len = buf.at::<u32>(yunq::message::field_offset(offset, #ind) + 4)?;

                    yunq::message::parse_repeated_message(buf, offset + rep_offset as usize, rep_len as usize, &caps)?
                };
            },
        }
    } else {
        match &field.field_type.inner_type {
            Type::String => quote! {
                let #name = {
                    let str_offset = buf.at::<u32>(yunq::message::field_offset(offset, #ind))?;
                    let str_length = buf.at::<u32>(yunq::message::field_offset(offset, #ind) + 4)?;

                    buf.str_at(offset + str_offset as usize, str_length as usize)?.to_string()
                };
            },
            Type::Capability => quote! {
                let #name = {
                    let cap_ind = buf.at::<u64>(yunq::message::field_offset(offset, #ind))?;
                    caps[cap_ind as usize]
                };
            },
            Type::U64 => quote! {
                let #name = buf.at::<u64>(yunq::message::field_offset(offset, #ind))?;
            },
            Type::I64 => quote! {
                let #name = buf.at::<i64>(yunq::message::field_offset(offset, #ind))?;
            },
            Type::Bytes => {
                unimplemented!();
            }
            Type::Message(s) => {
                let m_type = ident(s);
                quote! {
                    let #name = {
                        let msg_offset = buf.at::<u32>(yunq::message::field_offset(offset, #ind))? as usize;

                        #m_type::parse(buf, offset + msg_offset, caps)?
                    };
                }
            }
        }
    }
}

fn generate_serialize(message: &Message) -> TokenStream {
    let num_fields = message.fields.len();
    let serializers = message.fields.iter().map(serialize_field);
    quote! {
        #[allow(unused_variables)] // caps may be unused.
        fn serialize<const N: usize>(
            &self,
            buf: &mut yunq::ByteBuffer<N>,
            offset: usize,
            caps: &mut alloc::vec::Vec<z_cap_t>,
        ) -> Result<usize, ZError> {
            let num_fields = #num_fields;
            let core_size: u32 = (yunq::message::MESSAGE_HEADER_SIZE + 8 * num_fields) as u32;
            #[allow(unused_mut)]
            let mut next_extension = core_size;

            #(#serializers)*

            buf.write_at(offset + 0, yunq::message::MESSAGE_IDENT)?;
            buf.write_at(offset + 4, core_size)?;
            buf.write_at(offset + 8, next_extension)?;
            buf.write_at(offset + 12, 0 as u32)?;
            Ok(next_extension as usize)
        }
    }
}

fn generate_parse(message: &Message) -> TokenStream {
    let field_names = message.fields.iter().map(|field| ident(&field.name));
    let parsers = message.fields.iter().map(parse_field);
    quote! {
        #[allow(unused_variables)] // caps may be unused.
        fn parse<const N: usize>(
            buf: &yunq::ByteBuffer<N>,
            offset: usize,
            caps: &alloc::vec::Vec<z_cap_t>,
            ) -> Result<Self, ZError>
        where
            Self: Sized,
        {
            if buf.at::<u32>(offset + 0)? != yunq::message::MESSAGE_IDENT {
                mammoth::debug!("Expected IDENT at offest {:#x}, got {:#x}", offset, buf.at::<u32>(offset)?);
                return Err(ZError::INVALID_ARGUMENT);
            }
            // TODO: Parse core size.
            // TODO: Parse extension size.
            // TODO: Check CRC32
            // TODO: Parse options.

            #(#parsers)*

            Ok(Self {
                #(#field_names),*
            })

        }
    }
}

fn generate_message(message: &Message) -> TokenStream {
    let name = ident(&message.name);
    let field_names = message.fields.iter().map(|f| ident(&f.name));
    let field_types = message.fields.iter().map(|f| to_path(&f.field_type));

    let serialize = generate_serialize(message);
    let parse = generate_parse(message);
    quote! {
        #[derive(PartialEq, Eq, Clone, Debug)]
        pub struct #name {
            #(pub #field_names: #field_types),*
        }

        impl YunqMessage for #name {
            #serialize

            #parse
        }
    }
}

fn generate_method(method: &Method) -> TokenStream {
    let id = proc_macro2::Literal::u64_suffixed(method.number);
    let name = ident(&method.name.to_case(Case::Snake));
    let maybe_req = method.request.clone().map(|r| ident(&r));
    let maybe_resp = method.response.clone().map(|r| ident(&r));
    match (maybe_req, maybe_resp) {
        (Some(req), Some(resp)) => quote! {
            pub fn #name (&mut self, req: & #req) -> Result<#resp, ZError> {
                yunq::client::call_endpoint(#id, req, &mut self.byte_buffer, &self.endpoint_cap)
            }
        },
        (Some(req), None) => quote! {
            pub fn #name (&mut self, req: & #req) -> Result<yunq::message::Empty, ZError> {
                yunq::client::call_endpoint(#id, req, &mut self.byte_buffer, &self.endpoint_cap)
            }
        },
        (None, Some(resp)) => quote! {
            pub fn #name (&mut self) -> Result<#resp, ZError> {
                yunq::client::call_endpoint(#id, &yunq::message::Empty{}, &mut self.byte_buffer, &self.endpoint_cap)
            }
        },
        _ => unreachable!(),
    }
}

fn generate_client(interface: &Interface) -> TokenStream {
    let client_name = interface.name.clone() + "Client";
    let name = ident(&client_name);
    let methods = interface.methods.iter().map(generate_method);
    quote! {
        pub struct #name {
            endpoint_cap: Capability,
            byte_buffer: ByteBuffer<0x1000>,
        }

        impl #name {
            pub fn new(endpoint_cap: Capability) -> Self {
                Self {
                    endpoint_cap,
                    byte_buffer: ByteBuffer::new(),
                }
            }
            #(#methods)*
        }

    }
}

fn generate_server_case(method: &Method) -> TokenStream {
    let id = proc_macro2::Literal::u64_suffixed(method.number);
    let name = ident(&method.name.to_case(Case::Snake));
    let maybe_req = method.request.clone().map(|r| ident(&r));
    let maybe_resp = method.response.clone().map(|r| ident(&r));
    match (maybe_req, maybe_resp) {
        (Some(req), Some(_)) => quote! {
            #id => {
                let req = #req::parse_from_request(byte_buffer, cap_buffer)?;
                let resp = self.handler.#name(req)?;
                cap_buffer.resize(0, 0);
                let resp_len = resp.serialize_as_request(0, byte_buffer, cap_buffer)?;
                Ok(resp_len)
            },
        },
        (Some(req), None) => quote! {
            #id => {
                let req = #req::parse_from_request(byte_buffer, cap_buffer)?;
                self.handler.#name(req)?;
                cap_buffer.resize(0, 0);
                // TODO: Implement serialization for EmptyMessage so this is less hacky.
                yunq::message::serialize_error(byte_buffer, ZError::from(0));
                Ok(0x10)
            },
        },
        (None, Some(_)) => quote! {
            #id => {
                let resp = self.handler.#name()?;
                cap_buffer.resize(0, 0);
                let resp_len = resp.serialize_as_request(0, byte_buffer, cap_buffer)?;
                Ok(resp_len)
            },
        },
        _ => unreachable!(),
    }
}

fn generate_server_method(method: &Method) -> TokenStream {
    let name = ident(&method.name.to_case(Case::Snake));
    let maybe_req = method.request.clone().map(|r| ident(&r));
    let maybe_resp = method.response.clone().map(|r| ident(&r));
    match (maybe_req, maybe_resp) {
        (Some(req), Some(resp)) => quote! {
            fn #name (&mut self, req: #req) -> Result<#resp, ZError>;
        },
        (Some(req), None) => quote! {
            fn #name (&mut self, req: #req) -> Result<(), ZError>;
        },
        (None, Some(resp)) => quote! {
            fn #name (&mut self) -> Result<#resp, ZError>;
        },
        _ => unreachable!(),
    }
}

fn generate_server(interface: &Interface) -> TokenStream {
    let server_name = ident(&(interface.name.clone() + "Server"));
    let server_trait = ident(&(interface.name.clone() + "ServerHandler"));
    let server_trait_methods = interface.methods.iter().map(generate_server_method);
    let server_match_cases = interface.methods.iter().map(generate_server_case);
    quote! {
        pub trait #server_trait {
            #(#server_trait_methods)*
        }

        pub struct #server_name<T: #server_trait> {
            endpoint_cap: Capability,
            handler: T
        }

        impl<T: #server_trait> #server_name<T> {
            pub fn new(handler: T) -> Result<Self, ZError> {
                Ok(Self {
                    endpoint_cap: syscall::endpoint_create()?,
                    handler,
                })
            }
        }

        impl<T: #server_trait> yunq::server::YunqServer for #server_name<T> {
            fn endpoint_cap(&self) -> &Capability {
                &self.endpoint_cap
            }

            fn handle_request(
                &mut self,
                method_number: u64,
                byte_buffer: &mut ByteBuffer<1024>,
                cap_buffer: &mut Vec<z_cap_t>,
                ) -> Result<usize, ZError> {
                match method_number {
                    #(#server_match_cases)*

                    _ => Err(ZError::UNIMPLEMENTED)
                }
            }
        }
    }
}

fn generate_async_server_case(method: &Method) -> TokenStream {
    let id = proc_macro2::Literal::u64_suffixed(method.number);
    let name = ident(&method.name.to_case(Case::Snake));
    let maybe_req = method.request.clone().map(|r| ident(&r));
    let maybe_resp = method.response.clone().map(|r| ident(&r));
    match (maybe_req, maybe_resp) {
        (Some(req), Some(_)) => quote! {
            #id => {
                let req = #req::parse_from_request(byte_buffer, cap_buffer)?;
                let resp = self.handler.#name(req).await?;
                cap_buffer.resize(0, 0);
                let resp_len = resp.serialize_as_request(0, byte_buffer, cap_buffer)?;
                Ok(resp_len)
            },
        },
        (Some(req), None) => quote! {
            #id => {
                let req = #req::parse_from_request(byte_buffer, cap_buffer)?;
                self.handler.#name(req).await?;
                cap_buffer.resize(0, 0);
                // TODO: Implement serialization for EmptyMessage so this is less hacky.
                yunq::message::serialize_error(byte_buffer, ZError::from(0));
                Ok(0x10)
            },
        },
        (None, Some(_)) => quote! {
            #id => {
                let resp = self.handler.#name().await?;
                cap_buffer.resize(0, 0);
                let resp_len = resp.serialize_as_request(0, byte_buffer, cap_buffer)?;
                Ok(resp_len)
            },
        },
        _ => unreachable!(),
    }
}

fn generate_async_server_method(method: &Method) -> TokenStream {
    let name = ident(&method.name.to_case(Case::Snake));
    let maybe_req = method.request.clone().map(|r| ident(&r));
    let maybe_resp = method.response.clone().map(|r| ident(&r));
    match (maybe_req, maybe_resp) {
        (Some(req), Some(resp)) => quote! {
            fn #name (&self, req: #req) -> impl Future<Output= Result<#resp, ZError>> + Sync + Send;
        },
        (Some(req), None) => quote! {
            fn #name (&self, req: #req) -> impl Future<Output = Result<(), ZError>> + Sync + Send;
        },
        (None, Some(resp)) => quote! {
            fn #name (&self) -> impl Future<Output = Result<#resp, ZError>> + Sync + Send;
        },
        _ => unreachable!(),
    }
}

fn generate_async_server(interface: &Interface) -> TokenStream {
    let server_name = ident(&(String::from("Async") + &interface.name.clone() + "Server"));
    let server_trait = ident(&(String::from("Async") + &interface.name.clone() + "ServerHandler"));
    let server_trait_methods = interface.methods.iter().map(generate_async_server_method);
    let server_match_cases = interface.methods.iter().map(generate_async_server_case);
    quote! {
        pub trait #server_trait {
            #(#server_trait_methods)*
        }

        pub struct #server_name<T: #server_trait> {
            endpoint_cap: Capability,
            handler: T
        }

        impl<T: #server_trait> #server_name<T> {
            pub fn new(handler: T) -> Result<Self, ZError> {
                Ok(Self {
                    endpoint_cap: syscall::endpoint_create()?,
                    handler,
                })
            }
        }

        impl<T: #server_trait + Send + Sync + 'static> yunq::server::AsyncYunqServer for #server_name<T> {
            fn endpoint_cap(&self) -> &Capability {
                &self.endpoint_cap
            }

            async fn handle_request(
                &self,
                method_number: u64,
                byte_buffer: &mut ByteBuffer<1024>,
                cap_buffer: &mut Vec<z_cap_t>,
                ) -> Result<usize, ZError> {
                match method_number {
                    #(#server_match_cases)*

                    _ => Err(ZError::UNIMPLEMENTED)
                }
            }
        }
    }
}

fn generate_interface(interface: &Interface) -> TokenStream {
    let client = generate_client(interface);
    let server = generate_server(interface);
    let async_server = generate_async_server(interface);
    quote! {
        #client

        #server

        #async_server
    }
}

fn any_strings(ast: &[Decl]) -> bool {
    ast.iter()
        .filter_map(|decl| match decl {
            Decl::Message(m) => Some(m),
            _ => None,
        })
        .flat_map(|message| message.fields.iter())
        .any(|field| field.field_type.inner_type == Type::String)
}

fn any_interfaces(ast: &[Decl]) -> bool {
    ast.iter().any(|decl| matches!(decl, Decl::Interface(_)))
}

pub fn generate_code(ast: &[Decl]) -> String {
    let str_imports = if any_strings(ast) {
        quote! {
            use alloc::string::String;
            use alloc::string::ToString;
        }
    } else {
        quote! {}
    };

    let interface_imports = if any_interfaces(ast) {
        quote! {
            use core::future::Future;
            use mammoth::cap::Capability;
            use mammoth::syscall;
        }
    } else {
        quote! {}
    };

    let prelude = quote! {

    extern crate alloc;

    use alloc::vec::Vec;
    use mammoth::zion::z_cap_t;
    use mammoth::zion::ZError;
    use yunq::ByteBuffer;
    use yunq::YunqMessage;

    #str_imports

    #interface_imports

    };

    let message_decls = ast
        .iter()
        .filter_map(|decl| match decl {
            Decl::Message(m) => Some(m),
            _ => None,
        })
        .map(generate_message);

    let interface_decls = ast
        .iter()
        .filter_map(|decl| match decl {
            Decl::Interface(i) => Some(i),
            _ => None,
        })
        .map(generate_interface);

    let output = quote! {
        #prelude

        #(#message_decls)*

        #(#interface_decls)*
    }
    .to_string();

    let tokens = syn::parse_file(&output).unwrap_or_else(|e| {
        println!("{}", output);
        panic!("{}", e);
    });

    prettyplease::unparse(&tokens)
}