magnetar/magnetar_sdk/macros/src/lib.rs

use proc_macro::TokenStream;
use quote::ToTokens;
use std::collections::HashSet;
use syn::parse::Parse;
use syn::punctuated::Punctuated;
use syn::{
    Expr, ExprLit, ExprPath, GenericArgument, Ident, Lit, Meta, MetaNameValue, PathArguments,
    Token, Type, TypePath,
};

struct Field {
    name: Ident,
    ty: Type,
}

struct PackInput {
    struct_name: syn::Ident,
    fields: Vec<Field>,
}

impl Parse for PackInput {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let struct_name = input.parse()?;
        let _ = input.parse::<Token![,]>()?;

        let mut fields = Vec::new();

        loop {
            let ty = input.parse()?;
            let _alias_keyword = input.parse::<Token![as]>()?;
            let name = input.parse()?;

            fields.push(Field { name, ty });

            if input.is_empty() {
                break;
            }

            let _ = input.parse::<Token![&]>()?;
        }

        Ok(PackInput {
            struct_name,
            fields,
        })
    }
}

#[proc_macro]
pub fn pack(item: TokenStream) -> TokenStream {
    let parsed = syn::parse_macro_input!(item as PackInput);
    let struct_name = &parsed.struct_name;
    let fields = &parsed.fields;

    let names = fields.iter().map(|f| &f.name);
    let types_decl = fields.iter().map(|f| &f.ty);
    let types_struct = fields.iter().map(|f| &f.ty);
    let types_deps = fields.iter().map(|f| &f.ty);
    let types_packed = fields.iter().map(|f| &f.ty);

    let names_packed = fields.iter().map(|f| &f.name);
    let iota = (0..fields.len()).map(syn::Index::from);

    let export_path = format!("bindings/packed/{}.ts", struct_name);

    let tuple = quote::quote! {
        (#(#types_packed,)*)
    };

    let name = struct_name.to_string();

    let export_name = Ident::new(
        &format!("export_bindings_{}", struct_name.to_string().to_lowercase()),
        struct_name.span(),
    );

    let expanded = quote::quote! {
        #[derive(Clone, Debug, Deserialize, Serialize)]
        pub struct #struct_name {
            #(
            #[serde(flatten)]
            pub #names: #types_struct
            ),*
        }

        impl Packed for #struct_name {
            type Input = #tuple;

            fn pack_from(from: #tuple) -> Self {
                #struct_name {
                    #(
                    #names_packed: from.#iota
                    ),*
                }
            }
        }

        #[cfg(test)]
        #[test]
        fn #export_name() {
            #struct_name::export().expect("could not export type");
        }

        impl TS for #struct_name {
            const EXPORT_TO: Option<&'static str> = Some(#export_path);

            fn decl() -> String {
                format!(
                    "type {} = {};",
                    Self::name(),
                    vec![#(<#types_decl as TS>::inline()),*].join(" & ")
                )
            }

            fn name() -> String {
                #name.to_string()
            }

            fn dependencies() -> Vec<ts_rs::Dependency> {
                vec![
                    #(
                    <#types_deps as TS>::dependencies()
                    ),*
                ].into_iter().flatten().collect::<Vec<_>>()
            }

            fn transparent() -> bool {
                false
            }
        }
    };

    TokenStream::from(expanded)
}

