lyra-engine/lyra-scripting/lyra-scripting-derive/src/mat_wrapper.rs

use proc_macro2::Ident;
use quote::quote;
use syn::{Path, Token};

pub(crate) struct MatWrapper {
    pub column_type: Ident,
}

impl MatWrapper {
    pub fn to_field_tokens(&self, wrapped_path: &Path, wrapper_ident: &Ident) -> proc_macro2::TokenStream {
        quote! {
            builder.field("ZERO", #wrapper_ident(#wrapped_path::ZERO));
            builder.field("IDENTITY", #wrapper_ident(#wrapped_path::IDENTITY));
            builder.field("NAN", #wrapper_ident(#wrapped_path::NAN));
        }
    }

    pub fn to_method_tokens(&self, wrapped_path: &Path, wrapper_ident: &Ident) -> proc_macro2::TokenStream {
        let column_type = &self.column_type;

        let column_size = {
            let ty_str = column_type.to_string();
            ty_str[ty_str.len() - 1..].parse::<usize>()
                .expect("Failure to parse number from token type")
        };
        let column_size_xtwo = column_size * 2;

        let element_ty = quote!(f32);

        quote! {
            builder.function("from_cols", 
                    |_, (x_axis, y_axis): (#column_type, #column_type)| {
                Ok(#wrapper_ident(#wrapped_path::from_cols(x_axis.0, y_axis.0)))
            });

            builder.function("from_cols_array",
                    |_, (arr,): ([#element_ty; #column_size_xtwo],)| {
                Ok(#wrapper_ident(#wrapped_path::from_cols_array(&arr)))
            });

            builder.function("from_cols_array_2d",
                    |_, (arr,): ([[#element_ty; #column_size]; #column_size],)| {
                Ok(#wrapper_ident(#wrapped_path::from_cols_array_2d(&arr)))
            });

            builder.function("from_diagonal", 
                    |_, (diag,): (#column_type,)| {
                Ok(#wrapper_ident(#wrapped_path::from_diagonal(diag.0)))
            });

            builder.method("col",
                    |_, this, (idx,): (usize,)| {
                Ok(#column_type(this.col(idx)))
            });

            builder.method("row",
                    |_, this, (idx,): (usize,)| {
                Ok(#column_type(this.row(idx)))
            });

            builder.method_mut("set_col",
                    |_, this, (idx, newc): (usize, #column_type)| {
                let col = this.col_mut(idx);
                *col = newc.0;

                Ok(())
            });

            builder.method("is_finite",
                    |_, this, (): ()| {
                Ok(this.is_finite())
            });

            builder.method("is_nan",
                    |_, this, (): ()| {
                Ok(this.is_nan())
            });

            builder.method("transpose",
                    |_, this, (): ()| {
                Ok(#wrapper_ident(this.0.transpose()))
            });

            builder.method("determinant",
                    |_, this, (): ()| {
                Ok(this.determinant())
            });

            builder.method("inverse",
                    |_, this, (): ()| {
                Ok(#wrapper_ident(this.inverse()))
            });

            builder.method("abs_diff_eq",
                    |_, this, (rhs, max_abs_diff): (#wrapper_ident, f32)| {
                Ok(this.abs_diff_eq(rhs.0, max_abs_diff))
            });

            // TODO: After all DMat's are implemented
            /* builder.method("as_dmat",
                    |_, this, (rhs, max_abs_diff): (#wrapper_ident, f32)| {
                Ok(D#wrapper_ident(this.as_dmat))
            }); */
        }
    }
}

impl syn::parse::Parse for MatWrapper {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let mut column_type = None;

        // cba to remove the use of this bool
        let mut first = true;

        while input.peek(Token![,]) || first {
            if !first {
                let _: Token![,] = input.parse()?;
            }

            if input.peek(syn::Ident) {
                let ident: Ident = input.parse()?;
                let ident_str = ident.to_string();
                let ident_str = ident_str.as_str();

                match ident_str {
                    "col_type" => {
                        let _eq: Token![=] = input.parse()?;
                        column_type = Some(input.parse()?);
                    },
                    _ => return Err(syn::Error::new_spanned(ident, "unknown matrix wrapper command")),
                }
            }

            first = false;
        }

        Ok(Self {
            column_type: column_type.ok_or_else(|| syn::Error::new(input.span(),
                "expected `col_type`"))?,
        })
    }
}