use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use quote::{quote, format_ident};
use syn::{
parse_macro_input, FnArg, Ident, ItemFn, Meta, Pat, ReturnType, Token, Type,
parse::Parser, punctuated::Punctuated,
};
#[proc_macro_attribute]
pub fn polydat_node(attr: TokenStream, item: TokenStream) -> TokenStream {
let func = parse_macro_input!(item as ItemFn);
let attrs = match parse_attrs(attr.into()) {
Ok(a) => a,
Err(e) => return e.to_compile_error().into(),
};
if let Some(adapter) = &attrs.adapter {
let fn_ident = &func.sig.ident;
let prefix = format!("{adapter}_");
if !fn_ident.to_string().starts_with(&prefix) {
let msg = format!(
"#[polydat_node(adapter = \"{adapter}\")] requires the node \
name to start with \"{prefix}\" (found \"{fn_ident}\")",
);
return syn::Error::new_spanned(fn_ident, msg)
.to_compile_error()
.into();
}
}
if attrs.instantiate.is_empty() {
return match generate(func, attrs, None) {
Ok(ts) => ts.into(),
Err(e) => e.to_compile_error().into(),
};
}
match instantiate_and_generate(func, attrs) {
Ok(ts) => ts.into(),
Err(e) => e.to_compile_error().into(),
}
}
fn instantiate_and_generate(
func: ItemFn,
attrs: NodeAttrs,
) -> syn::Result<TokenStream2> {
let generics = &func.sig.generics;
let type_params: Vec<&syn::TypeParam> = generics.type_params().collect();
if type_params.len() != 1 {
return Err(syn::Error::new_spanned(
&func.sig,
format!(
"#[polydat_node(instantiate(...))] requires exactly one type \
parameter on the function (got {}). Declare the function as \
`fn name<T: Wire>(...) -> ...` and list concrete `Wire`-impl \
types in the `instantiate(...)` clause.",
type_params.len(),
),
));
}
let type_param_ident = type_params[0].ident.clone();
let dsl_name = func.sig.ident.to_string();
let mut out = TokenStream2::new();
let mut shared_attrs = attrs.clone();
let instantiations = std::mem::take(&mut shared_attrs.instantiate);
for concrete in instantiations {
let mut inst_func = func.clone();
inst_func.sig.generics.params.clear();
inst_func.sig.generics.where_clause = None;
let mut subst = TypeSubst {
type_param: type_param_ident.clone(),
concrete: concrete.clone(),
};
syn::visit_mut::VisitMut::visit_item_fn_mut(&mut subst, &mut inst_func);
let suffix = type_suffix(&concrete);
let new_ident = syn::Ident::new(
&format!("{dsl_name}_{}", suffix.to_lowercase()),
inst_func.sig.ident.span(),
);
inst_func.sig.ident = new_ident;
let emit = generate(inst_func, shared_attrs.clone(), Some(dsl_name.clone()))?;
out.extend(emit);
}
Ok(out)
}
fn type_suffix(ty: &Type) -> String {
let raw = type_to_string(ty);
let mut out = String::new();
let mut capitalize_next = true;
for c in raw.chars() {
if c.is_alphanumeric() {
if capitalize_next {
out.extend(c.to_uppercase());
capitalize_next = false;
} else {
out.push(c);
}
} else {
capitalize_next = true;
}
}
if out.is_empty() { "Inst".to_string() } else { out }
}
struct TypeSubst {
type_param: syn::Ident,
concrete: Type,
}
impl syn::visit_mut::VisitMut for TypeSubst {
fn visit_type_mut(&mut self, ty: &mut Type) {
if let Type::Path(p) = ty
&& p.qself.is_none() && p.path.is_ident(&self.type_param) {
*ty = self.concrete.clone();
return;
}
syn::visit_mut::visit_type_mut(self, ty);
}
}
#[derive(Clone)]
struct NodeAttrs {
category: Ident,
no_jit: bool,
compiled_u64_override: Option<syn::ExprPath>,
jit_constants_override: Option<syn::ExprPath>,
decompose: Option<syn::ExprPath>,
purity: Option<syn::Expr>,
identity: Option<syn::Expr>,
commutativity: Option<Ident>,
variadic_min: Option<syn::LitInt>,
output_names: Option<Vec<Ident>>,
instantiate: Vec<Type>,
adapter: Option<String>,
}
fn parse_attrs(attr: TokenStream2) -> syn::Result<NodeAttrs> {
if attr.is_empty() {
return Err(syn::Error::new(
proc_macro2::Span::call_site(),
"#[polydat_node] requires `category = <FuncCategory variant>`. \
Example: #[polydat_node(category = Comparison)]",
));
}
let parser = Punctuated::<Meta, Token![,]>::parse_terminated;
let items = parser.parse2(attr)?;
let mut category: Option<Ident> = None;
let mut no_jit = false;
let mut compiled_u64_override: Option<syn::ExprPath> = None;
let mut jit_constants_override: Option<syn::ExprPath> = None;
let mut decompose: Option<syn::ExprPath> = None;
let mut purity: Option<syn::Expr> = None;
let mut identity: Option<syn::Expr> = None;
let mut commutativity: Option<Ident> = None;
let mut variadic_min: Option<syn::LitInt> = None;
let mut output_names: Option<Vec<Ident>> = None;
let mut instantiate: Vec<Type> = Vec::new();
let mut adapter: Option<String> = None;
for item in items {
match item {
Meta::Path(p) => {
let key = p.get_ident()
.ok_or_else(|| syn::Error::new_spanned(
&p,
"#[polydat_node] flag keys must be bare identifiers",
))?
.clone();
match key.to_string().as_str() {
"no_jit" => { no_jit = true; }
other => {
return Err(syn::Error::new_spanned(
&key,
format!(
"#[polydat_node] does not recognize flag `{other}`. \
PR B.7 flags: `no_jit`.",
),
));
}
}
}
Meta::NameValue(nv) => {
let key = nv.path.get_ident()
.ok_or_else(|| syn::Error::new_spanned(
&nv.path,
"#[polydat_node] parameter keys must be bare identifiers",
))?
.clone();
match key.to_string().as_str() {
"category" => {
let syn::Expr::Path(p) = &nv.value else {
return Err(syn::Error::new_spanned(
&nv.value,
"`category` value must be a bare identifier \
(a polydat `FuncCategory` variant name).",
));
};
category = Some(p.path.get_ident()
.ok_or_else(|| syn::Error::new_spanned(
&nv.value,
"`category` value must be a single identifier.",
))?
.clone());
}
"compiled_u64" => {
let syn::Expr::Path(p) = &nv.value else {
return Err(syn::Error::new_spanned(
&nv.value,
"`compiled_u64` value must be a path to a free \
function with signature `fn(&[u64], &mut [u64])`.",
));
};
compiled_u64_override = Some(p.clone());
}
"jit_constants" => {
let syn::Expr::Path(p) = &nv.value else {
return Err(syn::Error::new_spanned(
&nv.value,
"`jit_constants` value must be a path to a free \
function with signature `fn(&Node) -> Vec<u64>`.",
));
};
jit_constants_override = Some(p.clone());
}
"decompose" => {
let syn::Expr::Path(p) = &nv.value else {
return Err(syn::Error::new_spanned(
&nv.value,
"`decompose` value must be a path to a free \
function with signature \
`fn(&Self) -> DecomposedGraph`.",
));
};
decompose = Some(p.clone());
}
"purity" => {
match &nv.value {
syn::Expr::Path(_) | syn::Expr::Call(_) => {
purity = Some(nv.value.clone());
}
_ => {
return Err(syn::Error::new_spanned(
&nv.value,
"`purity` value must be a Purity variant: \
`Pure`, `Nondeterministic`, or \
`SideChannel(<sink>)` where `<sink>` is a \
`SideChannelSink` variant ident.",
));
}
}
}
"identity" => {
identity = Some(nv.value.clone());
}
"commutativity" => {
let syn::Expr::Path(p) = &nv.value else {
return Err(syn::Error::new_spanned(
&nv.value,
"`commutativity` value must be a `Commutativity` \
variant ident (Positional / AllCommutative / ...).",
));
};
commutativity = Some(p.path.get_ident()
.ok_or_else(|| syn::Error::new_spanned(
&nv.value,
"`commutativity` value must be a single identifier.",
))?
.clone());
}
"variadic_min" => {
let syn::Expr::Lit(syn::ExprLit { lit: syn::Lit::Int(n), .. }) = &nv.value else {
return Err(syn::Error::new_spanned(
&nv.value,
"`variadic_min` value must be an integer literal.",
));
};
variadic_min = Some(n.clone());
}
"adapter" => {
let syn::Expr::Lit(syn::ExprLit { lit: syn::Lit::Str(s), .. }) = &nv.value else {
return Err(syn::Error::new_spanned(
&nv.value,
"`adapter` value must be a string literal \
(the adapter's canonical registered name, \
e.g. `adapter = \"cql\"`).",
));
};
adapter = Some(s.value());
}
other => {
return Err(syn::Error::new_spanned(
&key,
format!(
"#[polydat_node] does not recognize parameter `{other}`. \
PR B.2 keys: `category = ...`. PR B.7 keys: \
`no_jit`, `compiled_u64 = ...`, \
`jit_constants = ...`, `purity = ...`. \
PR B.9 keys: `identity = ...`, \
`commutativity = ...`, `variadic_min = ...`. \
Namespacing: `adapter = \"...\"`.",
),
));
}
}
}
Meta::List(list) => {
let key = list.path.get_ident()
.ok_or_else(|| syn::Error::new_spanned(
&list.path,
"#[polydat_node] list-form keys must be bare identifiers",
))?
.clone();
match key.to_string().as_str() {
"output_names" => {
let names: Punctuated<Ident, Token![,]> =
list.parse_args_with(Punctuated::parse_terminated)?;
if names.is_empty() {
return Err(syn::Error::new_spanned(
&list,
"`output_names(...)` requires at least one name.",
));
}
output_names = Some(names.into_iter().collect());
}
"instantiate" => {
let types: Punctuated<Type, Token![,]> =
list.parse_args_with(Punctuated::parse_terminated)?;
if types.is_empty() {
return Err(syn::Error::new_spanned(
&list,
"`instantiate(...)` requires at least one type. \
List the concrete `Wire`-impl types that should \
get their own per-instantiation registrations.",
));
}
instantiate = types.into_iter().collect();
}
other => {
return Err(syn::Error::new_spanned(
&key,
format!(
"#[polydat_node] does not recognize list-form key `{other}`. \
Recognised: `output_names(...)`, `instantiate(...)`.",
),
));
}
}
}
}
}
let category = category.ok_or_else(|| syn::Error::new(
proc_macro2::Span::call_site(),
"#[polydat_node] requires `category = <FuncCategory variant>`.",
))?;
Ok(NodeAttrs {
category,
no_jit,
compiled_u64_override,
jit_constants_override,
decompose,
purity,
identity,
commutativity,
variadic_min,
output_names,
instantiate,
adapter,
})
}
struct ClassifiedArg {
name: syn::Ident,
declared_ty: Type,
kind: ArgKind,
default_value: Option<syn::Expr>,
wire_constraint: Option<Ident>,
}
#[derive(Clone)]
enum ArgKind {
Wire,
Const(ConstShape),
ConstVec(ConstShape),
Setup(Box<SetupSpec>),
PolyWire,
Variadic(VariadicElement),
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum VariadicElement {
U64,
#[allow(dead_code)]
F64,
Bool,
BorrowedStr,
OwnedString,
Value,
}
impl VariadicElement {
fn port_type_tokens(self) -> TokenStream2 {
match self {
VariadicElement::U64 => quote!(polydat::ast::PortType::U64),
VariadicElement::F64 => quote!(polydat::ast::PortType::F64),
VariadicElement::Bool => quote!(polydat::ast::PortType::Bool),
VariadicElement::BorrowedStr => quote!(polydat::ast::PortType::Str),
VariadicElement::OwnedString => quote!(polydat::ast::PortType::Str),
VariadicElement::Value => quote!(polydat::ast::PortType::Str),
}
}
fn extract_from_value(self) -> TokenStream2 {
match self {
VariadicElement::U64 => quote!(|v: &polydat::ast::Value| v.as_u64()),
VariadicElement::F64 => quote!(|v: &polydat::ast::Value| v.as_f64()),
VariadicElement::Bool => quote!(|v: &polydat::ast::Value| v.as_bool()),
VariadicElement::BorrowedStr => quote!(|v: &polydat::ast::Value| v.as_str()),
VariadicElement::OwnedString => quote!(|v: &polydat::ast::Value| v.as_str().to_string()),
VariadicElement::Value => quote!(|v: &polydat::ast::Value| v.clone()),
}
}
}
#[derive(Clone)]
struct SetupSpec {
inner_ty: Type,
setup_fn: syn::Expr,
source_args: Vec<syn::Ident>,
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum ConstShape {
U64,
F64,
Bool,
Str,
}
impl ConstShape {
fn slot_type_tokens(self) -> TokenStream2 {
match self {
ConstShape::U64 => quote!(polydat::ast::SlotType::ConstU64),
ConstShape::F64 => quote!(polydat::ast::SlotType::ConstF64),
ConstShape::Bool => quote!(polydat::ast::SlotType::ConstU64),
ConstShape::Str => quote!(polydat::ast::SlotType::ConstStr),
}
}
fn field_type_tokens(self) -> TokenStream2 {
match self {
ConstShape::U64 => quote!(u64),
ConstShape::F64 => quote!(f64),
ConstShape::Bool => quote!(bool),
ConstShape::Str => quote!(String),
}
}
fn extract_from_const_arg(self, c: TokenStream2) -> TokenStream2 {
match self {
ConstShape::U64 => quote!(#c.as_u64()),
ConstShape::F64 => quote!(#c.as_f64()),
ConstShape::Bool => quote!(#c.as_u64() != 0),
ConstShape::Str => quote!(#c.as_str().to_string()),
}
}
fn wrap_as_const(self, field_ref: TokenStream2) -> TokenStream2 {
match self {
ConstShape::U64 => quote!(polydat::derive_support::Const(#field_ref)),
ConstShape::F64 => quote!(polydat::derive_support::Const(#field_ref)),
ConstShape::Bool => quote!(polydat::derive_support::Const(#field_ref)),
ConstShape::Str => quote!(polydat::derive_support::Const(#field_ref.as_str())),
}
}
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum JitType {
U64,
I64,
F64,
Bool,
U8,
U16,
U32,
I8,
I16,
I32,
F32,
F16,
Str,
Bytes,
U128,
I128,
RegRaw,
RegI8x16,
RegI16x8,
RegI32x4,
RegI64x2,
RegF16x8,
RegF32x4,
RegF64x2,
}
impl JitType {
fn width(self) -> usize {
match self {
JitType::U128 | JitType::I128 | JitType::RegRaw
| JitType::RegI8x16 | JitType::RegI16x8 | JitType::RegI32x4
| JitType::RegI64x2 | JitType::RegF16x8 | JitType::RegF32x4
| JitType::RegF64x2 => 2,
_ => 1,
}
}
fn read_from_u64_buffer(self, idx: usize) -> TokenStream2 {
let i = syn::Index::from(idx);
let i1 = syn::Index::from(idx + 1);
let limbs = quote!(polydat::ast::Bits128([inputs[#i], inputs[#i1]]));
match self {
JitType::U64 => quote!(inputs[#i]),
JitType::I64 => quote!(inputs[#i] as i64),
JitType::F64 => quote!(f64::from_bits(inputs[#i])),
JitType::Bool => quote!(inputs[#i] != 0),
JitType::U8 => quote!(inputs[#i] as u8),
JitType::U16 => quote!(inputs[#i] as u16),
JitType::U32 => quote!(inputs[#i] as u32),
JitType::I8 => quote!((inputs[#i] as i64) as i8),
JitType::I16 => quote!((inputs[#i] as i64) as i16),
JitType::I32 => quote!((inputs[#i] as i64) as i32),
JitType::F32 => quote!(f32::from_bits(inputs[#i] as u32)),
JitType::F16 => quote!(polydat::half::f16::from_bits(inputs[#i] as u16)),
JitType::Str => quote!(polydat::kernel::resolve_thread_str(inputs[#i]).into()),
JitType::Bytes => quote!(polydat::kernel::resolve_thread_bytes(inputs[#i]).into()),
JitType::U128 => quote!((#limbs).as_u128()),
JitType::I128 => quote!((#limbs).as_i128()),
JitType::RegRaw => limbs,
JitType::RegI8x16 => quote!((#limbs).lanes_i8()),
JitType::RegI16x8 => quote!((#limbs).lanes_i16()),
JitType::RegI32x4 => quote!((#limbs).lanes_i32()),
JitType::RegI64x2 => quote!((#limbs).lanes_i64()),
JitType::RegF16x8 => quote!((#limbs).lanes_f16()),
JitType::RegF32x4 => quote!((#limbs).lanes_f32()),
JitType::RegF64x2 => quote!((#limbs).lanes_f64()),
}
}
fn write_to_u64_buffer_at(self, base: usize, result: TokenStream2) -> TokenStream2 {
let o = syn::Index::from(base);
let o1 = syn::Index::from(base + 1);
let write_limbs = |from: TokenStream2| {
quote! {{
let __limbs = #from;
outputs[#o] = __limbs.0[0];
outputs[#o1] = __limbs.0[1];
}}
};
match self {
JitType::U64 => quote!(outputs[#o] = #result;),
JitType::I64 => quote!(outputs[#o] = (#result) as u64;),
JitType::F64 => quote!(outputs[#o] = (#result).to_bits();),
JitType::Bool => quote!(outputs[#o] = if #result { 1 } else { 0 };),
JitType::U8 | JitType::U16 | JitType::U32
=> quote!(outputs[#o] = (#result) as u64;),
JitType::I8 | JitType::I16 | JitType::I32
=> quote!(outputs[#o] = ((#result) as i64) as u64;),
JitType::F32 => quote!(outputs[#o] = (#result).to_bits() as u64;),
JitType::F16 => quote!(outputs[#o] = (#result).to_bits() as u64;),
JitType::Str => quote!(outputs[#o] = polydat::kernel::put_thread_str((#result).as_ref());),
JitType::Bytes => quote!(outputs[#o] = polydat::kernel::put_thread_bytes((#result).as_ref());),
JitType::U128 => write_limbs(quote!(polydat::ast::Bits128::from_u128(#result))),
JitType::I128 => write_limbs(quote!(polydat::ast::Bits128::from_i128(#result))),
JitType::RegRaw => write_limbs(quote!(#result)),
JitType::RegI8x16 => write_limbs(quote!(polydat::ast::Bits128::from_lanes_i8(#result))),
JitType::RegI16x8 => write_limbs(quote!(polydat::ast::Bits128::from_lanes_i16(#result))),
JitType::RegI32x4 => write_limbs(quote!(polydat::ast::Bits128::from_lanes_i32(#result))),
JitType::RegI64x2 => write_limbs(quote!(polydat::ast::Bits128::from_lanes_i64(#result))),
JitType::RegF16x8 => write_limbs(quote!(polydat::ast::Bits128::from_lanes_f16(#result))),
JitType::RegF32x4 => write_limbs(quote!(polydat::ast::Bits128::from_lanes_f32(#result))),
JitType::RegF64x2 => write_limbs(quote!(polydat::ast::Bits128::from_lanes_f64(#result))),
}
}
fn write_to_u64_buffer(self, result: TokenStream2) -> TokenStream2 {
self.write_to_u64_buffer_at(0, result)
}
fn const_field_as_u64(self, field_ref: TokenStream2) -> TokenStream2 {
match self {
JitType::U64 => quote!(#field_ref),
JitType::I64 => quote!((#field_ref) as u64),
JitType::F64 => quote!((#field_ref).to_bits()),
JitType::Bool => quote!(if #field_ref { 1 } else { 0 }),
JitType::U8 | JitType::U16 | JitType::U32
=> quote!((#field_ref) as u64),
JitType::I8 | JitType::I16 | JitType::I32
=> quote!(((#field_ref) as i64) as u64),
JitType::F32 | JitType::F16
=> quote!((#field_ref).to_bits() as u64),
JitType::Str | JitType::Bytes
=> quote!(polydat::kernel::StaticInterner::intern((#field_ref).as_ref())),
JitType::U128 | JitType::I128 | JitType::RegRaw
| JitType::RegI8x16 | JitType::RegI16x8 | JitType::RegI32x4
| JitType::RegI64x2 | JitType::RegF16x8 | JitType::RegF32x4
| JitType::RegF64x2 => {
unreachable!("128-bit/register types have no const shape")
}
}
}
}
fn const_shape_to_jit_type(s: ConstShape) -> Option<JitType> {
match s {
ConstShape::U64 => Some(JitType::U64),
ConstShape::F64 => Some(JitType::F64),
ConstShape::Bool => Some(JitType::Bool),
ConstShape::Str => Some(JitType::Str),
}
}
fn wire_type_to_jit_type(ty: &Type) -> Option<JitType> {
let s = type_to_string(ty);
match s.as_str() {
"u64" => Some(JitType::U64),
"i64" => Some(JitType::I64),
"f64" => Some(JitType::F64),
"bool" => Some(JitType::Bool),
"u8" => Some(JitType::U8),
"u16" => Some(JitType::U16),
"u32" => Some(JitType::U32),
"i8" => Some(JitType::I8),
"i16" => Some(JitType::I16),
"i32" => Some(JitType::I32),
"f32" => Some(JitType::F32),
"u128" => Some(JitType::U128),
"i128" => Some(JitType::I128),
"String" | "& str" | "&str" => Some(JitType::Str),
"Vec < u8 >" | "Vec<u8>" | "& [ u8 ]" | "&[u8]" => Some(JitType::Bytes),
_ => {
let flat: String = s.split_whitespace().collect();
match flat.as_str() {
"Arc<str>" | "std::sync::Arc<str>" => Some(JitType::Str),
"Arc<[u8]>" | "std::sync::Arc<[u8]>" => Some(JitType::Bytes),
"half::f16" | "f16" => Some(JitType::F16),
"Bits128" | "crate::ast::Bits128" | "polydat::ast::Bits128"
| "ast::Bits128" => Some(JitType::RegRaw),
"[i8;16]" => Some(JitType::RegI8x16),
"[i16;8]" => Some(JitType::RegI16x8),
"[i32;4]" => Some(JitType::RegI32x4),
"[i64;2]" => Some(JitType::RegI64x2),
"[half::f16;8]" | "[f16;8]" => Some(JitType::RegF16x8),
"[f32;4]" => Some(JitType::RegF32x4),
"[f64;2]" => Some(JitType::RegF64x2),
_ => None,
}
}
}
}
fn classify_type(ty: &Type) -> Option<ConstShape> {
let syn::Type::Path(p) = ty else { return None; };
let last = p.path.segments.last()?;
if last.ident != "Const" { return None; }
let syn::PathArguments::AngleBracketed(args) = &last.arguments else { return None; };
let inner = args.args.iter().find_map(|a| {
if let syn::GenericArgument::Type(t) = a { Some(t) } else { None }
})?;
let s = type_to_string(inner);
match s.as_str() {
"u64" => Some(ConstShape::U64),
"f64" => Some(ConstShape::F64),
"bool" => Some(ConstShape::Bool),
"& str" | "&str" => Some(ConstShape::Str),
_ => None,
}
}
fn classify_const_vec(ty: &Type) -> Option<ConstShape> {
let syn::Type::Path(p) = ty else { return None; };
let last = p.path.segments.last()?;
if last.ident != "Const" { return None; }
let syn::PathArguments::AngleBracketed(args) = &last.arguments else { return None; };
let inner = args.args.iter().find_map(|a| {
if let syn::GenericArgument::Type(t) = a { Some(t) } else { None }
})?;
let syn::Type::Path(vp) = inner else { return None; };
let vlast = vp.path.segments.last()?;
if vlast.ident != "Vec" { return None; }
let syn::PathArguments::AngleBracketed(vargs) = &vlast.arguments else { return None; };
let velem = vargs.args.iter().find_map(|a| {
if let syn::GenericArgument::Type(t) = a { Some(t) } else { None }
})?;
let s = type_to_string(velem);
match s.as_str() {
"u64" => Some(ConstShape::U64),
"f64" => Some(ConstShape::F64),
"bool" => Some(ConstShape::Bool),
"String" => Some(ConstShape::Str),
"& str" | "&str" => Some(ConstShape::Str),
_ => None,
}
}
fn classify_dynamic_outputs(ty: &Type) -> Option<Type> {
let syn::Type::Path(p) = ty else { return None; };
let last = p.path.segments.last()?;
if last.ident != "DynamicOutputs" { return None; }
let syn::PathArguments::AngleBracketed(args) = &last.arguments else { return None; };
args.args.iter().find_map(|a| {
if let syn::GenericArgument::Type(t) = a { Some(t.clone()) } else { None }
})
}
fn parse_poly_default(attrs: &[syn::Attribute]) -> syn::Result<Option<syn::Expr>> {
for attr in attrs {
if !attr.path().is_ident("poly_default") { continue; }
let expr: syn::Expr = attr.parse_args()?;
return Ok(Some(expr));
}
Ok(None)
}
fn parse_wire_constraint(attrs: &[syn::Attribute]) -> syn::Result<Option<Ident>> {
for attr in attrs {
if !attr.path().is_ident("constraint") { continue; }
let variant: Ident = attr.parse_args()?;
return Ok(Some(variant));
}
Ok(None)
}
fn parse_poly_const(attrs: &[syn::Attribute]) -> syn::Result<Option<(syn::Expr, Vec<syn::Ident>)>> {
for attr in attrs {
if !attr.path().is_ident("poly_const") { continue; }
let parser = |input: syn::parse::ParseStream| -> syn::Result<(syn::Expr, Vec<syn::Ident>)> {
let fn_expr: syn::Expr = input.parse()?;
let _comma: Token![,] = input.parse()?;
let from_kw: syn::Ident = input.parse()?;
if from_kw != "from" {
return Err(syn::Error::new_spanned(
from_kw,
"#[poly_const(...)] requires a `from = <source>` clause. \
Supported shapes: `from = ()` (empty), `from = ident` \
(single), `from = (a, b, c)` (multi-source).",
));
}
let _eq: Token![=] = input.parse()?;
if input.peek(syn::token::Paren) {
let inner;
let _paren = syn::parenthesized!(inner in input);
if inner.is_empty() {
return Ok((fn_expr, Vec::new()));
}
let parsed: Punctuated<syn::Ident, Token![,]> =
Punctuated::parse_terminated(&inner)?;
if parsed.is_empty() {
return Err(syn::Error::new_spanned(
from_kw,
"#[poly_const(..., from = (...))] — the parenthesised \
form expects a comma-separated list of source-arg \
identifiers, or an empty `()` for session-static \
setup.",
));
}
return Ok((fn_expr, parsed.into_iter().collect()));
}
let source: syn::Ident = input.parse()?;
Ok((fn_expr, vec![source]))
};
let parsed = attr.parse_args_with(parser)?;
return Ok(Some(parsed));
}
Ok(None)
}
fn classify_borrowed(ty: &Type) -> Option<Type> {
let syn::Type::Reference(r) = ty else { return None; };
if r.mutability.is_some() { return None; }
Some((*r.elem).clone())
}
fn classify_polywire(ty: &Type) -> bool {
let syn::Type::Path(p) = ty else { return false; };
p.path.segments.last().map(|s| s.ident == "Value").unwrap_or(false)
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum WrapperWire {
Bytes,
Json,
Handle,
VecF32, VecI32, VecF64, VecI64, VecF16, VecI16, VecI8,
}
fn classify_wrapper_wire(ty: &Type) -> Option<WrapperWire> {
if let Some(kind) = classify_vec_wire(ty) {
return Some(kind);
}
if let Some(inner) = strip_arc(ty)
&& let syn::Type::Slice(slc) = inner
&& let syn::Type::Path(p) = &*slc.elem
&& p.path.is_ident("u8")
{
return Some(WrapperWire::Bytes);
}
if let Some(inner) = strip_arc(ty)
&& let syn::Type::Path(p) = inner
&& last_segment_is(p, "Value")
&& path_contains_segment(p, "serde_json")
{
return Some(WrapperWire::Json);
}
if let Some(inner) = strip_arc(ty)
&& let syn::Type::Path(p) = inner
&& p.path.is_ident("str")
{
return None;
}
if let Some(inner) = strip_arc(ty)
&& matches!(inner, syn::Type::TraitObject(_))
{
return None;
}
if strip_arc(ty).is_some() {
return Some(WrapperWire::Handle);
}
if let syn::Type::Path(p) = ty
&& let Some(last) = p.path.segments.last()
&& last.ident == "Vec"
&& let syn::PathArguments::AngleBracketed(args) = &last.arguments
&& let Some(syn::GenericArgument::Type(syn::Type::Path(elem))) = args.args.first()
&& elem.path.is_ident("u8")
{
return Some(WrapperWire::Bytes);
}
if let syn::Type::Reference(r) = ty
&& r.mutability.is_none()
&& let syn::Type::Slice(slc) = &*r.elem
&& let syn::Type::Path(p) = &*slc.elem
&& p.path.is_ident("u8")
{
return Some(WrapperWire::Bytes);
}
if let syn::Type::Reference(r) = ty
&& r.mutability.is_none()
&& let syn::Type::Path(p) = &*r.elem
&& last_segment_is(p, "Value")
&& path_contains_segment(p, "serde_json")
{
return Some(WrapperWire::Json);
}
None
}
fn strip_arc(ty: &Type) -> Option<&Type> {
let syn::Type::Path(p) = ty else { return None; };
let last = p.path.segments.last()?;
if last.ident != "Arc" { return None; }
let syn::PathArguments::AngleBracketed(args) = &last.arguments else { return None; };
args.args.iter().find_map(|a| match a {
syn::GenericArgument::Type(t) => Some(t),
_ => None,
})
}
fn last_segment_is(p: &syn::TypePath, name: &str) -> bool {
p.path.segments.last().map(|s| s.ident == name).unwrap_or(false)
}
fn path_contains_segment(p: &syn::TypePath, name: &str) -> bool {
p.path.segments.iter().any(|s| s.ident == name)
}
fn extract_handle_inner(ty: &Type) -> Option<Type> {
strip_arc(ty).cloned()
}
fn is_option_arg(ty: &Type) -> bool {
let syn::Type::Path(p) = ty else { return false; };
let Some(last) = p.path.segments.last() else { return false; };
if last.ident != "Option" { return false; }
matches!(&last.arguments,
syn::PathArguments::AngleBracketed(args)
if args.args.iter().any(|a| matches!(a, syn::GenericArgument::Type(_))))
}
#[derive(Clone)]
enum BorrowWire {
Str,
Bytes,
Json,
Vec(&'static str , TokenStream2 ),
}
fn is_borrow_wire_shape(ty: &Type) -> Option<BorrowWire> {
let syn::Type::Reference(r) = ty else { return None; };
if r.mutability.is_some() { return None; }
match &*r.elem {
syn::Type::Path(p) if p.path.is_ident("str") => Some(BorrowWire::Str),
syn::Type::Slice(slc) => {
if let syn::Type::Path(p) = &*slc.elem {
if p.path.is_ident("u8") {
return Some(BorrowWire::Bytes);
}
let elem_name = p.path.segments.last()?.ident.to_string();
let (variant, port_expr) = match elem_name.as_str() {
"f32" => ("VecF32", quote!(polydat::ast::PortType::VecF32)),
"i32" => ("VecI32", quote!(polydat::ast::PortType::VecI32)),
"f64" => ("VecF64", quote!(polydat::ast::PortType::VecF64)),
"i64" => ("VecI64", quote!(polydat::ast::PortType::VecI64)),
"f16" => ("VecF16", quote!(polydat::ast::PortType::VecF16)),
"i16" => ("VecI16", quote!(polydat::ast::PortType::VecI16)),
"i8" => ("VecI8", quote!(polydat::ast::PortType::VecI8)),
_ => return None,
};
return Some(BorrowWire::Vec(variant, port_expr));
}
None
}
syn::Type::Path(p) if last_segment_is(p, "Value")
&& path_contains_segment(p, "serde_json") => Some(BorrowWire::Json),
_ => None,
}
}
fn borrow_extract_tokens(shape: BorrowWire, input_expr: TokenStream2) -> TokenStream2 {
match shape {
BorrowWire::Str => quote! {
match #input_expr {
polydat::ast::Value::Str(__arc) => __arc.as_ref(),
__other => panic!("expected Str wire, got {__other:?}"),
}
},
BorrowWire::Bytes => quote! {
match #input_expr {
polydat::ast::Value::Bytes(__arc) => __arc.as_ref(),
__other => panic!("expected Bytes wire, got {__other:?}"),
}
},
BorrowWire::Json => quote! {
match #input_expr {
polydat::ast::Value::Json(__arc) => __arc.as_ref(),
__other => panic!("expected Json wire, got {__other:?}"),
}
},
BorrowWire::Vec(variant, _port) => {
let v = syn::Ident::new(variant, proc_macro2::Span::call_site());
quote! {
match #input_expr {
polydat::ast::Value::#v(__arc) => __arc.as_slice(),
__other => panic!(
concat!("expected ", stringify!(#v), " wire, got {:?}"),
__other),
}
}
}
}
}
fn borrow_port_type(shape: &BorrowWire) -> TokenStream2 {
match shape {
BorrowWire::Str => quote!(polydat::ast::PortType::Str),
BorrowWire::Bytes => quote!(polydat::ast::PortType::Bytes),
BorrowWire::Json => quote!(polydat::ast::PortType::Json),
BorrowWire::Vec(_, port_expr) => port_expr.clone(),
}
}
fn classify_vec_wire(ty: &Type) -> Option<WrapperWire> {
let elem: Type = if let Some(elem) = strip_vec(ty) {
elem.clone()
} else if let Some(elem) = strip_slice_arc(ty) {
elem.clone()
} else if let Some(elem) = strip_borrowed_slice(ty) {
elem.clone()
} else {
return None;
};
let syn::Type::Path(p) = &elem else { return None; };
let last = p.path.segments.last()?;
match last.ident.to_string().as_str() {
"f32" => Some(WrapperWire::VecF32),
"i32" => Some(WrapperWire::VecI32),
"f64" => Some(WrapperWire::VecF64),
"i64" => Some(WrapperWire::VecI64),
"f16" => Some(WrapperWire::VecF16),
"i16" => Some(WrapperWire::VecI16),
"i8" => Some(WrapperWire::VecI8),
_ => None,
}
}
fn strip_vec(ty: &Type) -> Option<&Type> {
let syn::Type::Path(p) = ty else { return None; };
let last = p.path.segments.last()?;
if last.ident != "Vec" { return None; }
let syn::PathArguments::AngleBracketed(args) = &last.arguments else { return None; };
args.args.iter().find_map(|a| match a {
syn::GenericArgument::Type(t) => Some(t),
_ => None,
})
}
fn strip_slice_arc(ty: &Type) -> Option<&Type> {
let syn::Type::Path(p) = ty else { return None; };
let last = p.path.segments.last()?;
if last.ident != "SliceArc" { return None; }
let syn::PathArguments::AngleBracketed(args) = &last.arguments else { return None; };
args.args.iter().find_map(|a| match a {
syn::GenericArgument::Type(t) => Some(t),
_ => None,
})
}
fn strip_borrowed_slice(ty: &Type) -> Option<&Type> {
let syn::Type::Reference(r) = ty else { return None; };
if r.mutability.is_some() { return None; }
let syn::Type::Slice(slc) = &*r.elem else { return None; };
Some(&slc.elem)
}
fn classify_variadic(ty: &Type) -> Option<VariadicElement> {
let syn::Type::Reference(r) = ty else { return None; };
if r.mutability.is_some() { return None; }
let syn::Type::Slice(s) = &*r.elem else { return None; };
if let syn::Type::Reference(inner_r) = &*s.elem
&& inner_r.mutability.is_none()
&& let syn::Type::Path(p) = &*inner_r.elem
&& p.path.is_ident("str")
{
return Some(VariadicElement::BorrowedStr);
}
let syn::Type::Path(p) = &*s.elem else { return None; };
let last = p.path.segments.last()?;
if !last.arguments.is_empty() { return None; }
match last.ident.to_string().as_str() {
"u64" => Some(VariadicElement::U64),
"bool" => Some(VariadicElement::Bool),
"String" => Some(VariadicElement::OwnedString),
"Value" => Some(VariadicElement::Value),
_ => None,
}
}
fn classify_result_return(ty: &Type) -> Option<Type> {
let syn::Type::Path(p) = ty else { return None; };
let last = p.path.segments.last()?;
if last.ident != "Result" { return None; }
let syn::PathArguments::AngleBracketed(args) = &last.arguments else { return None; };
let mut tys = args.args.iter().filter_map(|a| match a {
syn::GenericArgument::Type(t) => Some(t.clone()),
_ => None,
});
tys.next()
}
fn generate(
func: ItemFn,
attrs: NodeAttrs,
dsl_name_override: Option<String>,
) -> syn::Result<TokenStream2> {
let fn_name = &func.sig.ident;
let fn_name_raw = fn_name.to_string();
let rust_name_str = fn_name_raw
.strip_prefix("r#")
.unwrap_or(&fn_name_raw)
.to_string();
let struct_name = format_ident!("{}", to_camel_case(&rust_name_str));
let is_instantiation = dsl_name_override.is_some();
let func_name_str = dsl_name_override.unwrap_or_else(|| rust_name_str.clone());
let category = &attrs.category;
let mut args: Vec<ClassifiedArg> = Vec::new();
for input in &func.sig.inputs {
match input {
FnArg::Receiver(r) => {
return Err(syn::Error::new_spanned(
r,
"#[polydat_node] does not support `self` parameters yet; \
state-bearing nodes are deferred to a later PR.",
));
}
FnArg::Typed(pat_ty) => {
let ident = match &*pat_ty.pat {
Pat::Ident(p) => p.ident.clone(),
other => {
return Err(syn::Error::new_spanned(
other,
"#[polydat_node] requires plain identifier parameters; \
pattern matching in argument position isn't supported.",
));
}
};
let declared_ty = (*pat_ty.ty).clone();
let default_value = parse_poly_default(&pat_ty.attrs)?;
let setup_attr = parse_poly_const(&pat_ty.attrs)?;
let wire_constraint = parse_wire_constraint(&pat_ty.attrs)?;
let is_polywire = classify_polywire(&declared_ty);
let variadic_elem = classify_variadic(&declared_ty);
let kind = if let Some(elem) = variadic_elem {
if default_value.is_some() || setup_attr.is_some() || is_polywire {
return Err(syn::Error::new_spanned(
pat_ty,
"variadic `&[T]` args don't combine with \
#[poly_default(...)], #[poly_const(...)], or `Value`.",
));
}
ArgKind::Variadic(elem)
} else if is_polywire {
if default_value.is_some() || setup_attr.is_some() {
return Err(syn::Error::new_spanned(
pat_ty,
"`Value` args (PolyWire) don't combine with \
#[poly_default(...)] or #[poly_const(...)]; \
the runtime port type comes from the upstream wire \
at construction time.",
));
}
ArgKind::PolyWire
} else if let Some((setup_fn, source_args)) = setup_attr {
let inner_ty = classify_borrowed(&declared_ty)
.ok_or_else(|| syn::Error::new_spanned(
&declared_ty,
"#[poly_const(...)] requires the argument type to be \
a borrow `&T` — the macro stores the computed `T` \
in a struct field and hands the body a borrow each \
eval.",
))?;
if default_value.is_some() {
return Err(syn::Error::new_spanned(
pat_ty,
"#[poly_default(...)] cannot combine with \
#[poly_const(...)]; defaults belong on the source \
Const arg, not on the derived setup arg.",
));
}
ArgKind::Setup(Box::new(SetupSpec { inner_ty, setup_fn, source_args }))
} else if let Some(inner) = classify_const_vec(&declared_ty) {
if default_value.is_some() {
return Err(syn::Error::new_spanned(
pat_ty,
"#[poly_default(...)] cannot combine with \
`Const<Vec<C>>`; the empty Vec IS the implicit \
default. Use `Const<C>` with a poly_default \
literal for a single-value default instead.",
));
}
if setup_attr.is_some() {
return Err(syn::Error::new_spanned(
pat_ty,
"`Const<Vec<C>>` doesn't combine with \
#[poly_const(...)]; route the derived state \
from a scalar `Const<C>` source instead.",
));
}
ArgKind::ConstVec(inner)
} else {
match classify_type(&declared_ty) {
Some(shape) => ArgKind::Const(shape),
None => {
if default_value.is_some() {
return Err(syn::Error::new_spanned(
pat_ty,
"#[poly_default(...)] only applies to const args \
(`Const<T>`); bare-type wire args don't have \
assembly-time defaults.",
));
}
ArgKind::Wire
}
}
};
args.push(ClassifiedArg { name: ident, declared_ty, kind, default_value, wire_constraint });
}
}
}
{
let const_vec_positions: Vec<usize> = args.iter().enumerate()
.filter_map(|(i, a)| if matches!(a.kind, ArgKind::ConstVec(_)) { Some(i) } else { None })
.collect();
if const_vec_positions.len() > 1 {
return Err(syn::Error::new_spanned(
&args[const_vec_positions[1]].declared_ty,
"#[polydat_node] supports at most one `Const<Vec<C>>` arg \
per function; the variadic-const surface consumes the \
tail of the consts slice and a second one would have no \
entries to claim.",
));
}
if let Some(&pos) = const_vec_positions.first() {
for later in &args[pos + 1..] {
if matches!(later.kind, ArgKind::Const(_)) {
return Err(syn::Error::new_spanned(
&later.declared_ty,
"scalar `Const<T>` arg declared after a \
`Const<Vec<C>>` arg is unreachable — the variadic \
consumes everything from its position to the end \
of the consts slice. Move the scalar consts BEFORE \
the `Const<Vec<C>>` in the function signature.",
));
}
}
}
}
let wire_port_type_for = |ty: &Type| -> syn::Result<TokenStream2> {
if let Some(kind) = classify_wrapper_wire(ty) {
return Ok(match kind {
WrapperWire::Bytes => quote!(polydat::ast::PortType::Bytes),
WrapperWire::Json => quote!(polydat::ast::PortType::Json),
WrapperWire::Handle => quote!(polydat::ast::PortType::Handle),
WrapperWire::VecF32 => quote!(polydat::ast::PortType::VecF32),
WrapperWire::VecI32 => quote!(polydat::ast::PortType::VecI32),
WrapperWire::VecF64 => quote!(polydat::ast::PortType::VecF64),
WrapperWire::VecI64 => quote!(polydat::ast::PortType::VecI64),
WrapperWire::VecF16 => quote!(polydat::ast::PortType::VecF16),
WrapperWire::VecI16 => quote!(polydat::ast::PortType::VecI16),
WrapperWire::VecI8 => quote!(polydat::ast::PortType::VecI8),
});
}
if let Some(borrow) = is_borrow_wire_shape(ty) {
return Ok(borrow_port_type(&borrow));
}
Ok(quote!(<#ty as polydat::derive_support::Wire>::PORT))
};
let mut slot_exprs: Vec<TokenStream2> = Vec::new();
for a in &args {
let name_str = a.name.to_string();
match &a.kind {
ArgKind::Wire => {
let pt = wire_port_type_for(&a.declared_ty)?;
let ty = &a.declared_ty;
let constraint_chain = if let Some(variant) = &a.wire_constraint {
quote! {
.with_constraint(
polydat::dsl::const_constraints::ConstConstraint::#variant)
}
} else {
quote!()
};
let cost_chain = if is_borrow_wire_shape(ty).is_none()
&& classify_wrapper_wire(ty) != Some(WrapperWire::Handle)
{
quote! {
.with_cost(<#ty as polydat::derive_support::Wire>::WIRE_COST)
}
} else {
quote!()
};
slot_exprs.push(quote! {
polydat::ast::Slot::Wire(
polydat::ast::Port::new(#name_str, #pt)
#constraint_chain
#cost_chain
)
});
}
ArgKind::Const(shape) => {
let field_name = &a.name;
let const_value_ctor = match shape {
ConstShape::U64 => quote!(polydat::ast::ConstValue::U64(#field_name)),
ConstShape::F64 => quote!(polydat::ast::ConstValue::F64(#field_name)),
ConstShape::Bool => quote!(polydat::ast::ConstValue::U64(if #field_name { 1 } else { 0 })),
ConstShape::Str => quote!(polydat::ast::ConstValue::Str(#field_name.clone())),
};
slot_exprs.push(quote! {
polydat::ast::Slot::Const {
name: #name_str.into(),
value: #const_value_ctor,
}
});
}
ArgKind::Setup(_) => {
}
ArgKind::PolyWire => {
let pt_param = format_ident!("{}_type", a.name);
slot_exprs.push(quote! {
polydat::ast::Slot::Wire(polydat::ast::Port::new(
#name_str, #pt_param))
});
}
ArgKind::Variadic(_) => {
}
ArgKind::ConstVec(inner) => {
let field_name = &a.name;
match inner {
ConstShape::U64 => slot_exprs.push(quote! {
polydat::ast::Slot::Const {
name: #name_str.into(),
value: polydat::ast::ConstValue::VecU64(#field_name.clone()),
}
}),
ConstShape::F64 => slot_exprs.push(quote! {
polydat::ast::Slot::Const {
name: #name_str.into(),
value: polydat::ast::ConstValue::VecF64(#field_name.clone()),
}
}),
_ => {}
}
}
}
}
let variadic_slot_extends: Vec<TokenStream2> = args.iter()
.filter_map(|a| match &a.kind {
ArgKind::Variadic(elem) => {
let name_str = a.name.to_string();
let pt = elem.port_type_tokens();
Some(quote! {
for __i in 0..n_wires {
ins.push(polydat::ast::Slot::Wire(
polydat::ast::Port::new(
format!("{}_{__i}", #name_str),
#pt,
)));
}
})
}
_ => None,
})
.collect();
let param_specs: Vec<TokenStream2> = args.iter()
.filter_map(|a| {
let name_str = a.name.to_string();
let required = match &a.kind {
ArgKind::Variadic(_) | ArgKind::ConstVec(_) => false,
_ => a.default_value.is_none(),
};
let slot_type = match &a.kind {
ArgKind::Wire | ArgKind::PolyWire | ArgKind::Variadic(_) => quote!(polydat::ast::SlotType::Wire),
ArgKind::Const(shape) => shape.slot_type_tokens(),
ArgKind::ConstVec(inner) => inner.slot_type_tokens(),
ArgKind::Setup(_) => return None,
};
Some(quote! {
polydat::dsl::registry::ParamSpec {
name: #name_str,
slot_type: #slot_type,
required: #required,
example: #name_str,
constraint: None,
}
})
})
.collect();
let declared_ret_ty = match &func.sig.output {
ReturnType::Default => {
return Err(syn::Error::new_spanned(
&func.sig,
"#[polydat_node] requires an explicit return type; \
nodes always produce a value.",
));
}
ReturnType::Type(_, t) => (**t).clone(),
};
let fallible_inner_ty: Option<Type> = classify_result_return(&declared_ret_ty);
let is_fallible = fallible_inner_ty.is_some();
let ret_ty = fallible_inner_ty.clone().unwrap_or_else(|| declared_ret_ty.clone());
let ret_is_polywire = classify_polywire(&ret_ty);
if is_fallible {
for a in &args {
match &a.kind {
ArgKind::Wire | ArgKind::PolyWire | ArgKind::Variadic(_) => {
return Err(syn::Error::new_spanned(
&a.declared_ty,
"fallible-construction nodes (-> Result<T, E>) must \
have only Const args. Wire/PolyWire/variadic inputs \
can't be evaluated at construction time. Use the \
#[poly_const(setup_fn, from = ...)] shape instead \
when per-eval inputs are needed.",
));
}
ArgKind::Setup(_) | ArgKind::Const(_) | ArgKind::ConstVec(_) => {}
}
}
}
let tuple_ret_elems: Option<Vec<Type>> = match &ret_ty {
syn::Type::Tuple(t) => Some(t.elems.iter().cloned().collect()),
_ => None,
};
let dynamic_outputs_inner: Option<Type> = classify_dynamic_outputs(&ret_ty);
let dynamic_outputs_count_arg: Option<syn::Ident> = if dynamic_outputs_inner.is_some() {
let const_vec_args: Vec<&syn::Ident> = args.iter()
.filter_map(|a| match &a.kind {
ArgKind::ConstVec(_) => Some(&a.name),
_ => None,
})
.collect();
if const_vec_args.len() != 1 {
return Err(syn::Error::new_spanned(
&ret_ty,
format!(
"`DynamicOutputs<T>` return requires exactly one \
`Const<Vec<C>>` arg to drive the output port count \
(got {}). Declare one `Const<Vec<C>>` arg whose length \
determines the number of output ports.",
const_vec_args.len(),
),
));
}
Some(const_vec_args[0].clone())
} else {
None
};
if tuple_ret_elems.is_some() && ret_is_polywire {
return Err(syn::Error::new_spanned(
&ret_ty,
"tuple return + PolyWire don't compose (SameAsInput is a \
single-output dispatch).",
));
}
let first_polywire_idx: Option<usize> = args.iter()
.enumerate()
.find(|(_, a)| matches!(a.kind, ArgKind::PolyWire))
.map(|(i, _)| i);
let output_port_types: Vec<TokenStream2> = if let Some(elems) = &tuple_ret_elems {
elems.iter()
.map(wire_port_type_for)
.collect::<syn::Result<Vec<_>>>()?
} else if ret_is_polywire {
if let Some(polywire_arg) = args.iter().find(|a| matches!(a.kind, ArgKind::PolyWire)) {
let pt_ident = format_ident!("{}_type", polywire_arg.name);
vec![quote!(#pt_ident)]
} else if args.iter().any(|a| matches!(&a.kind, ArgKind::Variadic(VariadicElement::Value))) {
vec![quote!(polydat::ast::PortType::U64)]
} else {
return Err(syn::Error::new_spanned(
&ret_ty,
"function returns `Value` but has no `Value` arg — the macro \
needs at least one PolyWire (`Value`) arg or a `&[Value]` \
variadic to source the runtime port type for the output.",
));
}
} else if let Some(inner) = &dynamic_outputs_inner {
vec![wire_port_type_for(inner)?]
} else {
vec![wire_port_type_for(&ret_ty)?]
};
let output_names_strs: Vec<String> = match (&tuple_ret_elems, &attrs.output_names) {
(Some(elems), Some(names)) => {
if names.len() != elems.len() {
return Err(syn::Error::new_spanned(
&ret_ty,
format!(
"tuple return has {} elements but `output_names(...)` \
lists {}; lengths must match.",
elems.len(), names.len(),
),
));
}
names.iter().map(|n| n.to_string()).collect()
}
(Some(elems), None) => (0..elems.len()).map(|i| format!("out_{i}")).collect(),
(None, Some(names)) if names.len() != 1 => {
return Err(syn::Error::new_spanned(
&ret_ty,
"single-output return doesn't accept multi-name `output_names(...)`.",
));
}
(None, Some(names)) => vec![names[0].to_string()],
(None, None) => vec!["output".to_string()],
};
let output_port_field: TokenStream2 = if tuple_ret_elems.is_some()
|| ret_is_polywire
|| dynamic_outputs_inner.is_some()
{
quote!(None)
} else {
let pt = &output_port_types[0];
quote!(Some(#pt))
};
let output_count = if dynamic_outputs_inner.is_some() { 0 } else { output_port_types.len() };
let output_count_lit = syn::LitInt::new(&output_count.to_string(), proc_macro2::Span::call_site());
let output_type_tokens: TokenStream2 = match (ret_is_polywire, first_polywire_idx) {
(true, Some(idx)) => {
let i = syn::Index::from(idx);
quote!(polydat::dsl::registry::OutputType::SameAsInput(#i))
}
_ => quote!(polydat::dsl::registry::OutputType::Fixed),
};
let struct_fields: Vec<TokenStream2> = args.iter()
.filter_map(|a| match &a.kind {
ArgKind::Wire | ArgKind::PolyWire | ArgKind::Variadic(_) => None,
ArgKind::Const(shape) => {
let n = &a.name;
let ft = shape.field_type_tokens();
Some(quote!(pub #n: #ft))
}
ArgKind::ConstVec(inner) => {
let n = &a.name;
let ft = inner.field_type_tokens();
Some(quote!(pub #n: Vec<#ft>))
}
ArgKind::Setup(spec) => {
let n = &a.name;
let ty = &spec.inner_ty;
Some(quote!(pub #n: #ty))
}
})
.collect();
let new_params: Vec<TokenStream2> = args.iter()
.filter_map(|a| match &a.kind {
ArgKind::Wire => None,
ArgKind::Const(shape) => {
let n = &a.name;
let ft = shape.field_type_tokens();
Some(quote!(#n: #ft))
}
ArgKind::ConstVec(inner) => {
let n = &a.name;
let ft = inner.field_type_tokens();
Some(quote!(#n: Vec<#ft>))
}
ArgKind::Setup(_) => None,
ArgKind::PolyWire => {
let n = format_ident!("{}_type", a.name);
Some(quote!(#n: polydat::ast::PortType))
}
ArgKind::Variadic(_) => None,
})
.collect();
let has_variadic = args.iter().any(|a| matches!(a.kind, ArgKind::Variadic(_)));
let variadic_count = args.iter().filter(|a| matches!(a.kind, ArgKind::Variadic(_))).count();
if variadic_count > 2 {
return Err(syn::Error::new_spanned(
&func.sig,
"`#[polydat_node]` supports at most two variadic `&[T]` args (split-halves shape). \
Functions declaring more than two are not expressible in any SRD-80b shape.",
));
}
let is_split_halves = variadic_count == 2;
let variadic_positions: std::collections::HashMap<String, usize> = args.iter()
.filter(|a| matches!(a.kind, ArgKind::Variadic(_)))
.enumerate()
.map(|(i, a)| (a.name.to_string(), i))
.collect();
let new_params: Vec<TokenStream2> = if has_variadic {
let mut v = new_params;
v.push(quote!(n_wires: usize));
v
} else {
new_params
};
#[derive(Clone, Copy)]
enum ConstSourceShape {
ScalarValue,
ScalarStr,
VecValues,
}
let const_shape_by_name: std::collections::HashMap<String, ConstSourceShape> = args.iter()
.filter_map(|a| match &a.kind {
ArgKind::Const(ConstShape::Str) => Some((a.name.to_string(), ConstSourceShape::ScalarStr)),
ArgKind::Const(_) => Some((a.name.to_string(), ConstSourceShape::ScalarValue)),
ArgKind::ConstVec(_) => Some((a.name.to_string(), ConstSourceShape::VecValues)),
_ => None,
})
.collect();
let setup_precomputes: Vec<TokenStream2> = args.iter()
.filter_map(|a| match &a.kind {
ArgKind::Wire | ArgKind::Const(_) | ArgKind::ConstVec(_) | ArgKind::PolyWire | ArgKind::Variadic(_) => None,
ArgKind::Setup(spec) => {
let n = &a.name;
let setup_fn = &spec.setup_fn;
let mut src_exprs: Vec<TokenStream2> = Vec::new();
let mut err: Option<TokenStream2> = None;
for src in &spec.source_args {
let shape = const_shape_by_name.get(&src.to_string());
let expr = match shape {
Some(ConstSourceShape::ScalarStr) => quote!(#src.as_str()),
Some(ConstSourceShape::ScalarValue) => quote!(#src),
Some(ConstSourceShape::VecValues) => quote!(&#src),
None => {
err = Some(syn::Error::new(
src.span(),
format!(
"#[poly_const(... from = ... {src} ...)] — \
`{src}` is not declared as a `Const<T>` \
arg in the same function signature."),
).to_compile_error());
break;
}
};
src_exprs.push(expr);
}
if let Some(e) = err { return Some(e); }
let call = quote!(#setup_fn( #( #src_exprs ),* ));
Some(quote! {
let #n = #call;
})
}
})
.collect();
let new_field_inits: Vec<TokenStream2> = args.iter()
.filter_map(|a| match &a.kind {
ArgKind::Wire | ArgKind::PolyWire | ArgKind::Variadic(_) => None,
ArgKind::Const(_) | ArgKind::ConstVec(_) | ArgKind::Setup(_) => {
let n = &a.name;
Some(quote!(#n))
}
})
.collect();
let mut wire_idx = 0usize;
let arg_bindings: Vec<TokenStream2> = args.iter()
.map(|a| {
let n = &a.name;
match &a.kind {
ArgKind::Wire => {
let idx = syn::Index::from(wire_idx);
wire_idx += 1;
let ty = &a.declared_ty;
if classify_wrapper_wire(ty) == Some(WrapperWire::Handle) {
let inner = extract_handle_inner(ty)
.expect("Handle classification implies Arc<T> shape");
quote! {
let #n: std::sync::Arc<#inner> = match &inputs[#idx] {
polydat::ast::Value::Handle(arc) => arc.clone()
.downcast::<#inner>()
.expect("Handle type mismatch — wiring bug"),
other => panic!("expected Handle, got {other:?}"),
};
}
} else if let Some(borrow) = is_borrow_wire_shape(ty) {
let extract = borrow_extract_tokens(borrow, quote!(&inputs[#idx]));
quote! {
let #n = #extract;
}
} else {
quote! {
let #n = <#ty as polydat::derive_support::Wire>::extract(&inputs[#idx]);
}
}
}
ArgKind::Const(shape) => {
let wrap = shape.wrap_as_const(quote!(self.#n));
quote! {
let #n = #wrap;
}
}
ArgKind::Setup(_) => {
quote! {
let #n = &self.#n;
}
}
ArgKind::PolyWire => {
let idx = syn::Index::from(wire_idx);
wire_idx += 1;
quote! {
let #n: polydat::ast::Value = inputs[#idx].clone();
}
}
ArgKind::Variadic(elem) => {
let extractor = elem.extract_from_value();
let owned = format_ident!("__{}_owned", a.name);
let slice_expr = if is_split_halves {
let pos = variadic_positions[&a.name.to_string()];
if pos == 0 {
quote!({ let __half = inputs.len() / 2; &inputs[..__half] })
} else {
quote!({ let __half = inputs.len() / 2; &inputs[__half..] })
}
} else {
quote!(inputs)
};
quote! {
let #owned: Vec<_> = #slice_expr.iter().map(#extractor).collect();
let #n: &[_] = #owned.as_slice();
}
}
ArgKind::ConstVec(_) => {
quote! {
let #n = polydat::derive_support::Const(self.#n.clone());
}
}
}
})
.collect();
let mut const_idx_for_extract = 0usize;
let const_extracts: Vec<TokenStream2> = args.iter()
.filter_map(|a| match &a.kind {
ArgKind::Wire | ArgKind::Setup(_) | ArgKind::PolyWire | ArgKind::Variadic(_) => None,
ArgKind::Const(shape) => {
let n = &a.name;
let i = const_idx_for_extract;
const_idx_for_extract += 1;
let i_lit = syn::Index::from(i);
let extract_present = shape.extract_from_const_arg(quote!(c));
let fallback = match &a.default_value {
Some(default_expr) => {
match shape {
ConstShape::Str => quote!((#default_expr).to_string()),
_ => quote!(#default_expr),
}
}
None => {
let msg = format!(
"missing required const arg '{n}' for function '{func_name_str}'");
quote!(return Some(Err(#msg.to_string())))
}
};
Some(quote! {
let #n: _ = match consts.get(#i_lit) {
Some(c) => #extract_present,
None => #fallback,
};
})
}
ArgKind::ConstVec(inner) => {
let n = &a.name;
let i = const_idx_for_extract;
let i_lit = syn::LitInt::new(&i.to_string(), proc_macro2::Span::call_site());
let extract_one = inner.extract_from_const_arg(quote!(c));
Some(quote! {
let #n: Vec<_> = consts[#i_lit..].iter()
.map(|c| #extract_one)
.collect();
})
}
})
.collect();
let mut new_call_args: Vec<TokenStream2> = args.iter()
.filter_map(|a| match &a.kind {
ArgKind::Wire | ArgKind::Setup(_) | ArgKind::Variadic(_) => None,
ArgKind::Const(_) | ArgKind::ConstVec(_) => {
let n = &a.name;
Some(quote!(#n))
}
ArgKind::PolyWire => {
let n = format_ident!("{}_type", a.name);
Some(quote!(#n))
}
})
.collect();
if has_variadic {
new_call_args.push(quote!(n_wires));
}
let variadic_n_wires_extract: TokenStream2 = if has_variadic {
if is_split_halves {
quote! { let n_wires: usize = _wires.len() / 2; }
} else {
quote! { let n_wires: usize = _wires.len(); }
}
} else {
quote!()
};
let polywire_extracts: Vec<TokenStream2> = {
let mut wire_idx = 0usize;
let mut out = Vec::new();
for a in &args {
match &a.kind {
ArgKind::Wire => { wire_idx += 1; }
ArgKind::Variadic(_) => {
wire_idx += 0; }
ArgKind::PolyWire => {
let pt_ident = format_ident!("{}_type", a.name);
let i = syn::Index::from(wire_idx);
let n_str = a.name.to_string();
let err = format!(
"polywire arg '{n_str}' for '{func_name_str}': assembler \
did not resolve a port type at wire index {wire_idx}");
out.push(quote! {
let #pt_ident: polydat::ast::PortType = match _wire_types.get(#i) {
Some(t) => *t,
None => return Some(Err(#err.to_string())),
};
});
wire_idx += 1;
}
ArgKind::Const(_) | ArgKind::ConstVec(_) | ArgKind::Setup(_) => {}
}
}
out
};
let block = &func.block;
let default_resolver_field: TokenStream2 = {
let wire_tys: Vec<&Type> = args.iter()
.filter_map(|a| match &a.kind {
ArgKind::Wire if is_borrow_wire_shape(&a.declared_ty).is_none()
&& classify_wrapper_wire(&a.declared_ty) != Some(WrapperWire::Handle)
=> Some(&a.declared_ty),
_ => None,
})
.collect();
if wire_tys.is_empty() {
quote!(None)
} else {
let mut acc = quote!(None);
for ty in wire_tys.iter().rev() {
acc = quote! {
match <#ty as polydat::derive_support::Wire>::RESOLVER {
Some(__r) => Some(__r),
None => #acc,
}
};
}
acc
}
};
let port_guard: TokenStream2 = if is_instantiation {
let mut wi: usize = 0;
let mut checks: Vec<TokenStream2> = Vec::new();
for a in &args {
match &a.kind {
ArgKind::Wire | ArgKind::PolyWire => {
let i = syn::Index::from(wi);
let ty = &a.declared_ty;
if !classify_polywire(ty) {
checks.push(quote! {
if _wire_types.get(#i) != Some(&<#ty as polydat::derive_support::Wire>::PORT) {
return None;
}
});
}
wi += 1;
}
ArgKind::Variadic(_) => {
}
_ => {}
}
}
quote! { #( #checks )* }
} else {
quote!()
};
let has_non_wire = args.iter().any(|a| !matches!(a.kind, ArgKind::Wire));
let default_impl = if has_non_wire {
quote!()
} else {
quote! {
impl Default for #struct_name {
fn default() -> Self { Self::new() }
}
}
};
let fused_node_impl: TokenStream2 = if let Some(path) = &attrs.decompose {
quote! {
impl polydat::compile::fusion::FusedNode for #struct_name {
fn decomposed(&self) -> polydat::compile::fusion::DecomposedGraph {
#path(self)
}
}
}
} else {
quote!()
};
let has_setup = args.iter().any(|a| matches!(a.kind, ArgKind::Setup(_)));
let ret_jit_type = wire_type_to_jit_type(&ret_ty);
let arg_jit_types: Option<Vec<JitType>> = if has_setup {
None
} else {
args.iter()
.map(|a| match &a.kind {
ArgKind::Wire => wire_type_to_jit_type(&a.declared_ty),
ArgKind::Const(shape) => const_shape_to_jit_type(*shape),
ArgKind::Setup(_) | ArgKind::PolyWire | ArgKind::ConstVec(_) => None,
ArgKind::Variadic(elem) => match elem {
VariadicElement::U64 => Some(JitType::U64),
_ => None,
},
})
.collect()
};
let tuple_ret_jit_types: Option<Vec<JitType>> = tuple_ret_elems.as_ref()
.and_then(|elems| {
elems.iter()
.map(wire_type_to_jit_type)
.collect::<Option<Vec<_>>>()
});
let jit_eligible = !is_fallible
&& arg_jit_types.is_some()
&& (ret_jit_type.is_some() || tuple_ret_jit_types.is_some());
let slice_arg_elem = |ty: &Type| -> Option<&'static str> {
match is_borrow_wire_shape(ty) {
Some(BorrowWire::Vec(variant, _)) => match variant {
"VecF32" => Some("F32"),
"VecF64" => Some("F64"),
"VecF16" => Some("F16"),
"VecI8" => Some("I8"),
"VecI16" => Some("I16"),
"VecI32" => Some("I32"),
"VecI64" => Some("I64"),
_ => None,
},
_ => None,
}
};
let vec_ret_elem: Option<&'static str> = {
let flat: String = type_to_string(&ret_ty).split_whitespace().collect();
match flat.as_str() {
"Vec<f32>" => Some("F32"),
"Vec<f64>" => Some("F64"),
"Vec<half::f16>" | "Vec<f16>" => Some("F16"),
"Vec<i8>" => Some("I8"),
"Vec<i16>" => Some("I16"),
"Vec<i32>" => Some("I32"),
"Vec<i64>" => Some("I64"),
_ => None,
}
};
enum SlotArgRead {
Jit(JitType),
Slice(&'static str),
Const,
}
let slot_arg_reads: Option<Vec<SlotArgRead>> = if has_setup
|| tuple_ret_elems.is_some()
|| is_fallible
{
None
} else {
args.iter()
.map(|a| match &a.kind {
ArgKind::Wire => wire_type_to_jit_type(&a.declared_ty)
.map(SlotArgRead::Jit)
.or_else(|| slice_arg_elem(&a.declared_ty).map(SlotArgRead::Slice)),
ArgKind::Const(_) => Some(SlotArgRead::Const),
_ => None,
})
.collect()
};
let has_slice_shape = slot_arg_reads
.as_ref()
.map(|v| v.iter().any(|r| matches!(r, SlotArgRead::Slice(_))))
.unwrap_or(false)
|| vec_ret_elem.is_some();
let slot_eligible = !jit_eligible
&& has_slice_shape
&& slot_arg_reads.is_some()
&& (ret_jit_type.is_some() || vec_ret_elem.is_some());
let emit_compiled_u64 = attrs.compiled_u64_override.is_some()
|| (jit_eligible && !attrs.no_jit);
let emit_jit_constants = attrs.jit_constants_override.is_some()
|| (jit_eligible && !attrs.no_jit);
let use_shared_body = (jit_eligible && (emit_compiled_u64 || !attrs.no_jit))
|| (slot_eligible && !attrs.no_jit);
let body_params: Vec<TokenStream2> = args.iter()
.map(|a| {
let n = &a.name;
let t = &a.declared_ty;
quote!(#n: #t)
})
.collect();
let body_fn_def: TokenStream2 = if is_fallible {
quote! {
#[inline(always)]
#[allow(unused_variables)]
fn __polydat_body( #( #body_params ),* ) -> #declared_ret_ty #block
}
} else if use_shared_body {
quote! {
#[inline(always)]
#[allow(unused_variables)]
fn __polydat_body( #( #body_params ),* ) -> #ret_ty #block
}
} else {
quote!()
};
let output_assign = |idx_lit: TokenStream2, elem_ty: &Type, local: TokenStream2| -> TokenStream2 {
if classify_wrapper_wire(elem_ty) == Some(WrapperWire::Handle) {
quote! {
outputs[#idx_lit] = polydat::ast::Value::handle(#local);
}
} else if classify_polywire(elem_ty) {
quote! {
outputs[#idx_lit] = #local;
}
} else if let Some(borrow) = is_borrow_wire_shape(elem_ty) {
match borrow {
BorrowWire::Str => quote! {
outputs[#idx_lit] = polydat::ast::Value::Str((#local).into());
},
BorrowWire::Bytes => quote! {
outputs[#idx_lit] = polydat::ast::Value::Bytes((#local).into());
},
BorrowWire::Json => quote! {
outputs[#idx_lit] = polydat::ast::Value::Json(::std::sync::Arc::new((#local).clone()));
},
BorrowWire::Vec(variant, _) => {
let v = syn::Ident::new(variant, proc_macro2::Span::call_site());
quote! {
outputs[#idx_lit] = polydat::ast::Value::#v(polydat::ast::SliceArc::from_vec((#local).to_vec()));
}
}
}
} else {
quote! {
outputs[#idx_lit] = <#elem_ty as polydat::derive_support::Wire>::inject(#local);
}
}
};
let outs_build: TokenStream2 = if let (Some(inner), Some(count_arg)) =
(&dynamic_outputs_inner, &dynamic_outputs_count_arg)
{
quote! {
let outs: Vec<polydat::ast::Port> = (0..#count_arg.len())
.map(|__i| polydat::ast::Port::new(
format!("d{}", __i),
<#inner as polydat::derive_support::Wire>::PORT,
))
.collect();
}
} else {
quote! {
let outs = vec![ #(
polydat::ast::Port::new(#output_names_strs, #output_port_types)
),* ];
}
};
let result_to_outputs: TokenStream2 = if let Some(inner) = &dynamic_outputs_inner {
let inject_one = if classify_polywire(inner) {
quote!(__elem)
} else if let Some(borrow) = is_borrow_wire_shape(inner) {
match borrow {
BorrowWire::Str => quote!(polydat::ast::Value::Str((__elem).into())),
BorrowWire::Bytes => quote!(polydat::ast::Value::Bytes((__elem).into())),
BorrowWire::Json => quote!(polydat::ast::Value::Json(::std::sync::Arc::new((__elem).clone()))),
BorrowWire::Vec(variant, _) => {
let v = syn::Ident::new(variant, proc_macro2::Span::call_site());
quote!(polydat::ast::Value::#v(polydat::ast::SliceArc::from_vec((__elem).to_vec())))
}
}
} else {
quote!(<#inner as polydat::derive_support::Wire>::inject(__elem))
};
quote! {
for (__i, __elem) in result.0.into_iter().enumerate() {
outputs[__i] = #inject_one;
}
}
} else if let Some(elems) = &tuple_ret_elems {
let locals: Vec<Ident> = (0..elems.len())
.map(|i| format_ident!("__r_{}", i))
.collect();
let writes: Vec<TokenStream2> = elems.iter().enumerate()
.map(|(i, elem_ty)| {
let local = &locals[i];
let idx = syn::Index::from(i);
output_assign(quote!(#idx), elem_ty, quote!(#local))
})
.collect();
quote! {
let ( #( #locals ),* ) = result;
#( #writes )*
}
} else {
output_assign(quote!(0), &ret_ty, quote!(result))
};
let eval_body: TokenStream2 = if use_shared_body {
let arg_names: Vec<&syn::Ident> = args.iter().map(|a| &a.name).collect();
quote! {
#[allow(unused_variables)]
{
#( #arg_bindings )*
let result: #ret_ty = Self::__polydat_body( #( #arg_names ),* );
#result_to_outputs
}
}
} else {
quote! {
#[allow(unused_variables)]
{
#( #arg_bindings )*
let result: #ret_ty = (|| #block)();
#result_to_outputs
}
}
};
let compiled_u64_impl: TokenStream2 = if let Some(path) = &attrs.compiled_u64_override {
quote! {
fn compiled_u64(&self) -> Option<polydat::ast::CompiledU64Op> {
Some(#path(self))
}
}
} else if jit_eligible && !attrs.no_jit {
let jit_types = arg_jit_types.as_ref().unwrap();
let mut wire_buf_idx = 0usize;
let captures: Vec<TokenStream2> = args.iter()
.filter_map(|a| match &a.kind {
ArgKind::Wire | ArgKind::Variadic(_) => None,
ArgKind::Const(_) => {
let n = &a.name;
Some(quote!(let #n = self.#n.clone();))
}
ArgKind::Setup(_) | ArgKind::PolyWire | ArgKind::ConstVec(_) => {
unreachable!("setup/polywire/constvec excludes JIT eligibility")
}
})
.collect();
let arg_reads: Vec<TokenStream2> = args.iter().zip(jit_types.iter())
.map(|(a, jt)| {
let n = &a.name;
let _ = jt;
match &a.kind {
ArgKind::Wire => {
let read = jt.read_from_u64_buffer(wire_buf_idx);
wire_buf_idx += jt.width();
quote!(let #n = #read;)
}
ArgKind::Const(shape) => {
if *shape == ConstShape::Str {
quote!(let #n = polydat::derive_support::Const(#n.as_str());)
} else {
quote!(let #n = polydat::derive_support::Const(#n);)
}
}
ArgKind::Variadic(_) => {
quote!(let #n: &[u64] = inputs;)
}
ArgKind::Setup(_) | ArgKind::PolyWire | ArgKind::ConstVec(_) => unreachable!(),
}
})
.collect();
let arg_names: Vec<&syn::Ident> = args.iter().map(|a| &a.name).collect();
let write = if let Some(tuple_jits) = &tuple_ret_jit_types {
let locals: Vec<Ident> = (0..tuple_jits.len())
.map(|i| format_ident!("__jit_r_{}", i))
.collect();
let mut out_off = 0usize;
let writes: Vec<TokenStream2> = tuple_jits.iter().enumerate()
.map(|(i, jt)| {
let local = &locals[i];
let w = jt.write_to_u64_buffer_at(out_off, quote!(#local));
out_off += jt.width();
w
})
.collect();
quote! {
let ( #( #locals ),* ) = result;
#( #writes )*
}
} else {
let ret_jit = ret_jit_type.unwrap();
ret_jit.write_to_u64_buffer(quote!(result))
};
quote! {
fn compiled_u64(&self) -> Option<polydat::ast::CompiledU64Op> {
#( #captures )*
Some(Box::new(move |inputs: &[u64], outputs: &mut [u64]| {
#( #arg_reads )*
let result: #ret_ty = Self::__polydat_body( #( #arg_names ),* );
#write
}))
}
}
} else {
quote!()
};
let compiled_slot_impl: TokenStream2 = if slot_eligible && !attrs.no_jit {
let reads_spec = slot_arg_reads.as_ref().unwrap();
let elem_ty_tokens = |elem: &str| -> TokenStream2 {
match elem {
"F32" => quote!(f32),
"F64" => quote!(f64),
"F16" => quote!(polydat::half::f16),
"I8" => quote!(i8),
"I16" => quote!(i16),
"I32" => quote!(i32),
"I64" => quote!(i64),
_ => unreachable!(),
}
};
let captures: Vec<TokenStream2> = args.iter()
.filter_map(|a| match &a.kind {
ArgKind::Const(_) => {
let n = &a.name;
Some(quote!(let #n = self.#n.clone();))
}
_ => None,
})
.collect();
let mut off = 0usize;
let arg_reads: Vec<TokenStream2> = args.iter().zip(reads_spec.iter())
.map(|(a, spec)| {
let n = &a.name;
match spec {
SlotArgRead::Jit(jt) => {
let read = jt.read_from_u64_buffer(off);
off += jt.width();
quote!(let #n = #read;)
}
SlotArgRead::Slice(elem) => {
let et = elem_ty_tokens(elem);
let i = syn::Index::from(off);
let i1 = syn::Index::from(off + 1);
off += 2;
quote! {
let #n: &[#et] = unsafe {
::core::slice::from_raw_parts(
inputs[#i] as usize as *const #et,
inputs[#i1] as usize,
)
};
}
}
SlotArgRead::Const => {
quote!(let #n = polydat::derive_support::Const(#n.clone());)
}
}
})
.collect();
let arg_names: Vec<&syn::Ident> = args.iter().map(|a| &a.name).collect();
let (scratch_decl, write) = if let Some(elem) = vec_ret_elem {
let se = syn::Ident::new(elem, proc_macro2::Span::call_site());
(
quote!(vec![polydat::ast::ScratchElem::#se]),
quote! {
let polydat::ast::ScratchBuf::#se(__buf) = &mut scratch[0] else {
unreachable!("scratch element type mismatch");
};
*__buf = result;
outputs[0] = __buf.as_ptr() as usize as u64;
outputs[1] = __buf.len() as u64;
},
)
} else {
let ret_jit = ret_jit_type.unwrap();
(quote!(vec![]), ret_jit.write_to_u64_buffer(quote!(result)))
};
quote! {
fn compiled_slot(&self) -> Option<polydat::ast::CompiledSlotKit> {
#( #captures )*
Some(polydat::ast::CompiledSlotKit {
scratch: #scratch_decl,
op: Box::new(move |inputs: &[u64], outputs: &mut [u64], scratch: &mut [polydat::ast::ScratchBuf]| {
#( #arg_reads )*
let result: #ret_ty = Self::__polydat_body( #( #arg_names ),* );
#write
}),
})
}
}
} else {
quote!()
};
let jit_constants_impl: TokenStream2 = if let Some(path) = &attrs.jit_constants_override {
quote! {
fn jit_constants(&self) -> Vec<u64> {
#path(self)
}
}
} else if emit_jit_constants {
let const_encodings: Vec<TokenStream2> = args.iter()
.filter_map(|a| match &a.kind {
ArgKind::Const(shape) => {
let jt = const_shape_to_jit_type(*shape)?;
let n = &a.name;
Some(jt.const_field_as_u64(quote!(self.#n)))
}
_ => None,
})
.collect();
quote! {
fn jit_constants(&self) -> Vec<u64> {
vec![ #( #const_encodings ),* ]
}
}
} else {
quote!()
};
let purity_impl: TokenStream2 = match &attrs.purity {
None => quote!(),
Some(syn::Expr::Path(p)) => {
let variant = &p.path;
quote! {
fn purity(&self) -> polydat::ast::Purity {
polydat::ast::Purity::#variant
}
}
}
Some(syn::Expr::Call(c)) => {
let syn::Expr::Path(head_path) = &*c.func else {
return Err(syn::Error::new_spanned(
&c.func,
"purity call-form expects a Purity variant ident as the head.",
));
};
let head_ident = head_path.path.get_ident().ok_or_else(|| {
syn::Error::new_spanned(
&c.func,
"purity call-form head must be a single Purity variant ident.",
)
})?;
let arg = c.args.first().ok_or_else(|| {
syn::Error::new_spanned(
c,
"purity call-form requires one argument.")
})?;
match head_ident.to_string().as_str() {
"SideChannel" => quote! {
fn purity(&self) -> polydat::ast::Purity {
polydat::ast::Purity::SideChannel {
sink: polydat::ast::SideChannelSink::#arg,
}
}
},
"Nondeterministic" => quote! {
fn purity(&self) -> polydat::ast::Purity {
polydat::ast::Purity::Nondeterministic { reason: #arg }
}
},
other => return Err(syn::Error::new_spanned(
head_ident,
format!(
"purity call-form head `{other}` not recognized. \
Use `SideChannel(<sink>)` or `Nondeterministic(<reason>)`."),
)),
}
}
Some(other) => {
return Err(syn::Error::new_spanned(
other,
"purity attribute must be a Purity variant path or call form",
));
}
};
let _ = emit_compiled_u64;
let identity_field: TokenStream2 = if let Some(expr) = &attrs.identity {
quote!(Some(#expr))
} else {
quote!(None)
};
let has_const_arg = args.iter().any(|a| matches!(a.kind, ArgKind::Const(_)));
let has_polywire = args.iter().any(|a| matches!(a.kind, ArgKind::PolyWire));
let variadic_ctor_field: TokenStream2 = if has_variadic && !has_const_arg && !has_polywire {
if is_split_halves {
quote!(Some(|n| Box::new(#struct_name::new(n / 2))))
} else {
quote!(Some(|n| Box::new(#struct_name::new(n))))
}
} else {
quote!(None)
};
let has_none_aware_arg = args.iter().any(|a| match &a.kind {
ArgKind::Wire => is_option_arg(&a.declared_ty),
ArgKind::PolyWire => true,
_ => false,
});
let accepts_none_impl: TokenStream2 = if has_none_aware_arg {
quote! {
fn accepts_none_inputs(&self) -> bool { true }
}
} else {
quote!()
};
let has_const_vec = args.iter().any(|a| matches!(a.kind, ArgKind::ConstVec(_)));
let arity_field: TokenStream2 = if has_variadic {
let min_wires = match (&attrs.variadic_min, is_split_halves) {
(Some(v), true) => quote!(2 * (#v)),
(Some(v), false) => quote!(#v),
(None, _) => quote!(0),
};
quote!(polydat::dsl::registry::Arity::VariadicWires { min_wires: #min_wires })
} else if has_const_vec {
quote!(polydat::dsl::registry::Arity::VariadicConsts { min_consts: 0 })
} else {
quote!(polydat::dsl::registry::Arity::Fixed)
};
let commutativity_field: TokenStream2 = if let Some(c) = &attrs.commutativity {
quote!(polydat::ast::Commutativity::#c)
} else {
quote!(polydat::ast::Commutativity::Positional)
};
let (ctor_emission, eval_emission, build_call_emission): (TokenStream2, TokenStream2, TokenStream2) = if is_fallible {
let body_arg_passes: Vec<TokenStream2> = args.iter()
.map(|a| {
let n = &a.name;
match &a.kind {
ArgKind::Const(shape) => shape.wrap_as_const(quote!(#n)),
ArgKind::Setup(_) => quote!(&#n),
_ => quote!(#n),
}
})
.collect();
let try_new = quote! {
pub fn try_new( #( #new_params ),* ) -> ::std::result::Result<Self, String> {
#( #setup_precomputes )*
let mut ins: Vec<polydat::ast::Slot> = vec![ #( #slot_exprs ),* ];
#( #variadic_slot_extends )*
#outs_build
let __polydat_cached = match Self::__polydat_body( #( #body_arg_passes ),* ) {
Ok(v) => v,
Err(e) => return Err(Into::<String>::into(e)),
};
Ok(Self {
meta: polydat::ast::NodeMeta {
name: #func_name_str.into(),
ins,
outs,
},
#( #new_field_inits, )*
__polydat_cached,
})
}
};
let out_assign = output_assign(quote!(0), &ret_ty, quote!(self.__polydat_cached.clone()));
let ev = quote! {
#[allow(unused_variables)]
{ #out_assign }
};
let bc = quote! {
Some(match #struct_name::try_new( #( #new_call_args ),* ) {
Ok(n) => Ok(Box::new(n) as Box<dyn polydat::ast::PolydatNode>),
Err(e) => Err(e),
})
};
(try_new, ev, bc)
} else {
let ctor = quote! {
pub fn new( #( #new_params ),* ) -> Self {
#( #setup_precomputes )*
let mut ins: Vec<polydat::ast::Slot> = vec![ #( #slot_exprs ),* ];
#( #variadic_slot_extends )*
#outs_build
Self {
meta: polydat::ast::NodeMeta {
name: #func_name_str.into(),
ins,
outs,
},
#( #new_field_inits, )*
}
}
};
let ev = quote!(#eval_body);
let bc = quote! {
Some(match ::std::panic::catch_unwind(
::std::panic::AssertUnwindSafe(|| #struct_name::new( #( #new_call_args ),* ))
) {
Ok(node) => Ok(Box::new(node) as Box<dyn polydat::ast::PolydatNode>),
Err(panic) => {
let msg = panic.downcast_ref::<&str>().copied()
.or_else(|| panic.downcast_ref::<String>().map(|s| s.as_str()))
.unwrap_or("<non-string panic>");
Err(format!("{}: construction failed: {}", #func_name_str, msg))
}
})
};
(ctor, ev, bc)
};
let cached_field: TokenStream2 = if is_fallible {
quote!(__polydat_cached: #ret_ty,)
} else {
quote!()
};
let result = quote! {
pub struct #struct_name {
meta: polydat::ast::NodeMeta,
#( #struct_fields, )*
#cached_field
}
#default_impl
#fused_node_impl
impl #struct_name {
#ctor_emission
#body_fn_def
}
impl polydat::ast::PolydatNode for #struct_name {
fn meta(&self) -> &polydat::ast::NodeMeta { &self.meta }
fn eval(
&self,
inputs: &[polydat::ast::Value],
outputs: &mut [polydat::ast::Value],
) {
#eval_emission
}
#compiled_u64_impl
#compiled_slot_impl
#jit_constants_impl
#purity_impl
#accepts_none_impl
}
const _: () = {
static SIGS: &[polydat::dsl::registry::FuncSig] = &[
polydat::dsl::registry::FuncSig {
name: #func_name_str,
category: polydat::dsl::registry::FuncCategory::#category,
outputs: #output_count_lit,
description: "",
help: "",
identity: #identity_field,
variadic_ctor: #variadic_ctor_field,
params: &[ #( #param_specs ),* ],
arity: #arity_field,
commutativity: #commutativity_field,
default_resolver: #default_resolver_field,
output_type: #output_type_tokens,
output_port: #output_port_field,
},
];
fn signatures() -> &'static [polydat::dsl::registry::FuncSig] { SIGS }
fn build(
name: &str,
_wires: &[polydat::compile::assembly::WireRef],
_wire_types: &[polydat::ast::PortType],
consts: &[polydat::dsl::factory::ConstArg],
) -> Option<Result<Box<dyn polydat::ast::PolydatNode>, String>> {
if name != #func_name_str { return None; }
#port_guard
#( #const_extracts )*
#( #polywire_extracts )*
#variadic_n_wires_extract
#build_call_emission
}
::polydat::inventory::submit! {
polydat::dsl::registry::NodeRegistration {
signatures,
build,
validate: None,
}
}
};
};
Ok(result)
}
fn to_camel_case(s: &str) -> String {
let mut out = String::with_capacity(s.len());
let mut up = true;
for c in s.chars() {
if c == '_' { up = true; continue; }
if up { out.extend(c.to_uppercase()); up = false; }
else { out.push(c); }
}
out
}
fn type_to_string(ty: &Type) -> String {
use quote::ToTokens;
let mut s = String::new();
for t in ty.to_token_stream() {
s.push_str(&t.to_string());
s.push(' ');
}
s.trim().to_string()
}