#[proc_macro_derive(Endpoint, attributes(name, endpoint, method, request, response))]
pub fn derive_endpoint(item: TokenStream) -> TokenStream {
    let input = syn::parse_macro_input!(item as syn::DeriveInput);

    let struct_name = input.ident;
    assert!(input.generics.params.is_empty(), "generics not supported");
    let where_clause = input.generics.where_clause;
    assert!(where_clause.is_none(), "where clauses not supported");
    let data = match input.data {
        syn::Data::Struct(data) => data,
        _ => panic!("only unit structs are supported"),
    };
    assert!(
        matches!(data.fields, syn::Fields::Unit),
        "only unit structs are supported"
    );

    let mut name = struct_name.to_string();
    let mut endpoint = None;
    let mut method = None;
    let mut request = None;
    let mut response = None;
    let mut request_args = None;
    let mut response_args = None;

    let mut found = HashSet::new();

    let struct_attrs = input.attrs;

    for struct_attr in struct_attrs {
        if !struct_attr.meta.path().is_ident("endpoint") {
            continue;
        }

        if struct_attr.style != syn::AttrStyle::Outer {
            panic!("expected outer attribute");
        }

        let Meta::List(list) = struct_attr.meta else {
            panic!("expected a list of attributes");
        };

        let nested = list
            .parse_args_with(Punctuated::<MetaNameValue, Token![,]>::parse_terminated)
            .expect("expected a list of attributes");

        for kvp in nested {
            let key = kvp
                .path
                .get_ident()
                .expect("expected a single identifier")
                .to_owned();
            let value = kvp.value;

            if found.contains(&key) {
                panic!("duplicate attribute: {}", key);
            }

            match key.to_string().as_ref() {
                "name" => {
                    let Expr::Lit(ExprLit {
                        lit: Lit::Str(new_name),
                        ..
                    }) = value
                    else {
                        panic!("expected a string literal");
                    };

                    name = new_name.value();
                }
                "endpoint" => {
                    let Expr::Lit(ExprLit {
                        lit: Lit::Str(endpoint_val),
                        ..
                    }) = value
                    else {
                        panic!("expected a string literal");
                    };

                    endpoint = Some(endpoint_val.value());
                }
                "method" => {
                    let Expr::Path(ExprPath { path, .. }) = value else {
                        panic!("expected a an identifier");
                    };

                    method = Some(path);
                }
                "request" => {
                    let Expr::Lit(ExprLit {
                        lit: Lit::Str(str_lit),
                        ..
                    }) = value
                    else {
                        panic!("expected a an identifier");
                    };

                    let type_tok: TokenStream = str_lit.value().parse().unwrap();
                    let req_typ = syn::parse::<TypePath>(type_tok).unwrap();

                    if let Some(seg_last) = req_typ.path.segments.last() {
                        if let PathArguments::AngleBracketed(args) = &seg_last.arguments {
                            let args_res = args.args.iter().cloned().collect::<Vec<_>>();
                            request_args = Some(args_res);
                        }
                    }

                    request = Some(req_typ);
                }
                "response" => {
                    let Expr::Lit(ExprLit {
                        lit: Lit::Str(str_lit),
                        ..
                    }) = value
                    else {
                        panic!("expected a an identifier");
                    };

                    let type_tok: TokenStream = str_lit.value().parse().unwrap();
                    let res_typ = syn::parse::<TypePath>(type_tok).unwrap();
                    if let Some(seg_last) = res_typ.path.segments.last() {
                        if let PathArguments::AngleBracketed(args) = &seg_last.arguments {
                            let args_res = args.args.iter().cloned().collect::<Vec<_>>();
                            response_args = Some(args_res);
                        }
                    }
                    response = Some(res_typ);
                }
                _ => panic!("unknown attribute: {}", key),
            }

            found.insert(key);
        }
    }

    let endpoint = endpoint.expect("missing endpoint attribute");
    let method = method.expect("missing method attribute");
    let request = request.expect("missing request attribute");
    let response = response.expect("missing response attribute");

    let mut endpoint_chars = endpoint.chars();
    let mut path_params = String::from("[");
    while let Some(c) = endpoint_chars.next() {
        if c == ':' {
            if !path_params.ends_with('[') {
                path_params += ", ";
            }

            path_params += "\"";
            path_params += &endpoint_chars
                .clone()
                .take_while(|c| c.is_ascii_alphanumeric() || *c == '_')
                .collect::<String>();
            path_params += "\"";
        }
    }
    path_params += "]";

    let ts_path = format!("bindings/endpoints/{}.ts", name);
    let export_name = Ident::new(
        &format!("export_bindings_{}", struct_name.to_string().to_lowercase()),
        struct_name.span(),
    );

    let call_name_req = if let Some(ref args) = request_args {
        let args_str = args
            .iter()
            .map(|a| a.into_token_stream().to_string())
            .collect::<Vec<_>>();

        quote::quote! {
            <#struct_name as Endpoint>::Request::name_with_type_args(vec![#( #args_str.to_string() ),*])
        }
    } else {
        quote::quote! {
            <#struct_name as Endpoint>::Request::name()
        }
    };

    let req_args_flat = request_args.clone().unwrap_or_default();
    let mut res_args_flat = response_args.clone().unwrap_or_default();

    let call_name_res = if let Some(ref args) = response_args {
        let args_nested = args
            .iter()
            .map(|a| {
                let normal = Vec::new();

                let GenericArgument::Type(inner) = a else {
                    return normal;
                };
                let Type::Path(path_inner) = inner else {
                    return normal;
                };
                let Some(seg) = path_inner.path.segments.last() else {
                    return normal;
                };
                let PathArguments::AngleBracketed(angle_inner) = &seg.arguments else {
                    return normal;
                };
                angle_inner.args.iter().cloned().collect::<Vec<_>>()
            })
            .collect::<Vec<_>>();

        res_args_flat.extend(args_nested.iter().flatten().cloned());

        let arg_tokens_nested = args_nested
            .iter()
            .zip(args)
            .map(|(args, parent)| {
                if !args.is_empty() {
                    let names = args.iter().map(|a| a.to_token_stream().to_string());
                    quote::quote! {
                        <#parent as TS>::name_with_type_args(vec![#( #names.to_string() ),*])
                    }
                } else {
                    quote::quote! {
                        <#parent as TS>::name()
                    }
                }
            })
            .collect::<Vec<_>>();

        quote::quote! {
            <#struct_name as Endpoint>::Response::name_with_type_args(vec![#( #arg_tokens_nested ),*])
        }
    } else {
        quote::quote! {
            <#struct_name as Endpoint>::Response::name()
        }
    };

    let expanded = quote::quote! {
        impl Default for #struct_name {
            fn default() -> Self {
                Self
            }
        }

        impl Endpoint for #struct_name {
            const NAME: &'static str = #name;
            const ENDPOINT: &'static str = #endpoint;
            const METHOD: Method = #method;

            type Request = #request;
            type Response = #response;
        }

        #[cfg(test)]
        #[test]
        fn #export_name() {
            #struct_name::export().expect("could not export type");
        }

        impl TS for #struct_name {
            const EXPORT_TO: Option<&'static str> = Some(#ts_path);

            fn decl() -> String {
                format!(
                    "const {} = {{\n    \
                        endpoint: \"{}\",\n    \
                        pathParams: {} as {},\n    \
                        method: \"{}\" as \"GET\" | \"POST\" | \"PUT\" | \"DELETE\" | \"PATCH\",\n    \
                        request: undefined as unknown as {},\n    \
                        response: undefined as unknown as {}\n\
                    }}
                    ",
                    Self::name(),
                    Self::ENDPOINT,
                    #path_params,
                    #path_params,
                    Self::METHOD,
                    #call_name_req,
                    #call_name_res
                )
            }

            fn name() -> String {
                #struct_name::NAME.to_string()
            }

            fn dependencies() -> Vec<ts_rs::Dependency> {
                vec![
                    ts_rs::Dependency::from_ty::<<#struct_name as Endpoint>::Request>(),
                    ts_rs::Dependency::from_ty::<<#struct_name as Endpoint>::Response>(),
                    #(
                    ts_rs::Dependency::from_ty::<#req_args_flat>(),
                    )*
                    #(
                    ts_rs::Dependency::from_ty::<#res_args_flat>(),
                    )*
                ].into_iter().flatten().collect::<Vec<_>>()
            }

            fn transparent() -> bool {
                false
            }
        }

    };

    TokenStream::from(expanded)
}