Skip to main content

archmage_macros/
lib.rs

1//! Proc-macros for archmage SIMD capability tokens.
2//!
3//! Provides `#[arcane]` attribute (with `#[arcane]` alias) to make raw intrinsics
4//! safe via token proof.
5
6use proc_macro::TokenStream;
7use quote::{format_ident, quote, ToTokens};
8use syn::{
9    fold::Fold,
10    parse::{Parse, ParseStream},
11    parse_macro_input, parse_quote, Attribute, FnArg, GenericParam, Ident, ItemFn, PatType,
12    ReturnType, Signature, Token, Type, TypeParamBound,
13};
14
15/// A Fold implementation that replaces `Self` with a concrete type.
16struct ReplaceSelf<'a> {
17    replacement: &'a Type,
18}
19
20impl Fold for ReplaceSelf<'_> {
21    fn fold_type(&mut self, ty: Type) -> Type {
22        match ty {
23            Type::Path(ref type_path) if type_path.qself.is_none() => {
24                // Check if it's just `Self`
25                if type_path.path.is_ident("Self") {
26                    return self.replacement.clone();
27                }
28                // Otherwise continue folding
29                syn::fold::fold_type(self, ty)
30            }
31            _ => syn::fold::fold_type(self, ty),
32        }
33    }
34}
35
36/// Arguments to the `#[arcane]` macro.
37#[derive(Default)]
38struct ArcaneArgs {
39    /// Use `#[inline(always)]` instead of `#[inline]` for the inner function.
40    /// Requires nightly Rust with `#![feature(target_feature_inline_always)]`.
41    inline_always: bool,
42    /// The concrete type to use for `self` receiver.
43    /// When specified, `self`/`&self`/`&mut self` is transformed to `_self: Type`/`&Type`/`&mut Type`.
44    self_type: Option<Type>,
45}
46
47impl Parse for ArcaneArgs {
48    fn parse(input: ParseStream) -> syn::Result<Self> {
49        let mut args = ArcaneArgs::default();
50
51        while !input.is_empty() {
52            let ident: Ident = input.parse()?;
53            match ident.to_string().as_str() {
54                "inline_always" => args.inline_always = true,
55                "_self" => {
56                    let _: Token![=] = input.parse()?;
57                    args.self_type = Some(input.parse()?);
58                }
59                other => {
60                    return Err(syn::Error::new(
61                        ident.span(),
62                        format!("unknown arcane argument: `{}`", other),
63                    ))
64                }
65            }
66            // Consume optional comma
67            if input.peek(Token![,]) {
68                let _: Token![,] = input.parse()?;
69            }
70        }
71
72        Ok(args)
73    }
74}
75
76// Token-to-features and trait-to-features mappings are generated from
77// token-registry.toml by xtask. Regenerate with: cargo run -p xtask -- generate
78mod generated;
79use generated::{token_to_arch, token_to_features, trait_to_features};
80
81/// Result of extracting token info from a type.
82enum TokenTypeInfo {
83    /// Concrete token type (e.g., `Avx2Token`)
84    Concrete(String),
85    /// impl Trait with the trait names (e.g., `impl HasX64V2`)
86    ImplTrait(Vec<String>),
87    /// Generic type parameter name (e.g., `T`)
88    Generic(String),
89}
90
91/// Extract token type information from a type.
92fn extract_token_type_info(ty: &Type) -> Option<TokenTypeInfo> {
93    match ty {
94        Type::Path(type_path) => {
95            // Get the last segment of the path (e.g., "Avx2Token" from "archmage::Avx2Token")
96            type_path.path.segments.last().map(|seg| {
97                let name = seg.ident.to_string();
98                // Check if it's a known concrete token type
99                if token_to_features(&name).is_some() {
100                    TokenTypeInfo::Concrete(name)
101                } else {
102                    // Might be a generic type parameter like `T`
103                    TokenTypeInfo::Generic(name)
104                }
105            })
106        }
107        Type::Reference(type_ref) => {
108            // Handle &Token or &mut Token
109            extract_token_type_info(&type_ref.elem)
110        }
111        Type::ImplTrait(impl_trait) => {
112            // Handle `impl HasX64V2` or `impl HasX64V2 + HasNeon`
113            let traits: Vec<String> = extract_trait_names_from_bounds(&impl_trait.bounds);
114            if traits.is_empty() {
115                None
116            } else {
117                Some(TokenTypeInfo::ImplTrait(traits))
118            }
119        }
120        _ => None,
121    }
122}
123
124/// Extract trait names from type param bounds.
125fn extract_trait_names_from_bounds(
126    bounds: &syn::punctuated::Punctuated<TypeParamBound, Token![+]>,
127) -> Vec<String> {
128    bounds
129        .iter()
130        .filter_map(|bound| {
131            if let TypeParamBound::Trait(trait_bound) = bound {
132                trait_bound
133                    .path
134                    .segments
135                    .last()
136                    .map(|seg| seg.ident.to_string())
137            } else {
138                None
139            }
140        })
141        .collect()
142}
143
144/// Look up a generic type parameter in the function's generics.
145fn find_generic_bounds(sig: &Signature, type_name: &str) -> Option<Vec<String>> {
146    // Check inline bounds first (e.g., `fn foo<T: HasX64V2>(token: T)`)
147    for param in &sig.generics.params {
148        if let GenericParam::Type(type_param) = param {
149            if type_param.ident == type_name {
150                let traits = extract_trait_names_from_bounds(&type_param.bounds);
151                if !traits.is_empty() {
152                    return Some(traits);
153                }
154            }
155        }
156    }
157
158    // Check where clause (e.g., `fn foo<T>(token: T) where T: HasX64V2`)
159    if let Some(where_clause) = &sig.generics.where_clause {
160        for predicate in &where_clause.predicates {
161            if let syn::WherePredicate::Type(pred_type) = predicate {
162                if let Type::Path(type_path) = &pred_type.bounded_ty {
163                    if let Some(seg) = type_path.path.segments.last() {
164                        if seg.ident == type_name {
165                            let traits = extract_trait_names_from_bounds(&pred_type.bounds);
166                            if !traits.is_empty() {
167                                return Some(traits);
168                            }
169                        }
170                    }
171                }
172            }
173        }
174    }
175
176    None
177}
178
179/// Convert trait names to features, collecting all features from all traits.
180fn traits_to_features(trait_names: &[String]) -> Option<Vec<&'static str>> {
181    let mut all_features = Vec::new();
182
183    for trait_name in trait_names {
184        if let Some(features) = trait_to_features(trait_name) {
185            for &feature in features {
186                if !all_features.contains(&feature) {
187                    all_features.push(feature);
188                }
189            }
190        }
191    }
192
193    if all_features.is_empty() {
194        None
195    } else {
196        Some(all_features)
197    }
198}
199
200/// Result of finding a token parameter in a function signature.
201struct TokenParamInfo {
202    /// The parameter identifier (e.g., `token`)
203    ident: Ident,
204    /// Target features to enable (e.g., `["avx2", "fma"]`)
205    features: Vec<&'static str>,
206    /// Target architecture (Some for concrete tokens, None for traits/generics)
207    target_arch: Option<&'static str>,
208    /// Concrete token type name (Some for concrete tokens, None for traits/generics)
209    token_type_name: Option<String>,
210}
211
212/// Find the first token parameter in a function signature.
213fn find_token_param(sig: &Signature) -> Option<TokenParamInfo> {
214    for arg in &sig.inputs {
215        match arg {
216            FnArg::Receiver(_) => {
217                // Self receivers (self, &self, &mut self) are not yet supported.
218                // The macro creates an inner function, and Rust's inner functions
219                // cannot have `self` parameters. Supporting this would require
220                // AST rewriting to replace `self` with a regular parameter.
221                // See the module docs for the workaround.
222                continue;
223            }
224            FnArg::Typed(PatType { pat, ty, .. }) => {
225                if let Some(info) = extract_token_type_info(ty) {
226                    let (features, arch, token_name) = match info {
227                        TokenTypeInfo::Concrete(ref name) => {
228                            let features = token_to_features(name).map(|f| f.to_vec());
229                            let arch = token_to_arch(name);
230                            (features, arch, Some(name.clone()))
231                        }
232                        TokenTypeInfo::ImplTrait(trait_names) => {
233                            (traits_to_features(&trait_names), None, None)
234                        }
235                        TokenTypeInfo::Generic(type_name) => {
236                            // Look up the generic parameter's bounds
237                            let features = find_generic_bounds(sig, &type_name)
238                                .and_then(|traits| traits_to_features(&traits));
239                            (features, None, None)
240                        }
241                    };
242
243                    if let Some(features) = features {
244                        // Extract parameter name
245                        if let syn::Pat::Ident(pat_ident) = pat.as_ref() {
246                            return Some(TokenParamInfo {
247                                ident: pat_ident.ident.clone(),
248                                features,
249                                target_arch: arch,
250                                token_type_name: token_name,
251                            });
252                        }
253                    }
254                }
255            }
256        }
257    }
258    None
259}
260
261/// Represents the kind of self receiver and the transformed parameter.
262enum SelfReceiver {
263    /// `self` (by value/move)
264    Owned,
265    /// `&self` (shared reference)
266    Ref,
267    /// `&mut self` (mutable reference)
268    RefMut,
269}
270
271/// Shared implementation for arcane/arcane macros.
272fn arcane_impl(input_fn: ItemFn, macro_name: &str, args: ArcaneArgs) -> TokenStream {
273    // Check for self receiver
274    let has_self_receiver = input_fn
275        .sig
276        .inputs
277        .first()
278        .map(|arg| matches!(arg, FnArg::Receiver(_)))
279        .unwrap_or(false);
280
281    // If there's a self receiver, we need _self = Type
282    if has_self_receiver && args.self_type.is_none() {
283        let msg = format!(
284            "{} with self receiver requires `_self = Type` argument.\n\
285             Example: #[{}(_self = MyType)]\n\
286             Use `_self` (not `self`) in the function body to refer to self.",
287            macro_name, macro_name
288        );
289        return syn::Error::new_spanned(&input_fn.sig, msg)
290            .to_compile_error()
291            .into();
292    }
293
294    // Find the token parameter, its features, target arch, and token type name
295    let TokenParamInfo {
296        ident: _token_ident,
297        features,
298        target_arch,
299        token_type_name,
300    } = match find_token_param(&input_fn.sig) {
301        Some(result) => result,
302        None => {
303            let msg = format!(
304                "{} requires a token parameter. Supported forms:\n\
305                 - Concrete: `token: X64V3Token`\n\
306                 - impl Trait: `token: impl Has256BitSimd`\n\
307                 - Generic: `fn foo<T: HasX64V2>(token: T, ...)`\n\
308                 - With self: `#[{}(_self = Type)] fn method(&self, token: impl HasNeon, ...)`",
309                macro_name, macro_name
310            );
311            return syn::Error::new_spanned(&input_fn.sig, msg)
312                .to_compile_error()
313                .into();
314        }
315    };
316
317    // Build target_feature attributes
318    let target_feature_attrs: Vec<Attribute> = features
319        .iter()
320        .map(|feature| parse_quote!(#[target_feature(enable = #feature)]))
321        .collect();
322
323    // Extract function components
324    let vis = &input_fn.vis;
325    let sig = &input_fn.sig;
326    let fn_name = &sig.ident;
327    let generics = &sig.generics;
328    let where_clause = &generics.where_clause;
329    let inputs = &sig.inputs;
330    let output = &sig.output;
331    let body = &input_fn.block;
332    let attrs = &input_fn.attrs;
333
334    // Determine self receiver type if present
335    let self_receiver_kind: Option<SelfReceiver> = inputs.first().and_then(|arg| match arg {
336        FnArg::Receiver(receiver) => {
337            if receiver.reference.is_none() {
338                Some(SelfReceiver::Owned)
339            } else if receiver.mutability.is_some() {
340                Some(SelfReceiver::RefMut)
341            } else {
342                Some(SelfReceiver::Ref)
343            }
344        }
345        _ => None,
346    });
347
348    // Build inner function parameters, transforming self if needed
349    let inner_params: Vec<proc_macro2::TokenStream> = inputs
350        .iter()
351        .map(|arg| match arg {
352            FnArg::Receiver(_) => {
353                // Transform self receiver to _self parameter
354                let self_ty = args.self_type.as_ref().unwrap();
355                match self_receiver_kind.as_ref().unwrap() {
356                    SelfReceiver::Owned => quote!(_self: #self_ty),
357                    SelfReceiver::Ref => quote!(_self: &#self_ty),
358                    SelfReceiver::RefMut => quote!(_self: &mut #self_ty),
359                }
360            }
361            FnArg::Typed(pat_type) => quote!(#pat_type),
362        })
363        .collect();
364
365    // Build inner function call arguments
366    let inner_args: Vec<proc_macro2::TokenStream> = inputs
367        .iter()
368        .filter_map(|arg| match arg {
369            FnArg::Typed(pat_type) => {
370                if let syn::Pat::Ident(pat_ident) = pat_type.pat.as_ref() {
371                    let ident = &pat_ident.ident;
372                    Some(quote!(#ident))
373                } else {
374                    None
375                }
376            }
377            FnArg::Receiver(_) => Some(quote!(self)), // Pass self to inner as _self
378        })
379        .collect();
380
381    let inner_fn_name = format_ident!("__simd_inner_{}", fn_name);
382
383    // Choose inline attribute based on args
384    // Note: #[inline(always)] + #[target_feature] requires nightly with
385    // #![feature(target_feature_inline_always)]
386    let inline_attr: Attribute = if args.inline_always {
387        parse_quote!(#[inline(always)])
388    } else {
389        parse_quote!(#[inline])
390    };
391
392    // Transform output and body to replace Self with concrete type if needed
393    let (inner_output, inner_body): (ReturnType, syn::Block) =
394        if let Some(ref self_ty) = args.self_type {
395            let mut replacer = ReplaceSelf {
396                replacement: self_ty,
397            };
398            let transformed_output = replacer.fold_return_type(output.clone());
399            let transformed_body = replacer.fold_block((**body).clone());
400            (transformed_output, transformed_body)
401        } else {
402            (output.clone(), (**body).clone())
403        };
404
405    // Generate the expanded function
406    // If we know the target arch (concrete token), generate cfg-gated real impl + stub
407    let token_type_str = token_type_name.as_deref().unwrap_or("UnknownToken");
408    let expanded = if let Some(arch) = target_arch {
409        quote! {
410            // Real implementation for the correct architecture
411            #[cfg(target_arch = #arch)]
412            #(#attrs)*
413            #vis #sig {
414                #(#target_feature_attrs)*
415                #inline_attr
416                fn #inner_fn_name #generics (#(#inner_params),*) #inner_output #where_clause
417                #inner_body
418
419                // SAFETY: The token parameter proves the required CPU features are available.
420                // Calling a #[target_feature] function from a non-matching context requires
421                // unsafe because the CPU may not support those instructions. The token's
422                // existence proves summon() succeeded, so the features are available.
423                unsafe { #inner_fn_name(#(#inner_args),*) }
424            }
425
426            // Stub for other architectures - the token cannot be obtained, so this is unreachable
427            #[cfg(not(target_arch = #arch))]
428            #(#attrs)*
429            #vis #sig {
430                // This token type cannot be summoned on this architecture.
431                // If you're seeing this at runtime, there's a bug in dispatch logic
432                // or forge_token_dangerously() was used incorrectly.
433                let _ = (#(#inner_args),*); // suppress unused warnings
434                unreachable!(
435                    "BUG: {}() was called but requires {} (target_arch = \"{}\"). \
436                     {}::summon() returns None on this architecture, so this function \
437                     is unreachable in safe code. If you used forge_token_dangerously(), \
438                     that is the bug.",
439                    stringify!(#fn_name),
440                    #token_type_str,
441                    #arch,
442                    #token_type_str,
443                )
444            }
445        }
446    } else {
447        // No specific arch (trait bounds or generic) - generate without cfg guards
448        quote! {
449            #(#attrs)*
450            #vis #sig {
451                #(#target_feature_attrs)*
452                #inline_attr
453                fn #inner_fn_name #generics (#(#inner_params),*) #inner_output #where_clause
454                #inner_body
455
456                // SAFETY: Calling a #[target_feature] function from a non-matching context
457                // requires unsafe. The token proves the required CPU features are available.
458                unsafe { #inner_fn_name(#(#inner_args),*) }
459            }
460        }
461    };
462
463    expanded.into()
464}
465
466/// Mark a function as an arcane SIMD function.
467///
468/// This macro enables safe use of SIMD intrinsics by generating an inner function
469/// with the appropriate `#[target_feature(enable = "...")]` attributes based on
470/// the token parameter type. The outer function calls the inner function unsafely,
471/// which is justified because the token parameter proves the features are available.
472///
473/// **The token is passed through to the inner function**, so you can call other
474/// token-taking functions from inside `#[arcane]`.
475///
476/// # Token Parameter Forms
477///
478/// The macro supports four forms of token parameters:
479///
480/// ## Concrete Token Types
481///
482/// ```ignore
483/// #[arcane]
484/// fn process(token: Avx2Token, data: &[f32; 8]) -> [f32; 8] {
485///     // AVX2 intrinsics safe here
486/// }
487/// ```
488///
489/// ## impl Trait Bounds
490///
491/// ```ignore
492/// #[arcane]
493/// fn process(token: impl HasX64V2, data: &[f32; 8]) -> [f32; 8] {
494///     // Accepts any token with x86-64-v2 features (SSE4.2+)
495/// }
496/// ```
497///
498/// ## Generic Type Parameters
499///
500/// ```ignore
501/// #[arcane]
502/// fn process<T: HasX64V2>(token: T, data: &[f32; 8]) -> [f32; 8] {
503///     // Generic over any v2-capable token
504/// }
505///
506/// // Also works with where clauses:
507/// #[arcane]
508/// fn process<T>(token: T, data: &[f32; 8]) -> [f32; 8]
509/// where
510///     T: HasX64V2
511/// {
512///     // ...
513/// }
514/// ```
515///
516/// ## Methods with Self Receivers
517///
518/// Methods with `self`, `&self`, `&mut self` receivers are supported via the
519/// `_self = Type` argument. Use `_self` in the function body instead of `self`:
520///
521/// ```ignore
522/// use archmage::{X64V3Token, arcane};
523/// use wide::f32x8;
524///
525/// trait SimdOps {
526///     fn double(&self, token: X64V3Token) -> Self;
527///     fn square(self, token: X64V3Token) -> Self;
528///     fn scale(&mut self, token: X64V3Token, factor: f32);
529/// }
530///
531/// impl SimdOps for f32x8 {
532///     #[arcane(_self = f32x8)]
533///     fn double(&self, _token: X64V3Token) -> Self {
534///         // Use _self instead of self in the body
535///         *_self + *_self
536///     }
537///
538///     #[arcane(_self = f32x8)]
539///     fn square(self, _token: X64V3Token) -> Self {
540///         _self * _self
541///     }
542///
543///     #[arcane(_self = f32x8)]
544///     fn scale(&mut self, _token: X64V3Token, factor: f32) {
545///         *_self = *_self * f32x8::splat(factor);
546///     }
547/// }
548/// ```
549///
550/// **Why `_self`?** The macro generates an inner function where `self` becomes
551/// a regular parameter named `_self`. Using `_self` in your code reminds you
552/// that you're not using the normal `self` keyword.
553///
554/// **All receiver types are supported:**
555/// - `self` (by value/move) → `_self: Type`
556/// - `&self` (shared reference) → `_self: &Type`
557/// - `&mut self` (mutable reference) → `_self: &mut Type`
558///
559/// # Multiple Trait Bounds
560///
561/// When using `impl Trait` or generic bounds with multiple traits,
562/// all required features are enabled:
563///
564/// ```ignore
565/// #[arcane]
566/// fn fma_kernel(token: impl HasX64V2 + HasNeon, data: &[f32; 8]) -> [f32; 8] {
567///     // Cross-platform: SSE4.2 on x86, NEON on ARM
568/// }
569/// ```
570///
571/// # Expansion
572///
573/// The macro expands to approximately:
574///
575/// ```ignore
576/// fn process(token: Avx2Token, data: &[f32; 8]) -> [f32; 8] {
577///     #[target_feature(enable = "avx2")]
578///     #[inline]
579///     fn __simd_inner_process(token: Avx2Token, data: &[f32; 8]) -> [f32; 8] {
580///         let v = unsafe { _mm256_loadu_ps(data.as_ptr()) };
581///         let doubled = _mm256_add_ps(v, v);
582///         let mut out = [0.0f32; 8];
583///         unsafe { _mm256_storeu_ps(out.as_mut_ptr(), doubled) };
584///         out
585///     }
586///     // SAFETY: Calling #[target_feature] fn from non-matching context.
587///     // Token proves the required features are available.
588///     unsafe { __simd_inner_process(token, data) }
589/// }
590/// ```
591///
592/// # Profile Tokens
593///
594/// Profile tokens automatically enable all required features:
595///
596/// ```ignore
597/// #[arcane]
598/// fn kernel(token: X64V3Token, data: &mut [f32]) {
599///     // AVX2 + FMA + BMI1 + BMI2 intrinsics all safe here!
600/// }
601/// ```
602///
603/// # Supported Tokens
604///
605/// - **x86_64 tiers**: `X64V2Token`, `X64V3Token` / `Desktop64` / `Avx2FmaToken`,
606///   `X64V4Token` / `Avx512Token` / `Server64`, `Avx512ModernToken`, `Avx512Fp16Token`
607/// - **ARM**: `NeonToken` / `Arm64`, `NeonAesToken`, `NeonSha3Token`, `NeonCrcToken`
608/// - **WASM**: `Wasm128Token`
609///
610/// # Supported Trait Bounds
611///
612/// - **x86_64 tiers**: `HasX64V2`, `HasX64V4`
613/// - **ARM**: `HasNeon`, `HasNeonAes`, `HasNeonSha3`
614///
615/// **Preferred:** Use concrete tokens (`X64V3Token`, `Desktop64`, `NeonToken`) directly.
616/// Concrete token types also work as trait bounds (e.g., `impl X64V3Token`).
617///
618/// # Options
619///
620/// ## `inline_always`
621///
622/// Use `#[inline(always)]` instead of `#[inline]` for the inner function.
623/// This can improve performance by ensuring aggressive inlining, but requires
624/// nightly Rust with `#![feature(target_feature_inline_always)]` enabled in
625/// the crate using the macro.
626///
627/// ```ignore
628/// #![feature(target_feature_inline_always)]
629///
630/// #[arcane(inline_always)]
631/// fn fast_kernel(token: Avx2Token, data: &mut [f32]) {
632///     // Inner function will use #[inline(always)]
633/// }
634/// ```
635#[proc_macro_attribute]
636pub fn arcane(attr: TokenStream, item: TokenStream) -> TokenStream {
637    let args = parse_macro_input!(attr as ArcaneArgs);
638    let input_fn = parse_macro_input!(item as ItemFn);
639    arcane_impl(input_fn, "arcane", args)
640}
641
642/// Legacy alias for [`arcane`].
643///
644/// **Deprecated:** Use `#[arcane]` instead. This alias exists only for migration.
645#[proc_macro_attribute]
646#[doc(hidden)]
647pub fn simd_fn(attr: TokenStream, item: TokenStream) -> TokenStream {
648    let args = parse_macro_input!(attr as ArcaneArgs);
649    let input_fn = parse_macro_input!(item as ItemFn);
650    arcane_impl(input_fn, "simd_fn", args)
651}
652
653// ============================================================================
654// Rite macro for inner SIMD functions (inlines into matching #[target_feature] callers)
655// ============================================================================
656
657/// Annotate inner SIMD helpers called from `#[arcane]` functions.
658///
659/// Unlike `#[arcane]`, which creates an inner `#[target_feature]` function behind
660/// a safe boundary, `#[rite]` adds `#[target_feature]` and `#[inline]` directly.
661/// LLVM inlines it into any caller with matching features — no boundary crossing.
662///
663/// # When to Use
664///
665/// Use `#[rite]` for helper functions that are **only** called from within
666/// `#[arcane]` functions with matching or superset token types:
667///
668/// ```ignore
669/// use archmage::{arcane, rite, X64V3Token};
670///
671/// #[arcane]
672/// fn outer(token: X64V3Token, data: &[f32; 8]) -> f32 {
673///     // helper inlines — same target features, no boundary
674///     helper(token, data) * 2.0
675/// }
676///
677/// #[rite]
678/// fn helper(token: X64V3Token, data: &[f32; 8]) -> f32 {
679///     // Just has #[target_feature(enable = "avx2,fma,...")]
680///     // Called from #[arcane] context, so features are guaranteed
681///     let v = f32x8::from_array(token, *data);
682///     v.reduce_add()
683/// }
684/// ```
685///
686/// # Safety
687///
688/// `#[rite]` functions can only be safely called from contexts where the
689/// required CPU features are enabled:
690/// - From within `#[arcane]` functions with matching/superset tokens
691/// - From within other `#[rite]` functions with matching/superset tokens
692/// - From code compiled with `-Ctarget-cpu` that enables the features
693///
694/// Calling from other contexts requires `unsafe` and the caller must ensure
695/// the CPU supports the required features.
696///
697/// # Comparison with #[arcane]
698///
699/// | Aspect | `#[arcane]` | `#[rite]` |
700/// |--------|-------------|-----------|
701/// | Creates wrapper | Yes | No |
702/// | Entry point | Yes | No |
703/// | Inlines into caller | No (barrier) | Yes |
704/// | Safe to call anywhere | Yes (with token) | Only from feature-enabled context |
705#[proc_macro_attribute]
706pub fn rite(attr: TokenStream, item: TokenStream) -> TokenStream {
707    // Parse optional arguments (currently just inline_always)
708    let args = parse_macro_input!(attr as RiteArgs);
709    let input_fn = parse_macro_input!(item as ItemFn);
710    rite_impl(input_fn, args)
711}
712
713/// Arguments for the `#[rite]` macro.
714///
715/// Currently empty - `#[inline(always)]` is not supported because
716/// `#[inline(always)]` + `#[target_feature]` requires nightly Rust.
717/// The regular `#[inline]` hint is sufficient when called from
718/// matching `#[target_feature]` contexts.
719#[derive(Default)]
720struct RiteArgs {
721    // No options currently - inline_always doesn't work on stable
722}
723
724impl Parse for RiteArgs {
725    fn parse(input: ParseStream) -> syn::Result<Self> {
726        if !input.is_empty() {
727            let ident: Ident = input.parse()?;
728            return Err(syn::Error::new(
729                ident.span(),
730                "#[rite] takes no arguments. Note: inline_always is not supported \
731                 because #[inline(always)] + #[target_feature] requires nightly Rust.",
732            ));
733        }
734        Ok(RiteArgs::default())
735    }
736}
737
738/// Implementation for the `#[rite]` macro.
739fn rite_impl(mut input_fn: ItemFn, args: RiteArgs) -> TokenStream {
740    // Find the token parameter and its features
741    let TokenParamInfo {
742        features,
743        target_arch,
744        ..
745    } = match find_token_param(&input_fn.sig) {
746        Some(result) => result,
747        None => {
748            let msg = "rite requires a token parameter. Supported forms:\n\
749                 - Concrete: `token: X64V3Token`\n\
750                 - impl Trait: `token: impl HasX64V2`\n\
751                 - Generic: `fn foo<T: HasX64V2>(token: T, ...)`";
752            return syn::Error::new_spanned(&input_fn.sig, msg)
753                .to_compile_error()
754                .into();
755        }
756    };
757
758    // Build target_feature attributes
759    let target_feature_attrs: Vec<Attribute> = features
760        .iter()
761        .map(|feature| parse_quote!(#[target_feature(enable = #feature)]))
762        .collect();
763
764    // Always use #[inline] - #[inline(always)] + #[target_feature] requires nightly
765    let _ = args; // RiteArgs is currently empty but kept for future extensibility
766    let inline_attr: Attribute = parse_quote!(#[inline]);
767
768    // Prepend attributes to the function
769    let mut new_attrs = target_feature_attrs;
770    new_attrs.push(inline_attr);
771    new_attrs.append(&mut input_fn.attrs);
772    input_fn.attrs = new_attrs;
773
774    // If we know the target arch, generate cfg-gated impl + stub
775    if let Some(arch) = target_arch {
776        let vis = &input_fn.vis;
777        let sig = &input_fn.sig;
778        let attrs = &input_fn.attrs;
779        let block = &input_fn.block;
780
781        quote! {
782            #[cfg(target_arch = #arch)]
783            #(#attrs)*
784            #vis #sig
785            #block
786
787            #[cfg(not(target_arch = #arch))]
788            #vis #sig {
789                unreachable!(concat!(
790                    "This function requires ",
791                    #arch,
792                    " architecture"
793                ))
794            }
795        }
796        .into()
797    } else {
798        // No specific arch (trait bounds) - just emit the annotated function
799        quote!(#input_fn).into()
800    }
801}
802
803// =============================================================================
804// magetypes! macro - generate platform variants from generic function
805// =============================================================================
806
807/// Generate platform-specific variants from a function by replacing `Token`.
808///
809/// Use `Token` as a placeholder for the token type. The macro generates
810/// suffixed variants with `Token` replaced by the concrete token type, and
811/// each variant wrapped in the appropriate `#[cfg(target_arch = ...)]` guard.
812///
813/// # Default tiers
814///
815/// Without arguments, generates `_v3`, `_v4`, `_neon`, `_wasm128`, `_scalar`:
816///
817/// ```rust,ignore
818/// #[magetypes]
819/// fn process(token: Token, data: &[f32]) -> f32 {
820///     inner_simd_work(token, data)
821/// }
822/// ```
823///
824/// # Explicit tiers
825///
826/// Specify which tiers to generate:
827///
828/// ```rust,ignore
829/// #[magetypes(v1, v3, neon)]
830/// fn process(token: Token, data: &[f32]) -> f32 {
831///     inner_simd_work(token, data)
832/// }
833/// // Generates: process_v1, process_v3, process_neon, process_scalar
834/// ```
835///
836/// `scalar` is always included implicitly.
837///
838/// Known tiers: `v1`, `v2`, `v3`, `v4`, `modern`, `neon`, `neon_aes`,
839/// `neon_sha3`, `neon_crc`, `wasm128`, `scalar`.
840///
841/// # What gets replaced
842///
843/// **Only `Token`** is replaced — with the concrete token type for each variant
844/// (e.g., `archmage::X64V3Token`, `archmage::ScalarToken`). SIMD types like
845/// `f32x8` and constants like `LANES` are **not** replaced by this macro.
846///
847/// # Usage with incant!
848///
849/// The generated variants work with `incant!` for dispatch:
850///
851/// ```rust,ignore
852/// pub fn process_api(data: &[f32]) -> f32 {
853///     incant!(process(data))
854/// }
855///
856/// // Or with matching explicit tiers:
857/// pub fn process_api(data: &[f32]) -> f32 {
858///     incant!(process(data), [v1, v3, neon])
859/// }
860/// ```
861#[proc_macro_attribute]
862pub fn magetypes(attr: TokenStream, item: TokenStream) -> TokenStream {
863    let input_fn = parse_macro_input!(item as ItemFn);
864
865    // Parse optional tier list from attribute args
866    let tier_names: Vec<String> = if attr.is_empty() {
867        DEFAULT_TIER_NAMES.iter().map(|s| s.to_string()).collect()
868    } else {
869        let parser = |input: ParseStream| input.parse_terminated(Ident::parse, Token![,]);
870        let idents = match syn::parse::Parser::parse(parser, attr) {
871            Ok(p) => p,
872            Err(e) => return e.to_compile_error().into(),
873        };
874        idents.iter().map(|i| i.to_string()).collect()
875    };
876
877    let tiers = match resolve_tiers(&tier_names, input_fn.sig.ident.span()) {
878        Ok(t) => t,
879        Err(e) => return e.to_compile_error().into(),
880    };
881
882    magetypes_impl(input_fn, &tiers)
883}
884
885fn magetypes_impl(input_fn: ItemFn, tiers: &[&TierDescriptor]) -> TokenStream {
886    let fn_name = &input_fn.sig.ident;
887    let fn_attrs = &input_fn.attrs;
888
889    // Convert function to string for text substitution
890    let fn_str = input_fn.to_token_stream().to_string();
891
892    let mut variants = Vec::new();
893
894    for tier in tiers {
895        // Create suffixed function name
896        let suffixed_name = format!("{}_{}", fn_name, tier.suffix);
897
898        // Do text substitution
899        let mut variant_str = fn_str.clone();
900
901        // Replace function name
902        variant_str = variant_str.replacen(&fn_name.to_string(), &suffixed_name, 1);
903
904        // Replace Token type with concrete token
905        variant_str = variant_str.replace("Token", tier.token_path);
906
907        // Parse back to tokens
908        let variant_tokens: proc_macro2::TokenStream = match variant_str.parse() {
909            Ok(t) => t,
910            Err(e) => {
911                return syn::Error::new_spanned(
912                    &input_fn,
913                    format!(
914                        "Failed to parse generated variant `{}`: {}",
915                        suffixed_name, e
916                    ),
917                )
918                .to_compile_error()
919                .into();
920            }
921        };
922
923        // Add cfg guards
924        let cfg_guard = match (tier.target_arch, tier.cargo_feature) {
925            (Some(arch), Some(feature)) => {
926                quote! { #[cfg(all(target_arch = #arch, feature = #feature))] }
927            }
928            (Some(arch), None) => {
929                quote! { #[cfg(target_arch = #arch)] }
930            }
931            (None, Some(feature)) => {
932                quote! { #[cfg(feature = #feature)] }
933            }
934            (None, None) => {
935                quote! {} // No guard needed (scalar)
936            }
937        };
938
939        variants.push(quote! {
940            #cfg_guard
941            #variant_tokens
942        });
943    }
944
945    // Remove attributes from the list that should not be duplicated
946    let filtered_attrs: Vec<_> = fn_attrs
947        .iter()
948        .filter(|a| !a.path().is_ident("magetypes"))
949        .collect();
950
951    let output = quote! {
952        #(#filtered_attrs)*
953        #(#variants)*
954    };
955
956    output.into()
957}
958
959// =============================================================================
960// incant! macro - dispatch to platform-specific variants
961// =============================================================================
962
963// =============================================================================
964// Tier descriptors for incant! and #[magetypes]
965// =============================================================================
966
967/// Describes a dispatch tier for incant! and #[magetypes].
968struct TierDescriptor {
969    /// Tier name as written in user code (e.g., "v3", "neon")
970    name: &'static str,
971    /// Function suffix (e.g., "v3", "neon", "scalar")
972    suffix: &'static str,
973    /// Token type path (e.g., "archmage::X64V3Token")
974    token_path: &'static str,
975    /// IntoConcreteToken method name (e.g., "as_x64v3")
976    as_method: &'static str,
977    /// Target architecture for cfg guard (None = no guard)
978    target_arch: Option<&'static str>,
979    /// Required cargo feature (None = no feature guard)
980    cargo_feature: Option<&'static str>,
981    /// Dispatch priority (higher = tried first within same arch)
982    priority: u32,
983}
984
985/// All known tiers in dispatch-priority order (highest first within arch).
986const ALL_TIERS: &[TierDescriptor] = &[
987    // x86: highest to lowest
988    TierDescriptor {
989        name: "modern",
990        suffix: "modern",
991        token_path: "archmage::Avx512ModernToken",
992        as_method: "as_avx512_modern",
993        target_arch: Some("x86_64"),
994        cargo_feature: Some("avx512"),
995        priority: 50,
996    },
997    TierDescriptor {
998        name: "v4",
999        suffix: "v4",
1000        token_path: "archmage::X64V4Token",
1001        as_method: "as_x64v4",
1002        target_arch: Some("x86_64"),
1003        cargo_feature: Some("avx512"),
1004        priority: 40,
1005    },
1006    TierDescriptor {
1007        name: "v3",
1008        suffix: "v3",
1009        token_path: "archmage::X64V3Token",
1010        as_method: "as_x64v3",
1011        target_arch: Some("x86_64"),
1012        cargo_feature: None,
1013        priority: 30,
1014    },
1015    TierDescriptor {
1016        name: "v2",
1017        suffix: "v2",
1018        token_path: "archmage::X64V2Token",
1019        as_method: "as_x64v2",
1020        target_arch: Some("x86_64"),
1021        cargo_feature: None,
1022        priority: 20,
1023    },
1024    TierDescriptor {
1025        name: "v1",
1026        suffix: "v1",
1027        token_path: "archmage::X64V1Token",
1028        as_method: "as_x64v1",
1029        target_arch: Some("x86_64"),
1030        cargo_feature: None,
1031        priority: 10,
1032    },
1033    // ARM
1034    TierDescriptor {
1035        name: "neon_aes",
1036        suffix: "neon_aes",
1037        token_path: "archmage::NeonAesToken",
1038        as_method: "as_neon_aes",
1039        target_arch: Some("aarch64"),
1040        cargo_feature: None,
1041        priority: 30,
1042    },
1043    TierDescriptor {
1044        name: "neon_sha3",
1045        suffix: "neon_sha3",
1046        token_path: "archmage::NeonSha3Token",
1047        as_method: "as_neon_sha3",
1048        target_arch: Some("aarch64"),
1049        cargo_feature: None,
1050        priority: 30,
1051    },
1052    TierDescriptor {
1053        name: "neon_crc",
1054        suffix: "neon_crc",
1055        token_path: "archmage::NeonCrcToken",
1056        as_method: "as_neon_crc",
1057        target_arch: Some("aarch64"),
1058        cargo_feature: None,
1059        priority: 30,
1060    },
1061    TierDescriptor {
1062        name: "neon",
1063        suffix: "neon",
1064        token_path: "archmage::NeonToken",
1065        as_method: "as_neon",
1066        target_arch: Some("aarch64"),
1067        cargo_feature: None,
1068        priority: 20,
1069    },
1070    // WASM
1071    TierDescriptor {
1072        name: "wasm128",
1073        suffix: "wasm128",
1074        token_path: "archmage::Wasm128Token",
1075        as_method: "as_wasm128",
1076        target_arch: Some("wasm32"),
1077        cargo_feature: None,
1078        priority: 20,
1079    },
1080    // Scalar (always last)
1081    TierDescriptor {
1082        name: "scalar",
1083        suffix: "scalar",
1084        token_path: "archmage::ScalarToken",
1085        as_method: "as_scalar",
1086        target_arch: None,
1087        cargo_feature: None,
1088        priority: 0,
1089    },
1090];
1091
1092/// Default tiers (backwards-compatible with pre-explicit behavior).
1093const DEFAULT_TIER_NAMES: &[&str] = &["v4", "v3", "neon", "wasm128", "scalar"];
1094
1095/// Look up a tier by name, returning an error on unknown names.
1096fn find_tier(name: &str) -> Option<&'static TierDescriptor> {
1097    ALL_TIERS.iter().find(|t| t.name == name)
1098}
1099
1100/// Resolve tier names to descriptors, sorted by dispatch priority (highest first).
1101/// Always appends "scalar" if not already present.
1102fn resolve_tiers(
1103    tier_names: &[String],
1104    error_span: proc_macro2::Span,
1105) -> syn::Result<Vec<&'static TierDescriptor>> {
1106    let mut tiers = Vec::new();
1107    for name in tier_names {
1108        match find_tier(name) {
1109            Some(tier) => tiers.push(tier),
1110            None => {
1111                let known: Vec<&str> = ALL_TIERS.iter().map(|t| t.name).collect();
1112                return Err(syn::Error::new(
1113                    error_span,
1114                    format!("unknown tier `{}`. Known tiers: {}", name, known.join(", ")),
1115                ));
1116            }
1117        }
1118    }
1119
1120    // Always include scalar fallback
1121    if !tiers.iter().any(|t| t.name == "scalar") {
1122        tiers.push(find_tier("scalar").unwrap());
1123    }
1124
1125    // Sort by priority (highest first) for correct dispatch order
1126    tiers.sort_by(|a, b| b.priority.cmp(&a.priority));
1127
1128    Ok(tiers)
1129}
1130
1131// =============================================================================
1132// incant! macro - dispatch to platform-specific variants
1133// =============================================================================
1134
1135/// Input for the incant! macro
1136struct IncantInput {
1137    /// Function name to call
1138    func_name: Ident,
1139    /// Arguments to pass
1140    args: Vec<syn::Expr>,
1141    /// Optional token variable for passthrough mode
1142    with_token: Option<syn::Expr>,
1143    /// Optional explicit tier list (None = default tiers)
1144    tiers: Option<(Vec<String>, proc_macro2::Span)>,
1145}
1146
1147impl Parse for IncantInput {
1148    fn parse(input: ParseStream) -> syn::Result<Self> {
1149        // Parse: function_name(arg1, arg2, ...) [with token_expr] [, [tier1, tier2, ...]]
1150        let func_name: Ident = input.parse()?;
1151
1152        // Parse parenthesized arguments
1153        let content;
1154        syn::parenthesized!(content in input);
1155        let args = content
1156            .parse_terminated(syn::Expr::parse, Token![,])?
1157            .into_iter()
1158            .collect();
1159
1160        // Check for optional "with token"
1161        let with_token = if input.peek(Ident) {
1162            let kw: Ident = input.parse()?;
1163            if kw != "with" {
1164                return Err(syn::Error::new_spanned(kw, "expected `with` keyword"));
1165            }
1166            Some(input.parse()?)
1167        } else {
1168            None
1169        };
1170
1171        // Check for optional tier list: , [tier1, tier2, ...]
1172        let tiers = if input.peek(Token![,]) {
1173            let _: Token![,] = input.parse()?;
1174            let bracket_content;
1175            let bracket = syn::bracketed!(bracket_content in input);
1176            let tier_idents = bracket_content.parse_terminated(Ident::parse, Token![,])?;
1177            let tier_names: Vec<String> = tier_idents.iter().map(|i| i.to_string()).collect();
1178            Some((tier_names, bracket.span.join()))
1179        } else {
1180            None
1181        };
1182
1183        Ok(IncantInput {
1184            func_name,
1185            args,
1186            with_token,
1187            tiers,
1188        })
1189    }
1190}
1191
1192/// Dispatch to platform-specific SIMD variants.
1193///
1194/// # Entry Point Mode (no token yet)
1195///
1196/// Summons tokens and dispatches to the best available variant:
1197///
1198/// ```rust,ignore
1199/// pub fn public_api(data: &[f32]) -> f32 {
1200///     incant!(dot(data))
1201/// }
1202/// ```
1203///
1204/// Expands to runtime feature detection + dispatch to `dot_v3`, `dot_v4`,
1205/// `dot_neon`, `dot_wasm128`, or `dot_scalar`.
1206///
1207/// # Explicit Tiers
1208///
1209/// Specify which tiers to dispatch to:
1210///
1211/// ```rust,ignore
1212/// // Only dispatch to v1, v3, neon, and scalar
1213/// pub fn api(data: &[f32]) -> f32 {
1214///     incant!(process(data), [v1, v3, neon])
1215/// }
1216/// ```
1217///
1218/// `scalar` is always included implicitly. Unknown tier names cause a
1219/// compile error. Tiers are automatically sorted into correct dispatch
1220/// order (highest priority first).
1221///
1222/// Known tiers: `v1`, `v2`, `v3`, `v4`, `modern`, `neon`, `neon_aes`,
1223/// `neon_sha3`, `neon_crc`, `wasm128`, `scalar`.
1224///
1225/// # Passthrough Mode (already have token)
1226///
1227/// Uses compile-time dispatch via `IntoConcreteToken`:
1228///
1229/// ```rust,ignore
1230/// #[arcane]
1231/// fn outer(token: X64V3Token, data: &[f32]) -> f32 {
1232///     incant!(inner(data) with token)
1233/// }
1234/// ```
1235///
1236/// Also supports explicit tiers:
1237///
1238/// ```rust,ignore
1239/// fn inner<T: IntoConcreteToken>(token: T, data: &[f32]) -> f32 {
1240///     incant!(process(data) with token, [v3, neon])
1241/// }
1242/// ```
1243///
1244/// The compiler monomorphizes the dispatch, eliminating non-matching branches.
1245///
1246/// # Variant Naming
1247///
1248/// Functions must have suffixed variants matching the selected tiers:
1249/// - `_v1` for `X64V1Token`
1250/// - `_v2` for `X64V2Token`
1251/// - `_v3` for `X64V3Token`
1252/// - `_v4` for `X64V4Token` (requires `avx512` feature)
1253/// - `_modern` for `Avx512ModernToken` (requires `avx512` feature)
1254/// - `_neon` for `NeonToken`
1255/// - `_neon_aes` for `NeonAesToken`
1256/// - `_neon_sha3` for `NeonSha3Token`
1257/// - `_neon_crc` for `NeonCrcToken`
1258/// - `_wasm128` for `Wasm128Token`
1259/// - `_scalar` for `ScalarToken`
1260#[proc_macro]
1261pub fn incant(input: TokenStream) -> TokenStream {
1262    let input = parse_macro_input!(input as IncantInput);
1263    incant_impl(input)
1264}
1265
1266/// Legacy alias for [`incant!`].
1267#[proc_macro]
1268pub fn simd_route(input: TokenStream) -> TokenStream {
1269    let input = parse_macro_input!(input as IncantInput);
1270    incant_impl(input)
1271}
1272
1273fn incant_impl(input: IncantInput) -> TokenStream {
1274    let func_name = &input.func_name;
1275    let args = &input.args;
1276
1277    // Resolve tiers
1278    let tier_names: Vec<String> = match &input.tiers {
1279        Some((names, _)) => names.clone(),
1280        None => DEFAULT_TIER_NAMES.iter().map(|s| s.to_string()).collect(),
1281    };
1282    let error_span = input
1283        .tiers
1284        .as_ref()
1285        .map(|(_, span)| *span)
1286        .unwrap_or_else(|| func_name.span());
1287
1288    let tiers = match resolve_tiers(&tier_names, error_span) {
1289        Ok(t) => t,
1290        Err(e) => return e.to_compile_error().into(),
1291    };
1292
1293    // Group tiers by architecture for cfg-guarded blocks
1294    // Within each arch, tiers are already sorted by priority (highest first)
1295    if let Some(token_expr) = &input.with_token {
1296        gen_incant_passthrough(func_name, args, token_expr, &tiers)
1297    } else {
1298        gen_incant_entry(func_name, args, &tiers)
1299    }
1300}
1301
1302/// Generate incant! passthrough mode (already have a token).
1303fn gen_incant_passthrough(
1304    func_name: &Ident,
1305    args: &[syn::Expr],
1306    token_expr: &syn::Expr,
1307    tiers: &[&TierDescriptor],
1308) -> TokenStream {
1309    let mut dispatch_arms = Vec::new();
1310
1311    // Group non-scalar tiers by (target_arch, cargo_feature) for nested cfg blocks
1312    let mut arch_groups: Vec<(Option<&str>, Option<&str>, Vec<&TierDescriptor>)> = Vec::new();
1313    for tier in tiers {
1314        if tier.name == "scalar" {
1315            continue; // Handle scalar separately at the end
1316        }
1317        let key = (tier.target_arch, tier.cargo_feature);
1318        if let Some(group) = arch_groups.iter_mut().find(|(a, f, _)| (*a, *f) == key) {
1319            group.2.push(tier);
1320        } else {
1321            arch_groups.push((tier.target_arch, tier.cargo_feature, vec![tier]));
1322        }
1323    }
1324
1325    for (target_arch, cargo_feature, group_tiers) in &arch_groups {
1326        let mut tier_checks = Vec::new();
1327        for tier in group_tiers {
1328            let fn_suffixed = format_ident!("{}_{}", func_name, tier.suffix);
1329            let as_method = format_ident!("{}", tier.as_method);
1330            tier_checks.push(quote! {
1331                if let Some(__t) = __incant_token.#as_method() {
1332                    break '__incant #fn_suffixed(__t, #(#args),*);
1333                }
1334            });
1335        }
1336
1337        let inner = quote! { #(#tier_checks)* };
1338
1339        let guarded = match (target_arch, cargo_feature) {
1340            (Some(arch), Some(feat)) => quote! {
1341                #[cfg(target_arch = #arch)]
1342                {
1343                    #[cfg(feature = #feat)]
1344                    { #inner }
1345                }
1346            },
1347            (Some(arch), None) => quote! {
1348                #[cfg(target_arch = #arch)]
1349                { #inner }
1350            },
1351            (None, Some(feat)) => quote! {
1352                #[cfg(feature = #feat)]
1353                { #inner }
1354            },
1355            (None, None) => inner,
1356        };
1357
1358        dispatch_arms.push(guarded);
1359    }
1360
1361    // Scalar fallback (always last)
1362    let fn_scalar = format_ident!("{}_scalar", func_name);
1363    let scalar_arm = if tiers.iter().any(|t| t.name == "scalar") {
1364        quote! {
1365            if let Some(__t) = __incant_token.as_scalar() {
1366                break '__incant #fn_scalar(__t, #(#args),*);
1367            }
1368            unreachable!("Token did not match any known variant")
1369        }
1370    } else {
1371        quote! { unreachable!("Token did not match any known variant") }
1372    };
1373
1374    let expanded = quote! {
1375        '__incant: {
1376            use archmage::IntoConcreteToken;
1377            let __incant_token = #token_expr;
1378            #(#dispatch_arms)*
1379            #scalar_arm
1380        }
1381    };
1382    expanded.into()
1383}
1384
1385/// Generate incant! entry point mode (summon tokens).
1386fn gen_incant_entry(
1387    func_name: &Ident,
1388    args: &[syn::Expr],
1389    tiers: &[&TierDescriptor],
1390) -> TokenStream {
1391    let mut dispatch_arms = Vec::new();
1392
1393    // Group non-scalar tiers by target_arch for cfg blocks.
1394    // Within each arch group, further split by cargo_feature.
1395    let mut arch_groups: Vec<(Option<&str>, Vec<&TierDescriptor>)> = Vec::new();
1396    for tier in tiers {
1397        if tier.name == "scalar" {
1398            continue;
1399        }
1400        if let Some(group) = arch_groups.iter_mut().find(|(a, _)| *a == tier.target_arch) {
1401            group.1.push(tier);
1402        } else {
1403            arch_groups.push((tier.target_arch, vec![tier]));
1404        }
1405    }
1406
1407    for (target_arch, group_tiers) in &arch_groups {
1408        let mut tier_checks = Vec::new();
1409        for tier in group_tiers {
1410            let fn_suffixed = format_ident!("{}_{}", func_name, tier.suffix);
1411            let token_path: syn::Path = syn::parse_str(tier.token_path).unwrap();
1412
1413            let check = quote! {
1414                if let Some(__t) = #token_path::summon() {
1415                    break '__incant #fn_suffixed(__t, #(#args),*);
1416                }
1417            };
1418
1419            if let Some(feat) = tier.cargo_feature {
1420                tier_checks.push(quote! {
1421                    #[cfg(feature = #feat)]
1422                    { #check }
1423                });
1424            } else {
1425                tier_checks.push(check);
1426            }
1427        }
1428
1429        let inner = quote! { #(#tier_checks)* };
1430
1431        if let Some(arch) = target_arch {
1432            dispatch_arms.push(quote! {
1433                #[cfg(target_arch = #arch)]
1434                { #inner }
1435            });
1436        } else {
1437            dispatch_arms.push(inner);
1438        }
1439    }
1440
1441    // Scalar fallback
1442    let fn_scalar = format_ident!("{}_scalar", func_name);
1443
1444    let expanded = quote! {
1445        '__incant: {
1446            use archmage::SimdToken;
1447            #(#dispatch_arms)*
1448            #fn_scalar(archmage::ScalarToken, #(#args),*)
1449        }
1450    };
1451    expanded.into()
1452}
1453
1454// =============================================================================
1455// Unit tests for token/trait recognition maps
1456// =============================================================================
1457
1458#[cfg(test)]
1459mod tests {
1460    use super::*;
1461
1462    use super::generated::{ALL_CONCRETE_TOKENS, ALL_TRAIT_NAMES};
1463
1464    #[test]
1465    fn every_concrete_token_is_in_token_to_features() {
1466        for &name in ALL_CONCRETE_TOKENS {
1467            assert!(
1468                token_to_features(name).is_some(),
1469                "Token `{}` exists in runtime crate but is NOT recognized by \
1470                 token_to_features() in the proc macro. Add it!",
1471                name
1472            );
1473        }
1474    }
1475
1476    #[test]
1477    fn every_trait_is_in_trait_to_features() {
1478        for &name in ALL_TRAIT_NAMES {
1479            assert!(
1480                trait_to_features(name).is_some(),
1481                "Trait `{}` exists in runtime crate but is NOT recognized by \
1482                 trait_to_features() in the proc macro. Add it!",
1483                name
1484            );
1485        }
1486    }
1487
1488    #[test]
1489    fn token_aliases_map_to_same_features() {
1490        // Desktop64 = X64V3Token
1491        assert_eq!(
1492            token_to_features("Desktop64"),
1493            token_to_features("X64V3Token"),
1494            "Desktop64 and X64V3Token should map to identical features"
1495        );
1496
1497        // Server64 = X64V4Token = Avx512Token
1498        assert_eq!(
1499            token_to_features("Server64"),
1500            token_to_features("X64V4Token"),
1501            "Server64 and X64V4Token should map to identical features"
1502        );
1503        assert_eq!(
1504            token_to_features("X64V4Token"),
1505            token_to_features("Avx512Token"),
1506            "X64V4Token and Avx512Token should map to identical features"
1507        );
1508
1509        // Arm64 = NeonToken
1510        assert_eq!(
1511            token_to_features("Arm64"),
1512            token_to_features("NeonToken"),
1513            "Arm64 and NeonToken should map to identical features"
1514        );
1515    }
1516
1517    #[test]
1518    fn trait_to_features_includes_tokens_as_bounds() {
1519        // Tier tokens should also work as trait bounds
1520        // (for `impl X64V3Token` patterns, even though Rust won't allow it,
1521        // the macro processes AST before type checking)
1522        let tier_tokens = [
1523            "X64V2Token",
1524            "X64V3Token",
1525            "Desktop64",
1526            "Avx2FmaToken",
1527            "X64V4Token",
1528            "Avx512Token",
1529            "Server64",
1530            "Avx512ModernToken",
1531            "Avx512Fp16Token",
1532            "NeonToken",
1533            "Arm64",
1534            "NeonAesToken",
1535            "NeonSha3Token",
1536            "NeonCrcToken",
1537        ];
1538
1539        for &name in &tier_tokens {
1540            assert!(
1541                trait_to_features(name).is_some(),
1542                "Tier token `{}` should also be recognized in trait_to_features() \
1543                 for use as a generic bound. Add it!",
1544                name
1545            );
1546        }
1547    }
1548
1549    #[test]
1550    fn trait_features_are_cumulative() {
1551        // HasX64V4 should include all HasX64V2 features plus more
1552        let v2_features = trait_to_features("HasX64V2").unwrap();
1553        let v4_features = trait_to_features("HasX64V4").unwrap();
1554
1555        for &f in v2_features {
1556            assert!(
1557                v4_features.contains(&f),
1558                "HasX64V4 should include v2 feature `{}` but doesn't",
1559                f
1560            );
1561        }
1562
1563        // v4 should have more features than v2
1564        assert!(
1565            v4_features.len() > v2_features.len(),
1566            "HasX64V4 should have more features than HasX64V2"
1567        );
1568    }
1569
1570    #[test]
1571    fn x64v3_trait_features_include_v2() {
1572        // X64V3Token as trait bound should include v2 features
1573        let v2 = trait_to_features("HasX64V2").unwrap();
1574        let v3 = trait_to_features("X64V3Token").unwrap();
1575
1576        for &f in v2 {
1577            assert!(
1578                v3.contains(&f),
1579                "X64V3Token trait features should include v2 feature `{}` but don't",
1580                f
1581            );
1582        }
1583    }
1584
1585    #[test]
1586    fn has_neon_aes_includes_neon() {
1587        let neon = trait_to_features("HasNeon").unwrap();
1588        let neon_aes = trait_to_features("HasNeonAes").unwrap();
1589
1590        for &f in neon {
1591            assert!(
1592                neon_aes.contains(&f),
1593                "HasNeonAes should include NEON feature `{}`",
1594                f
1595            );
1596        }
1597    }
1598
1599    #[test]
1600    fn no_removed_traits_are_recognized() {
1601        // These traits were removed in 0.3.0 and should NOT be recognized
1602        let removed = [
1603            "HasSse",
1604            "HasSse2",
1605            "HasSse41",
1606            "HasSse42",
1607            "HasAvx",
1608            "HasAvx2",
1609            "HasFma",
1610            "HasAvx512f",
1611            "HasAvx512bw",
1612            "HasAvx512vl",
1613            "HasAvx512vbmi2",
1614            "HasSve",
1615            "HasSve2",
1616        ];
1617
1618        for &name in &removed {
1619            assert!(
1620                trait_to_features(name).is_none(),
1621                "Removed trait `{}` should NOT be in trait_to_features(). \
1622                 It was removed in 0.3.0 — users should migrate to tier traits.",
1623                name
1624            );
1625        }
1626    }
1627
1628    #[test]
1629    fn no_nonexistent_tokens_are_recognized() {
1630        // These tokens don't exist and should NOT be recognized
1631        let fake = [
1632            "SveToken",
1633            "Sve2Token",
1634            "Avx512VnniToken",
1635            "X64V4ModernToken",
1636            "NeonFp16Token",
1637        ];
1638
1639        for &name in &fake {
1640            assert!(
1641                token_to_features(name).is_none(),
1642                "Non-existent token `{}` should NOT be in token_to_features()",
1643                name
1644            );
1645        }
1646    }
1647}