archmage-macros 0.9.29

//! Shared utilities for all proc-macros.

use proc_macro2::Ident;
use quote::{ToTokens, format_ident, quote};
use syn::{Attribute, GenericParam, Signature, Type, parse::ParseStream, token};

/// A function parsed with the body left as an opaque TokenStream.
///
/// Only the signature is fully parsed into an AST — the body tokens are collected
/// without building any AST nodes (no expressions, statements, or patterns parsed).
/// Parsing cost therefore scales with the signature rather than a body AST.
#[derive(Clone)]
pub(crate) struct LightFn {
    pub attrs: Vec<Attribute>,
    pub vis: syn::Visibility,
    pub sig: Signature,
    pub brace_token: token::Brace,
    pub body: proc_macro2::TokenStream,
}

impl syn::parse::Parse for LightFn {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let attrs = input.call(Attribute::parse_outer)?;
        let vis: syn::Visibility = input.parse()?;
        let sig: Signature = input.parse()?;
        let content;
        let brace_token = syn::braced!(content in input);
        let body: proc_macro2::TokenStream = content.parse()?;
        Ok(LightFn {
            attrs,
            vis,
            sig,
            brace_token,
            body,
        })
    }
}

impl ToTokens for LightFn {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        for attr in &self.attrs {
            attr.to_tokens(tokens);
        }
        self.vis.to_tokens(tokens);
        self.sig.to_tokens(tokens);
        self.brace_token.surround(tokens, |tokens| {
            self.body.to_tokens(tokens);
        });
    }
}

/// Filter out `#[inline]`, `#[inline(always)]`, `#[inline(never)]` from attributes.
pub(crate) fn filter_inline_attrs(attrs: &[Attribute]) -> Vec<&Attribute> {
    attrs
        .iter()
        .filter(|attr| !attr.path().is_ident("inline"))
        .collect()
}

/// Check if an attribute is a lint-control attribute.
pub(crate) fn is_lint_attr(attr: &Attribute) -> bool {
    let path = attr.path();
    path.is_ident("allow")
        || path.is_ident("expect")
        || path.is_ident("deny")
        || path.is_ident("warn")
        || path.is_ident("forbid")
}

/// Extract lint-control attributes from a list of attributes.
pub(crate) fn filter_lint_attrs(attrs: &[Attribute]) -> Vec<&Attribute> {
    attrs.iter().filter(|attr| is_lint_attr(attr)).collect()
}

