fox-macros 0.2.0

use proc_macro::TokenStream;
use quote::quote;
use syn::{FnArg, Item, ItemFn, Pat, PatIdent, parse_macro_input};

/// Shorthand for deriving serde::{Deserialize, Serialize}
///
/// Prevents explicit `serde` dependency
#[proc_macro_attribute]
pub fn wire(_args: TokenStream, input: TokenStream) -> TokenStream {
    let item = parse_macro_input!(input as Item);

    quote! {
        #[derive(fox::serde::Serialize, fox::serde::Deserialize)]
        #[serde(crate = "fox::serde")]
        #item
    }
    .into()
}

/// Specifies an entrypoint.
///
/// - On `feature = "orchestrator"`
///   - Compiles to normal program.
/// - On `feature = "remote"`
///   - Runs a RPC loop instead.
#[proc_macro_attribute]
pub fn main(_args: TokenStream, input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as ItemFn);

    let (attrs, vis, sig, block) = (&input.attrs, &input.vis, &input.sig, &input.block);

    quote! {
        #(#attrs)*
        #[fox::tokio::main(crate = "fox::tokio")]
        #vis #sig {
            #[cfg(feature = "orchestrator")]
            {
                let (tx, rx) = fox::tokio::sync::mpsc::unbounded_channel();
                fox::ORCHESTRATOR_TX.with(|cell| cell.set(Some(tx)));
                fox::tokio::spawn(fox::orchestrator_rpc_engine(rx));
                #block
                ::std::process::exit(0);
            }
            #[cfg(feature = "remote")]
            fox::remote_rpc().await;
            ::std::process::exit(0);
        }
    }
    .into()
}

/// Marks a function as "running on the remote host"
///
/// - On `feature = "orchestrator"`
///   - Creates a communication shim
/// - On `feature = "remote"`
///   - Registers the function for running
#[proc_macro_attribute]
pub fn remote(_args: TokenStream, input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as ItemFn);

    let (attrs, vis, sig, block) = (&input.attrs, &input.vis, &input.sig, &input.block);
    let function_ident = &sig.ident;
    let inputs = &sig.inputs;
    // sig.generics.params.is_empty();
    let return_type = match &sig.output {
        syn::ReturnType::Default => quote! { () },
        syn::ReturnType::Type(_, ty) => quote! { #ty },
    };

    let mut arg_idents = vec![];
    let mut arg_types = vec![];
    for input_arg in &sig.inputs {
        match input_arg {
            FnArg::Receiver(..) => panic!("Argument `self` not allowed"),
            FnArg::Typed(pat_type) => match &*pat_type.pat {
                Pat::Ident(PatIdent { ident, .. }) => {
                    arg_idents.push(ident);
                    arg_types.push(&pat_type.ty)
                }
                other => panic!("Only standard named arguments are supported. Found: {other:?}"),
            },
        }
    }

    let type_signature = arg_types
        .iter()
        .map(|ty| quote! { #ty }.to_string().replace(' ', ""))
        .collect::<Vec<String>>()
        .join(",");
    let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap_or_default();
    let absolute_span_path = std::path::PathBuf::from(proc_macro::Span::call_site().file());
    let relative_path = absolute_span_path
        .strip_prefix(&manifest_dir)
        .unwrap_or(&absolute_span_path)
        .to_string_lossy()
        .to_string();
    let crate_name = std::env::var("CARGO_PKG_NAME").unwrap_or_default();

    let remote_fn_id = format!(
        "::{}::{}::{}({})",
        crate_name, relative_path, function_ident, type_signature
    );

    quote! {
        #[cfg(feature = "orchestrator")]
        #(#attrs)*
        #vis fn #function_ident (#inputs) -> fox::future::Future<#return_type> {
            let uuid = fox::uuid::Uuid::new_v4();
            let payload = fox::RemoteFnPayload {
                uuid,
                fn_id: #remote_fn_id.to_string(),
                data: fox::postcard::to_allocvec(&( #(#arg_idents),* )).expect("Failed to serialize args"),
            };
            let (tx, rx) = fox::tokio::sync::oneshot::channel();
            fox::ORCHESTRATOR_TX.with(|cell| {
                let opt_tx = cell.take();
                if let Some(ref sender) = opt_tx {
                    sender.send(fox::OutboundRegistration { payload, tx }).unwrap();
                }
                cell.set(opt_tx);
            });
            fox::future::Future::new(rx)
        }

        #[cfg(feature = "remote")]
        #(#attrs)*
        #vis fn #function_ident (#inputs) -> #return_type {
            #block
        }

        #[cfg(feature = "remote")]
        fox::inventory::submit! {
            fox::remote_fn::RemoteFn {
                id: #remote_fn_id,
                function: |args_bytes| {
                    let ( #(#arg_idents),* ): ( #(#arg_types),* ) =
                        fox::postcard::from_bytes(&args_bytes)
                        .expect("Failed to deserialize args");

                    fox::postcard::to_allocvec(&#function_ident( #(#arg_idents),* ))
                        .expect("Failed to serialize return value")
                },
            }
        }
    }
    .into()
}