Skip to main content

wingfoil_derive/
lib.rs

1use proc_macro::TokenStream;
2use proc_macro2::Span;
3use quote::quote;
4use syn::{
5    Attribute, Ident, ImplItem, ImplItemFn, ItemImpl, LitStr, Token, Type, Visibility, braced,
6    bracketed,
7    parse::{Parse, ParseStream},
8    parse_macro_input,
9    punctuated::Punctuated,
10};
11
12struct NodeArgs {
13    active: Vec<Ident>,
14    passive: Vec<Ident>,
15    output: Option<(Ident, Type)>,
16}
17
18impl Parse for NodeArgs {
19    fn parse(input: ParseStream) -> syn::Result<Self> {
20        let mut active: Option<Vec<Ident>> = None;
21        let mut passive: Option<Vec<Ident>> = None;
22        let mut output: Option<(Ident, Type)> = None;
23
24        while !input.is_empty() {
25            let key: Ident = input.parse()?;
26            input.parse::<Token![=]>()?;
27
28            match key.to_string().as_str() {
29                "active" => {
30                    if active.is_some() {
31                        return Err(syn::Error::new(key.span(), "duplicate key `active`"));
32                    }
33                    let content;
34                    bracketed!(content in input);
35                    let list = Punctuated::<Ident, Token![,]>::parse_terminated(&content)?;
36                    active = Some(list.into_iter().collect());
37                }
38                "passive" => {
39                    if passive.is_some() {
40                        return Err(syn::Error::new(key.span(), "duplicate key `passive`"));
41                    }
42                    let content;
43                    bracketed!(content in input);
44                    let list = Punctuated::<Ident, Token![,]>::parse_terminated(&content)?;
45                    passive = Some(list.into_iter().collect());
46                }
47                "output" => {
48                    if output.is_some() {
49                        return Err(syn::Error::new(key.span(), "duplicate key `output`"));
50                    }
51                    let field: Ident = input.parse()?;
52                    input.parse::<Token![:]>()?;
53                    let ty: Type = input.parse()?;
54                    output = Some((field, ty));
55                }
56                _ => {
57                    return Err(syn::Error::new(
58                        key.span(),
59                        format!("unknown key `{key}`; expected `active`, `passive`, or `output`"),
60                    ));
61                }
62            }
63
64            if input.peek(Token![,]) {
65                input.parse::<Token![,]>()?;
66            }
67        }
68
69        Ok(NodeArgs {
70            active: active.unwrap_or_default(),
71            passive: passive.unwrap_or_default(),
72            output,
73        })
74    }
75}
76
77/// Attribute macro that reduces boilerplate in [`MutableNode`] implementations.
78///
79/// Place `#[node(...)]` on an `impl MutableNode for MyType` block. It can:
80///
81/// - Inject `fn upstreams()` from `active = [field1, field2]` and/or `passive = [field3]`
82/// - Emit a separate `impl StreamPeekRef<T>` from `output = field_name: FieldType`
83///
84/// Fields listed as `active` or `passive` must implement `AsUpstreamNodes`
85/// (`Rc<dyn Node>`, `Rc<dyn Stream<T>>`, or `Vec` of either).
86///
87/// If neither `active` nor `passive` is specified, no `upstreams()` is injected —
88/// the default `UpStreams::none()` from [`MutableNode`] is used (source node), or you
89/// can write `upstreams()` manually in the impl block for complex cases (e.g. `Dep<T>`).
90///
91/// # Examples
92///
93/// ```rust,ignore
94/// // Transform node: active upstream, stream output
95/// #[node(active = [upstream], output = value: OUT)]
96/// impl<IN, OUT: Element> MutableNode for MapStream<IN, OUT> {
97///     fn cycle(&mut self, _state: &mut GraphState) -> anyhow::Result<bool> {
98///         self.value = (self.func)(self.upstream.peek_value());
99///         Ok(true)
100///     }
101/// }
102///
103/// // Sink node: active upstream, no output
104/// #[node(active = [upstream])]
105/// impl<IN> MutableNode for ConsumerNode<IN> {
106///     fn cycle(&mut self, state: &mut GraphState) -> anyhow::Result<bool> { ... }
107/// }
108///
109/// // Source node with output (upstreams default to none)
110/// #[node(output = value: T)]
111/// impl<T: Element> MutableNode for ConstantStream<T> {
112///     fn cycle(&mut self, _state: &mut GraphState) -> anyhow::Result<bool> { Ok(true) }
113/// }
114///
115/// // Mixed: passive + active upstreams, output
116/// #[node(passive = [upstream], active = [trigger], output = value: T)]
117/// impl<T: Element> MutableNode for SampleStream<T> { ... }
118///
119/// // Complex upstreams (Dep<T> etc.) — write upstreams() manually, use #[node] for output only
120/// #[node(output = value: OUT)]
121/// impl<IN1: 'static, IN2: 'static, OUT: Element> MutableNode for BiMapStream<IN1, IN2, OUT> {
122///     fn cycle(&mut self, ...) -> anyhow::Result<bool> { ... }
123///     fn upstreams(&self) -> UpStreams { /* Dep<T> logic */ }
124/// }
125/// ```
126#[proc_macro_attribute]
127pub fn node(attr: TokenStream, item: TokenStream) -> TokenStream {
128    let args = parse_macro_input!(attr as NodeArgs);
129    let mut impl_block = parse_macro_input!(item as ItemImpl);
130
131    let self_ty = impl_block.self_ty.clone();
132    let (impl_generics, _, where_clause) = impl_block.generics.split_for_impl();
133
134    // Inject fn upstreams() if active/passive fields are specified.
135    if !args.active.is_empty() || !args.passive.is_empty() {
136        let active_fields = &args.active;
137        let passive_fields = &args.passive;
138
139        // All wingfoil paths are fully qualified so the generated code is
140        // hygienic against user-defined traits/types of the same name.
141        // Inside the wingfoil crate itself, `::wingfoil` resolves via
142        // `extern crate self as wingfoil;` in lib.rs.
143        let upstreams_fn: ImplItemFn = syn::parse_quote! {
144            fn upstreams(&self) -> ::wingfoil::UpStreams {
145                let mut active: ::std::vec::Vec<::std::rc::Rc<dyn ::wingfoil::Node>> = ::std::vec::Vec::new();
146                let mut passive: ::std::vec::Vec<::std::rc::Rc<dyn ::wingfoil::Node>> = ::std::vec::Vec::new();
147                #(active.extend(::wingfoil::AsUpstreamNodes::as_upstream_nodes(&self.#active_fields));)*
148                #(passive.extend(::wingfoil::AsUpstreamNodes::as_upstream_nodes(&self.#passive_fields));)*
149                ::wingfoil::UpStreams::new(active, passive)
150            }
151        };
152        impl_block.items.push(ImplItem::Fn(upstreams_fn));
153    }
154
155    // Emit a StreamPeekRef impl if output is specified.
156    let peek_ref_impl = args.output.map(|(field, ty)| {
157        quote! {
158            impl #impl_generics ::wingfoil::StreamPeekRef<#ty> for #self_ty #where_clause {
159                fn peek_ref(&self) -> &#ty {
160                    &self.#field
161                }
162            }
163        }
164    });
165
166    quote! {
167        #impl_block
168        #peek_ref_impl
169    }
170    .into()
171}
172
173// =============================================================================
174// latency_stages! — declare a fixed-size, named-field latency record.
175// =============================================================================
176
177struct LatencyStagesInput {
178    visibility: Visibility,
179    name: Ident,
180    stages: Vec<Ident>,
181    type_name_override: Option<LitStr>,
182}
183
184impl Parse for LatencyStagesInput {
185    fn parse(input: ParseStream) -> syn::Result<Self> {
186        let attrs = input.call(Attribute::parse_outer)?;
187        let mut type_name_override: Option<LitStr> = None;
188        for attr in &attrs {
189            if attr.path().is_ident("type_name") {
190                if type_name_override.is_some() {
191                    return Err(syn::Error::new_spanned(
192                        attr,
193                        "duplicate #[type_name(...)] attribute",
194                    ));
195                }
196                let lit: LitStr = attr.parse_args().map_err(|_| {
197                    syn::Error::new_spanned(
198                        attr,
199                        "expected #[type_name(\"...\")] with a single string literal",
200                    )
201                })?;
202                type_name_override = Some(lit);
203            } else {
204                return Err(syn::Error::new_spanned(
205                    attr,
206                    "unrecognized attribute on latency_stages!; only #[type_name(\"...\")] is supported",
207                ));
208            }
209        }
210
211        let visibility: Visibility = input.parse()?;
212        let name: Ident = input.parse()?;
213        let content;
214        braced!(content in input);
215        let list = Punctuated::<Ident, Token![,]>::parse_terminated(&content)?;
216        let stages: Vec<Ident> = list.into_iter().collect();
217        if stages.is_empty() {
218            return Err(syn::Error::new(
219                name.span(),
220                "latency_stages! requires at least one stage",
221            ));
222        }
223        Ok(LatencyStagesInput {
224            visibility,
225            name,
226            stages,
227            type_name_override,
228        })
229    }
230}
231
232/// Convert `PascalCase` to `snake_case`. Used to derive the per-struct stage module name.
233fn pascal_to_snake(s: &str) -> String {
234    let mut out = String::with_capacity(s.len() + 4);
235    for (i, ch) in s.chars().enumerate() {
236        if ch.is_ascii_uppercase() {
237            if i != 0 {
238                out.push('_');
239            }
240            out.push(ch.to_ascii_lowercase());
241        } else {
242            out.push(ch);
243        }
244    }
245    out
246}
247
248/// Declare a fixed-size, named-field latency record suitable for embedding in
249/// a `#[repr(C)]` payload. Each field is a `u64` nanosecond timestamp.
250///
251/// The macro generates:
252/// - A `#[repr(C)]` struct with one `pub <field>: u64` per stage.
253/// - An `impl wingfoil::Latency` for the struct (gives slice access by index).
254/// - A nested module with one zero-sized marker per stage, each implementing
255///   `wingfoil::Stage<Self>` with the stage's compile-time index. Use those
256///   markers as the type parameter to `.stamp::<S>()`.
257///
258/// # Example
259///
260/// ```rust,ignore
261/// use wingfoil::*;
262///
263/// latency_stages! {
264///     pub TradeLatency {
265///         ingest,
266///         decode,
267///         strategy,
268///         publish,
269///     }
270/// }
271///
272/// // Markers live in a snake_case sub-module named after the struct:
273/// use trade_latency::strategy;
274/// let stamped = upstream.stamp::<strategy>();
275/// ```
276#[proc_macro]
277pub fn latency_stages(item: TokenStream) -> TokenStream {
278    let input = parse_macro_input!(item as LatencyStagesInput);
279    let LatencyStagesInput {
280        visibility,
281        name,
282        stages,
283        type_name_override,
284    } = input;
285
286    let n = stages.len();
287    let module_name = Ident::new(&pascal_to_snake(&name.to_string()), Span::call_site());
288    let stage_strs: Vec<String> = stages.iter().map(|i| i.to_string()).collect();
289    let stage_indices: Vec<usize> = (0..n).collect();
290    let field_names = &stages;
291    let marker_names = &stages;
292    let zero_copy_send_body = match type_name_override {
293        Some(lit) => quote! {
294            unsafe fn type_name() -> &'static str { #lit }
295        },
296        None => quote! {},
297    };
298
299    let expanded = quote! {
300        #[repr(C)]
301        #[derive(
302            ::std::clone::Clone, ::std::marker::Copy,
303            ::std::fmt::Debug, ::std::default::Default,
304            ::std::cmp::PartialEq, ::std::cmp::Eq,
305            ::std::hash::Hash,
306            ::serde::Serialize, ::serde::Deserialize,
307        )]
308        #visibility struct #name {
309            #( pub #field_names: u64, )*
310        }
311
312        impl Latency for #name {
313            const N: usize = #n;
314            fn stage_names() -> &'static [&'static str] {
315                &[ #( #stage_strs ),* ]
316            }
317            #[inline]
318            fn stamps(&self) -> &[u64] {
319                // SAFETY: `#[repr(C)]` struct of N consecutive u64 fields has
320                // the same memory layout as `[u64; N]`.
321                unsafe {
322                    ::std::slice::from_raw_parts(
323                        self as *const Self as *const u64,
324                        <Self as Latency>::N,
325                    )
326                }
327            }
328            #[inline]
329            fn stamp_mut(&mut self, idx: usize) -> &mut u64 {
330                assert!(idx < <Self as Latency>::N, "stage index out of bounds");
331                // SAFETY: see `stamps`; idx is bounds-checked above.
332                unsafe { &mut *((self as *mut Self as *mut u64).add(idx)) }
333            }
334        }
335
336        // SAFETY: `#[repr(C)]` packed `u64` fields are self-contained and
337        // have a uniform memory representation, satisfying `ZeroCopySend`'s
338        // invariants. Only emitted when the `iceoryx2` feature is on
339        // in the consuming crate.
340        #[cfg(feature = "iceoryx2")]
341        unsafe impl ::iceoryx2::prelude::ZeroCopySend for #name {
342            #zero_copy_send_body
343        }
344
345        #[allow(non_snake_case, non_camel_case_types)]
346        #visibility mod #module_name {
347            use super::*;
348            #(
349                /// Compile-time marker for a single latency stage.
350                pub struct #marker_names;
351                impl Stage<super::#name> for #marker_names {
352                    const NAME: &'static str = #stage_strs;
353                    const INDEX: usize = #stage_indices;
354                }
355            )*
356        }
357    };
358
359    expanded.into()
360}