init_space 0.0.3

//! # init_space_derive
//!
//! A procedural macro that automatically calculates the space required to store a struct
//! in Solana programs (or Borsh-serialized data).
//!
//! - Computes size for fixed types like `u8`, `u64`, `bool`, etc.
//! - Supports dynamic types like `String` and `Vec<u8>` by requiring a `#[max_len = N]` attribute.
//!
//! ## Example
//!
//! ```rust
//! use init_space_derive::InitSpace;
//! use borsh::{BorshSerialize, BorshDeserialize};
//!
//! #[derive(BorshSerialize, BorshDeserialize, InitSpace)]
//! pub struct Movie {
//!     pub title: u64,
//!     pub rating: u8,
//!     #[max_len = 1000]
//!     pub description: Vec<u8>,
//! }
//!
//! let space_needed = Movie::space();
//! assert_eq!(space_needed, 13 + 1004); // 13 for fixed, 4+1000 for Vec
//! ```




use proc_macro::TokenStream;
use quote::quote;
use syn::{
    parse_macro_input, DeriveInput, Data, Fields, Type, Attribute, Meta,
    Expr, ExprLit, Lit,
};

#[proc_macro_derive(InitSpace, attributes(max_len))]
pub fn derive_init_space(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    let name = input.ident;

    let mut total_size = quote! { 0 };

    if let Data::Struct(data_struct) = input.data {
        if let Fields::Named(fields_named) = data_struct.fields {
            for field in &fields_named.named {
                let ty = &field.ty;
                let attrs = &field.attrs;

                // If #[max_len = N] is provided, handle dynamic size
                if let Some(max_len)= get_max_len(attrs) {
                       // Only allow #[max_len = N] on Vec or String
                    if is_dynamic_type(ty) {
                      total_size = quote! { #total_size + 4 + #max_len };
                        continue;
                    } else {
                      panic!("`#[max_len = N]` is only allowed on Vec or String types");
                  }
                }

                // Otherwise, use fixed size or fallback to 4 for unknown
                let field_size = if let Some(size) = get_fixed_size(ty) {
                    quote! { #size }
                } else {
                    quote! { 4 }
                };

                total_size = quote! { #total_size + #field_size };
            }
        }
    }

    let expanded = quote! {
        impl #name {
            pub fn space() -> usize {
                #total_size
            }
        }
    };

    TokenStream::from(expanded)
}

// Known fixed-size types
fn get_fixed_size(ty: &Type) -> Option<usize> {
    if let Type::Path(typepath) = ty {
        //gets the last segment of the type path
        let ident = &typepath.path.segments.last()?.ident;
        match ident.to_string().as_str() {
            "u8" => Some(1),
            "u16" => Some(2),
            "u32" => Some(4),
            "u64" => Some(8),
            "u128" => Some(16),
            "i8" => Some(1),
            "i16" => Some(2),
            "i32" => Some(4),
            "i64" => Some(8),
            "i128" => Some(16),
            "bool" => Some(1),
            "Pubkey" => Some(32),
            _ => None,
        }
    } else {
        None
    }
}

// Parses #[max_len = 1000]
fn get_max_len(attrs: &[Attribute]) -> Option<usize> {
    for attr in attrs {
        if attr.path().is_ident("max_len") {
            if let Ok(meta) = attr.meta.clone().require_name_value() {
                if let Expr::Lit(ExprLit {
                    lit: Lit::Int(lit_int),
                    ..
                }) = &meta.value
                //here the meta.value is passed as arg to be used for lint_int
                {
                    return lit_int.base10_parse::<usize>().ok();
                }
            }
        }
    }
    None
}

fn is_dynamic_type(ty: &Type) -> bool {
    if let Type::Path(typepath) = ty {
        //return the last segment of the the type
        if let Some(segment) = typepath.path.segments.last() {
            let ident = segment.ident.to_string();
            return ident == "Vec" || ident == "String";
        }
    }
    false
}