use proc_macro::TokenStream;
use quote::quote;
use syn::{
FnArg, ItemFn, Lit, Meta, Token, parse::Parse, parse::ParseStream, parse_macro_input,
punctuated::Punctuated,
};
struct TaskArgs {
priority: String,
affinity: String,
kind: String,
stack: String,
switcher: String,
}
impl Parse for TaskArgs {
fn parse(input: ParseStream) -> syn::Result<Self> {
let vars = Punctuated::<Meta, Token![,]>::parse_terminated(input)?;
let mut priority = "Normal".to_string();
let mut affinity = "SameCore".to_string();
let mut kind = "Compute".to_string();
let mut stack = "2M".to_string();
let mut switcher = "CrossThreadFloat".to_string();
for var in vars {
if let Meta::NameValue(nv) = var {
if nv.path.is_ident("priority") {
if let syn::Expr::Lit(syn::ExprLit {
lit: Lit::Str(s), ..
}) = nv.value
{
priority = s.value();
}
} else if nv.path.is_ident("affinity") {
if let syn::Expr::Lit(syn::ExprLit {
lit: Lit::Str(s), ..
}) = nv.value
{
affinity = s.value();
}
} else if nv.path.is_ident("kind") {
if let syn::Expr::Lit(syn::ExprLit {
lit: Lit::Str(s), ..
}) = nv.value
{
kind = s.value();
}
} else if nv.path.is_ident("stack") {
if let syn::Expr::Lit(syn::ExprLit {
lit: Lit::Str(s), ..
}) = nv.value
{
stack = s.value();
}
} else if nv.path.is_ident("switcher")
&& let syn::Expr::Lit(syn::ExprLit {
lit: Lit::Str(s), ..
}) = nv.value
{
switcher = s.value();
}
}
}
Ok(TaskArgs {
priority,
affinity,
kind,
stack,
switcher,
})
}
}
#[proc_macro_attribute]
pub fn task(args: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(args as TaskArgs);
let input = parse_macro_input!(item as ItemFn);
let fn_name = &input.sig.ident;
let priority = &args.priority;
let affinity = &args.affinity;
let kind = &args.kind;
let stack = &args.stack;
let metadata_mod = syn::Ident::new(&format!("dtact_metadata_{}", fn_name), fn_name.span());
let priority_ident = syn::Ident::new(priority, fn_name.span());
let affinity_ident = syn::Ident::new(affinity, fn_name.span());
let kind_ident = syn::Ident::new(kind, fn_name.span());
let switcher = &args.switcher;
let switcher_ident = syn::Ident::new(switcher, fn_name.span());
let expanded = quote! {
#input
pub mod #metadata_mod {
pub const PRIORITY: dtact::Priority = dtact::Priority::#priority_ident;
pub const AFFINITY: dtact::topology::Affinity = dtact::topology::Affinity::#affinity_ident;
pub const KIND: dtact::WorkloadKind = dtact::WorkloadKind::#kind_ident;
pub const STACK_SIZE: &'static str = #stack;
pub type SWITCHER = dtact::#switcher_ident;
}
};
TokenStream::from(expanded)
}
#[proc_macro_attribute]
pub fn export_async(_args: TokenStream, item: TokenStream) -> TokenStream {
let input = parse_macro_input!(item as ItemFn);
let fn_name = &input.sig.ident;
let wrapper_name = syn::Ident::new(&format!("dtact_export_{}", fn_name), fn_name.span());
let mut c_params = Vec::new();
let mut call_args = Vec::new();
for input in &input.sig.inputs {
if let FnArg::Typed(pat_type) = input {
let pat = &pat_type.pat;
let ty = &pat_type.ty;
c_params.push(quote! { #pat: #ty });
call_args.push(quote! { #pat });
} else {
panic!("export_async does not support 'self' parameters");
}
}
let expanded = quote! {
#input
#[unsafe(no_mangle)]
pub extern "C" fn #wrapper_name(#(#c_params),*) -> dtact::dtact_handle_t {
dtact::spawn(#fn_name(#(#call_args),*))
}
};
TokenStream::from(expanded)
}
#[proc_macro_attribute]
pub fn export_fiber(_args: TokenStream, item: TokenStream) -> TokenStream {
let input = parse_macro_input!(item as ItemFn);
let fn_name = &input.sig.ident;
let wrapper_name = syn::Ident::new(&format!("dtact_export_fiber_{}", fn_name), fn_name.span());
let mut c_params = Vec::new();
let mut call_args = Vec::new();
for input in &input.sig.inputs {
if let FnArg::Typed(pat_type) = input {
let pat = &pat_type.pat;
let ty = &pat_type.ty;
c_params.push(quote! { #pat: #ty });
call_args.push(quote! { #pat });
} else {
panic!("export_fiber does not support 'self' parameters");
}
}
let expanded = quote! {
#input
#[unsafe(no_mangle)]
pub extern "C" fn #wrapper_name(#(#c_params),*) -> dtact::dtact_handle_t {
dtact::api::fiber::spawn_with_stack("2M", move || {
#fn_name(#(#call_args),*);
})
}
};
TokenStream::from(expanded)
}
struct InitArgs {
topology: String,
safety: String,
workers: usize,
capacity: u32,
stack: usize,
numa: usize,
}
impl Parse for InitArgs {
fn parse(input: ParseStream) -> syn::Result<Self> {
let vars = Punctuated::<Meta, Token![,]>::parse_terminated(input)?;
let mut topology = "P2PMesh".to_string();
let mut safety = "Safety1".to_string();
let mut workers = 0;
let mut capacity = 16384;
let mut stack = 2 * 1024 * 1024;
let mut numa = 0;
for var in vars {
if let Meta::NameValue(nv) = var {
if nv.path.is_ident("topology") {
if let syn::Expr::Lit(syn::ExprLit {
lit: Lit::Str(s), ..
}) = &nv.value
{
topology = s.value();
}
} else if nv.path.is_ident("safety") {
if let syn::Expr::Lit(syn::ExprLit {
lit: Lit::Str(s), ..
}) = &nv.value
{
safety = s.value();
}
} else if nv.path.is_ident("workers") {
if let syn::Expr::Lit(syn::ExprLit {
lit: Lit::Int(i), ..
}) = &nv.value
{
workers = i.base10_parse()?;
}
} else if nv.path.is_ident("capacity") {
if let syn::Expr::Lit(syn::ExprLit {
lit: Lit::Int(i), ..
}) = &nv.value
{
capacity = i.base10_parse()?;
}
} else if nv.path.is_ident("stack") {
if let syn::Expr::Lit(syn::ExprLit {
lit: Lit::Int(i), ..
}) = &nv.value
{
stack = i.base10_parse()?;
}
} else if nv.path.is_ident("numa")
&& let syn::Expr::Lit(syn::ExprLit {
lit: Lit::Int(i), ..
}) = &nv.value
{
numa = i.base10_parse()?;
}
}
}
Ok(InitArgs {
topology,
safety,
workers,
capacity,
stack,
numa,
})
}
}
#[proc_macro_attribute]
pub fn dtact_init(args: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(args as InitArgs);
let input = parse_macro_input!(item as ItemFn);
let topology = &args.topology;
let safety = &args.safety;
let workers = args.workers;
let capacity = args.capacity;
let stack = args.stack;
let numa = args.numa;
let topology_ident = syn::Ident::new(topology, input.sig.ident.span());
let safety_ident = syn::Ident::new(safety, input.sig.ident.span());
let autostart_fn_name = syn::Ident::new("dtact_autostart", input.sig.ident.span());
let attrs = &input.attrs;
let vis = &input.vis;
let sig = &input.sig;
let block = &input.block;
let expanded = quote! {
#[unsafe(no_mangle)]
extern "C" fn #autostart_fn_name() {
let runtime = dtact::GLOBAL_RUNTIME.get_or_init(|| {
let mut workers_count = #workers;
if workers_count == 0 {
workers_count = std::thread::available_parallelism().map(|n| n.get()).unwrap_or(1);
}
let scheduler = dtact::dta_scheduler::DtaScheduler::new(
workers_count,
dtact::dta_scheduler::TopologyMode::#topology_ident
);
let pool = dtact::memory_management::ContextPool::new(
#capacity,
#stack,
dtact::memory_management::SafetyLevel::#safety_ident,
#numa
).expect("DTA-V3 Hardware Initialization Failed");
dtact::Runtime {
scheduler,
pool,
started: core::sync::atomic::AtomicBool::new(false),
shutdown: core::sync::atomic::AtomicBool::new(false),
}
});
runtime.start();
}
#(#attrs)* #vis #sig {
#autostart_fn_name();
#block
}
};
TokenStream::from(expanded)
}