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}