Skip to main content

fox_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{FnArg, Item, ItemFn, Pat, PatIdent, parse_macro_input};
4
5/// Shorthand for deriving serde::{Deserialize, Serialize}
6///
7/// Prevents explicit `serde` dependency
8#[proc_macro_attribute]
9pub fn rpc(_args: TokenStream, input: TokenStream) -> TokenStream {
10    let item = parse_macro_input!(input as Item);
11
12    quote! {
13        #[derive(fox::serde::Serialize, fox::serde::Deserialize)]
14        #[serde(crate = "fox::serde")]
15        #item
16    }
17    .into()
18}
19
20/// Specifies an entrypoint.
21///
22/// - On `feature = "orchestrator"`
23///   - Compiles to normal program.
24/// - On `feature = "remote"`
25///   - Runs a RPC loop instead.
26#[proc_macro_attribute]
27pub fn main(_args: TokenStream, input: TokenStream) -> TokenStream {
28    let input = parse_macro_input!(input as ItemFn);
29
30    let (attrs, vis, sig, block) = (&input.attrs, &input.vis, &input.sig, &input.block);
31
32    quote! {
33        #(#attrs)*
34        #vis #sig {
35            #[cfg(not(any(feature = "orchestrator", feature = "remote")))]
36            compile_error!("None of the required features were selected");
37
38            if cfg!(feature = "orchestrator") {
39                #block
40            } else {
41                fox::remote_rpc()
42            }
43        }
44    }
45    .into()
46}
47
48/// Marks a boundary of orchestrator<->remote communication
49///
50/// - On `feature = "orchestrator"`
51///   - Creates a communication shim
52/// - On `feature = "remote"`
53///   - Registers the function
54#[proc_macro_attribute]
55pub fn remote(_args: TokenStream, input: TokenStream) -> TokenStream {
56    let input = parse_macro_input!(input as ItemFn);
57
58    let (attrs, vis, sig, block) = (&input.attrs, &input.vis, &input.sig, &input.block);
59    let function_ident = &sig.ident;
60
61    let mut arg_idents = vec![];
62    let mut arg_types = vec![];
63    for input_arg in &sig.inputs {
64        match input_arg {
65            FnArg::Receiver(..) => panic!("Argument `self` not allowed"),
66            FnArg::Typed(pat_type) => match &*pat_type.pat {
67                Pat::Ident(PatIdent { ident, .. }) => {
68                    arg_idents.push(ident);
69                    arg_types.push(&pat_type.ty)
70                }
71                other => panic!("Only standard named arguments are supported. Found: {other:?}"),
72            },
73        }
74    }
75
76    let type_signature = arg_types
77        .iter()
78        .map(|ty| quote! { #ty }.to_string().replace(' ', ""))
79        .collect::<Vec<String>>()
80        .join(",");
81    let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap_or_default();
82    let absolute_span_path = std::path::PathBuf::from(proc_macro::Span::call_site().file());
83    let relative_path = absolute_span_path
84        .strip_prefix(&manifest_dir)
85        .unwrap_or(&absolute_span_path)
86        .to_string_lossy()
87        .to_string();
88    let crate_name = std::env::var("CARGO_PKG_NAME").unwrap_or_default();
89    let bulletproof_seed = format!(
90        "{}::{}::{} Colonial_({})",
91        crate_name, relative_path, function_ident, type_signature
92    );
93    let uuid =
94        uuid::Uuid::new_v5(&uuid::Uuid::NAMESPACE_DNS, bulletproof_seed.as_bytes()).to_string();
95
96    quote! {
97        #(#attrs)*
98        #vis #sig {
99            #[cfg(not(any(feature = "orchestrator", feature = "remote")))]
100            compile_error!("None of the required features were selected");
101
102            if cfg!(feature = "orchestrator") {
103                let rpc_payload = fox::RpcPayload {
104                    uuid: #uuid.to_string(),
105                    data: fox::postcard::to_allocvec(&( #(#arg_idents),* ))
106                            .expect("Failed to serialize args"),
107                };
108
109                let outbound_bytes =
110                    fox::postcard::to_allocvec_cobs(&rpc_payload)
111                        .expect("Failed to serialize payload");
112
113                use ::std::io::Write as _;
114                let mut stdout = ::std::io::stdout().lock();
115                stdout.write_all(&outbound_bytes).expect("Failed to write to stdout");
116                stdout.flush().expect("Failed to flush stdout");
117
118                use std::io::BufRead as _;
119                let mut inbound_bytes = ::std::vec::Vec::new();
120                ::std::io::stdin().lock()
121                    .read_until(0x00, &mut inbound_bytes)
122                    .expect("Failed to read stdin");
123                fox::postcard::from_bytes_cobs(&mut inbound_bytes)
124                    .expect("Failed to deserialize return value")
125            } else {
126                #block
127            }
128        }
129
130        #[cfg(feature = "remote")]
131        fox::inventory::submit! {
132            fox::RemoteFn {
133                uuid: #uuid,
134                function: |args_bytes| {
135                    let ( #(#arg_idents),* ): ( #(#arg_types),* ) =
136                        fox::postcard::from_bytes(&args_bytes)
137                        .expect("Failed to deserialize args");
138
139                    fox::postcard::to_allocvec_cobs(&#function_ident( #(#arg_idents),* ))
140                        .expect("Failed to serialize return value")
141                },
142            }
143        }
144    }
145    .into()
146}