rs/macros/src/step.rs

//! `#[step]` — generates `StepReset` impl for structs and reset
//! functions for statics.
//!
//! Reset rules by field/value type:
//! - Integer types (u8..u128, i8..i128) → 0
//! - bool → false
//! - Option<T> → None
//! - BoundedVec<T, N> / BoundedMap<K, V, N> → clear()
//! - AtomicU32 / AtomicU64 → store(0, Ordering::SeqCst)
//! - Custom types → StepReset::reset(&mut self)

use proc_macro2::{Span, TokenStream};
use quote::{format_ident, quote};
use syn::{
    parse2, Data, DeriveInput, Error, Fields, Result, Type,
};

pub fn expand(_attr: TokenStream, item: TokenStream) -> Result<TokenStream> {
    // Try struct first.
    if let Ok(input) = parse2::<DeriveInput>(item.clone()) {
        return expand_struct(input);
    }

    // Try static.
    if let Ok(input) = parse2::<syn::ItemStatic>(item) {
        return expand_static(input);
    }

    Err(Error::new(
        Span::call_site(),
        "#[step] can only be applied to structs with named fields or static items",
    ))
}

fn expand_struct(input: DeriveInput) -> Result<TokenStream> {
    let fields = match &input.data {
        Data::Struct(data) => match &data.fields {
            Fields::Named(named) => &named.named,
            Fields::Unnamed(_) => {
                return Err(Error::new_spanned(
                    &input.ident,
                    "#[step] does not support tuple structs",
                ));
            }
            Fields::Unit => {
                return Err(Error::new_spanned(
                    &input.ident,
                    "#[step] does not support unit structs",
                ));
            }
        },
        _ => {
            return Err(Error::new_spanned(
                &input.ident,
                "#[step] can only be applied to structs with named fields",
            ));
        }
    };

    let name = &input.ident;
    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();

    let reset_stmts: Vec<TokenStream> = fields
        .iter()
        .map(|f| {
            let field_name = f.ident.as_ref().unwrap();
            field_reset_expr(&f.ty, field_name)
        })
        .collect();

    let attrs = &input.attrs;
    let vis = &input.vis;
    let generics = &input.generics;
    let fields_def = match &input.data {
        Data::Struct(data) => &data.fields,
        _ => unreachable!(),
    };

    Ok(quote! {
        #(#attrs)*
        #vis struct #name #generics #where_clause #fields_def

        impl #impl_generics rs_lang::StepReset for #name #ty_generics #where_clause {
            fn reset(&mut self) {
                #(#reset_stmts)*
            }
        }
    })
}

fn expand_static(item: syn::ItemStatic) -> Result<TokenStream> {
    let name = &item.ident;
    let ty = &item.ty;
    let vis = &item.vis;
    let expr = &item.expr;
    let mutability = &item.mutability;
    let attrs = &item.attrs;

    let reset_fn_name = format_ident!(
        "__rs_step_reset_{}",
        name.to_string().to_lowercase()
    );

    let reset_body = static_reset_expr(ty, name, expr);

    Ok(quote! {
        #(#attrs)*
        #vis static #mutability #name: #ty = #expr;

        /// Generated by `#[step]` — resets the static to its initial value.
        #[doc(hidden)]
        #[allow(dead_code)]
        unsafe fn #reset_fn_name() {
            #reset_body
        }
    })
}

/// Generate a reset statement for a struct field (`self.field = ...`).
fn field_reset_expr(ty: &Type, field_name: &syn::Ident) -> TokenStream {
    let type_name = extract_type_name(ty);

    match type_name.as_deref() {
        // Integer types → 0
        Some("u8") | Some("u16") | Some("u32") | Some("u64") | Some("u128")
        | Some("i8") | Some("i16") | Some("i32") | Some("i64") | Some("i128") => {
            quote! { self.#field_name = 0; }
        }

        // bool → false
        Some("bool") => {
            quote! { self.#field_name = false; }
        }

        // Option<T> → None
        Some("Option") => {
            quote! { self.#field_name = None; }
        }

        // BoundedVec / BoundedMap → clear()
        Some("BoundedVec") | Some("BoundedMap") => {
            quote! { self.#field_name.clear(); }
        }

        // AtomicU32 / AtomicU64 → store(0)
        Some("AtomicU32") | Some("AtomicU64") => {
            quote! {
                self.#field_name.store(0, ::core::sync::atomic::Ordering::SeqCst);
            }
        }

        // Custom types → delegate to StepReset::reset
        _ => {
            quote! {
                rs_lang::StepReset::reset(&mut self.#field_name);
            }
        }
    }
}

/// Generate the reset body for a static item.
///
/// Atomics use `store(0, SeqCst)`. All other types re-assign the
/// original initializer expression via `addr_of_mut!`.
fn static_reset_expr(
    ty: &Type,
    name: &syn::Ident,
    init_expr: &syn::Expr,
) -> TokenStream {
    let type_name = extract_type_name(ty);

    match type_name.as_deref() {
        Some("AtomicU32") | Some("AtomicU64") => {
            quote! {
                #name.store(0, ::core::sync::atomic::Ordering::SeqCst);
            }
        }

        Some("BoundedVec") | Some("BoundedMap") => {
            quote! {
                (*::core::ptr::addr_of_mut!(#name)).clear();
            }
        }

        _ => {
            quote! {
                *::core::ptr::addr_of_mut!(#name) = #init_expr;
            }
        }
    }
}

/// Extract the outermost type name from a Type node.
fn extract_type_name(ty: &Type) -> Option<String> {
    match ty {
        Type::Path(type_path) => {
            let last_seg = type_path.path.segments.last()?;
            Some(last_seg.ident.to_string())
        }
        _ => None,
    }
}