/// Generate a cfg guard combining target_arch and an optional feature gate.
pub(crate) fn gen_cfg_guard(
    target_arch: Option<&str>,
    cfg_feature: Option<&str>,
) -> proc_macro2::TokenStream {
    match (target_arch, cfg_feature) {
        (Some(arch), Some(feat)) => {
            quote! { #[cfg(all(target_arch = #arch, feature = #feat))] }
        }
        (Some(arch), None) => quote! { #[cfg(target_arch = #arch)] },
        (None, Some(feat)) => quote! { #[cfg(feature = #feat)] },
        (None, None) => quote! {},
    }
}

/// Build a turbofish token stream from a function's generics.
pub(crate) fn build_turbofish(generics: &syn::Generics) -> proc_macro2::TokenStream {
    let params: Vec<proc_macro2::TokenStream> = generics
        .params
        .iter()
        .filter_map(|param| match param {
            GenericParam::Type(tp) => {
                let ident = &tp.ident;
                Some(quote! { #ident })
            }
            GenericParam::Const(cp) => {
                let ident = &cp.ident;
                Some(quote! { #ident })
            }
            GenericParam::Lifetime(_) => None,
        })
        .collect();
    if params.is_empty() {
        quote! {}
    } else {
        quote! { ::<#(#params),*> }
    }
}

/// Conservative token-level presence check. Literals and partial identifier
/// matches do not count. A false positive only costs a rewrite pass; a false
/// negative would change expansion, so all delimiter kinds are traversed.
pub(crate) fn tokens_contain_ident(tokens: &proc_macro2::TokenStream, names: &[&str]) -> bool {
    tokens.clone().into_iter().any(|token| match token {
        proc_macro2::TokenTree::Ident(id) => names.iter().any(|name| id == *name),
        proc_macro2::TokenTree::Group(group) => tokens_contain_ident(&group.stream(), names),
        _ => false,
    })
}

/// Replace all occurrences of a named identifier in a token stream.
///
/// Recurses into groups (braces, parens, brackets). Each matching `Ident` is
/// replaced with `replacement` — which can be multiple tokens (e.g., a path
/// like `archmage::X64V3Token`). Non-matching tokens pass through unchanged.
pub(crate) fn replace_ident_in_tokens(
    tokens: proc_macro2::TokenStream,
    target: &str,
    replacement: &proc_macro2::TokenStream,
) -> proc_macro2::TokenStream {
    let mut result = proc_macro2::TokenStream::new();
    for tt in tokens {
        match tt {
            proc_macro2::TokenTree::Ident(ref ident) if *ident == target => {
                result.extend(replacement.clone());
            }
            proc_macro2::TokenTree::Group(group) => {
                let stream = group.stream();
                if !tokens_contain_ident(&stream, &[target]) {
                    result.extend(std::iter::once(proc_macro2::TokenTree::Group(group)));
                    continue;
                }
                let new_stream = replace_ident_in_tokens(stream, target, replacement);
                let mut new_group = proc_macro2::Group::new(group.delimiter(), new_stream);
                new_group.set_span(group.span());
                result.extend(std::iter::once(proc_macro2::TokenTree::Group(new_group)));
            }
            other => {
                result.extend(std::iter::once(other));
            }
        }
    }
    result
}

/// Replace all `Self` identifier tokens with a concrete type in a token stream.
pub(crate) fn replace_self_in_tokens(
    tokens: proc_macro2::TokenStream,
    replacement: &Type,
) -> proc_macro2::TokenStream {
    replace_ident_in_tokens(tokens, "Self", &replacement.to_token_stream())
}

/// Generate import statements for intrinsics and/or magetypes.
pub(crate) fn generate_imports(
    target_arch: Option<&str>,
    magetypes_namespace: Option<&str>,
    import_intrinsics: bool,
    import_magetypes: bool,
) -> proc_macro2::TokenStream {
    let mut imports = proc_macro2::TokenStream::new();

    if import_intrinsics && let Some(arch) = target_arch {
        let arch_ident = quote::format_ident!("{}", arch);
        imports.extend(quote! {
            #[allow(unused_imports)]
            use archmage::intrinsics::#arch_ident::*;
        });
    }

    if import_magetypes && let Some(ns) = magetypes_namespace {
        let ns_ident = quote::format_ident!("{}", ns);
        imports.extend(quote! {
            #[allow(unused_imports)]
            use magetypes::simd::#ns_ident::*;
            #[allow(unused_imports)]
            use magetypes::simd::backends::*;
        });
    }

    imports
}

/// Check if any argument expression contains the `Token` identifier.
/// Check if an expression is a bare ident matching a given name.
pub(crate) fn is_bare_ident_pub(expr: &syn::Expr, name: &str) -> bool {
    is_bare_ident(expr, name)
}

fn is_bare_ident(expr: &syn::Expr, name: &str) -> bool {
    match expr {
        syn::Expr::Path(p) => {
            p.qself.is_none() && p.path.segments.len() == 1 && p.path.segments[0].ident == name
        }
        _ => false,
    }
}

/// Detect how the token is placed in incant! arguments.
///
/// Returns the kind of token placement found:
/// - `TokenPlacement::Explicit(index)`: arg at `index` is the `Token` marker
/// - `TokenPlacement::Variable(index)`: arg at `index` is the caller's token variable
/// - `TokenPlacement::None`: no token in args (will be prepended, deprecated)
pub(crate) enum TokenPlacement {
    /// `Token` marker found in args
    Explicit,
    /// Caller's token variable name found at this arg index
    Variable(usize),
    /// No token in args — will prepend (deprecated)
    None,
}

/// Find where the token is placed in incant! arguments.
///
/// Checks for:
/// 1. `Token` marker (explicit placeholder for summon mode)
/// 2. A bare ident matching `caller_token_ident` (the actual variable)
///
/// Falls back to `None` (token will be prepended).
pub(crate) fn find_token_placement(
    args: &[syn::Expr],
    caller_token_ident: Option<&str>,
) -> TokenPlacement {
    for arg in args.iter() {
        if is_bare_ident(arg, "Token") {
            return TokenPlacement::Explicit;
        }
    }
    if let Some(ident_name) = caller_token_ident {
        for (i, arg) in args.iter().enumerate() {
            if is_bare_ident(arg, ident_name) {
                return TokenPlacement::Variable(i);
            }
        }
    }
    TokenPlacement::None
}

/// Build call arguments with the token in the correct position.
///
/// Handles three cases:
/// - `Token` marker in args: replace with `token_expr`
/// - Caller's token variable in args: replace with `token_expr` (may downcast)
/// - Neither: prepend `token_expr` (backward compat, deprecated)
pub(crate) fn build_call_args(
    args: &[syn::Expr],
    token_expr: &proc_macro2::TokenStream,
) -> proc_macro2::TokenStream {
    build_call_args_with_ident(args, token_expr, None)
}

/// Build call arguments, also checking for the caller's token variable name.
pub(crate) fn build_call_args_with_ident(
    args: &[syn::Expr],
    token_expr: &proc_macro2::TokenStream,
    caller_token_ident: Option<&str>,
) -> proc_macro2::TokenStream {
    match find_token_placement(args, caller_token_ident) {
        TokenPlacement::Explicit => {
            // Replace Token marker with token expression
            let replaced = args
                .iter()
                .map(|arg| replace_ident_in_tokens(arg.to_token_stream(), "Token", token_expr));
            quote! { #(#replaced),* }
        }
        TokenPlacement::Variable(idx) => {
            // Replace the caller's token variable with the target token expression
            let replaced = args.iter().enumerate().map(|(i, arg)| {
                if i == idx {
                    token_expr.clone()
                } else {
                    arg.to_token_stream()
                }
            });
            quote! { #(#replaced),* }
        }
        TokenPlacement::None => {
            // No token in args — prepend (backward compat)
            quote! { #token_expr, #(#args),* }
        }
    }
}

/// Build call arguments for scalar fallback.
///
/// If `Token` marker or caller's token variable is in args, replace with `ScalarToken`.
/// If neither, prepend `ScalarToken`.
pub(crate) fn build_scalar_call_args(args: &[syn::Expr]) -> proc_macro2::TokenStream {
    let scalar = quote! { archmage::ScalarToken };
    build_call_args(args, &scalar)
}

/// Suffix the last segment of a path: `process` → `process_v3`.
pub(crate) fn suffix_path(path: &syn::Path, suffix: &str) -> syn::Path {
    let mut suffixed = path.clone();
    if let Some(last) = suffixed.segments.last_mut() {
        last.ident = quote::format_ident!("{}_{}", last.ident, suffix);
    }
    suffixed
}

/// Emit the tier-trait identity assertion for a trait or generic token bound.
///
/// `#[arcane]`/`#[rite]` resolve a tier trait **by name**: seeing `HasX64V2` in
/// the signature is what selects the SSE4.2 `#[target_feature]` list. On its own
/// that proves nothing, because a downstream crate can declare
///
/// ```ignore
/// pub trait HasX64V2 {}
/// impl HasX64V2 for NotAToken {}
/// ```
///
/// and receive the whole feature set with no token, no sealing and no runtime
/// detection — in ordinary safe code, with no diagnostic. On a CPU without the
/// features that is SIGILL.
///
/// A concrete token is already covered: the wrapper asserts its generated
/// `__ARCHMAGE_ASSERT_TIER_<tag>` const, which a same-named local struct does
/// not have. This is the equivalent for the bound case. It re-states each tier
/// trait through an **absolute** `::archmage::` path, which a local trait cannot
/// shadow, and requires the token value's type to satisfy it:
///
/// ```ignore
/// const fn __archmage_assert_tier_trait<__T: ?Sized + ::archmage::HasX64V2>(_: &__T) {}
/// __archmage_assert_tier_trait(&token);
/// ```
///
/// Because archmage's tier traits are sealed through `SimdToken`, satisfying the
/// bound is only possible for a genuine token of that tier or stronger. A
/// same-named local trait fails with rustc's own diagnostic, which names both:
/// "`impl HasX64V2` implements similarly named trait `HasX64V2`, but not
/// `archmage::HasX64V2`". A real but *weaker* token fails too, closing the
/// trait-bound analogue of the token-aliasing hole.
///
/// It is a `const fn` so it provably has no runtime body, and its parameter is
/// taken by reference so the form works for `impl Trait` in argument position,
/// where the type cannot be named. All bounds are emitted on one helper, so a
/// multi-trait bound (`impl HasNeon + HasNeonAes`) is checked as the same union
/// of tiers the `#[target_feature]` list was built from.
///
/// Returns nothing for a concrete token (`tier_traits` empty) — that path keeps
/// its cheaper const assertion.
pub(crate) fn gen_tier_trait_assertion(
    tier_traits: &[String],
    token_ident: &Ident,
) -> proc_macro2::TokenStream {
    if tier_traits.is_empty() {
        return quote! {};
    }
    let bounds = tier_traits.iter().map(|name| {
        let ident = format_ident!("{}", name);
        quote! { + ::archmage::#ident }
    });
    quote! {
        {
            #[inline(always)]
            const fn __archmage_assert_tier_trait<__T: ?Sized #(#bounds)*>(_: &__T) {}
            __archmage_assert_tier_trait(&#token_ident);
        }
    }
}