shank-parse-macro 2.0.0

use crate::idl::*;
use proc_macro2::{Ident, Literal, TokenStream};
use quote::{format_ident, quote};
use std::collections::HashMap;

pub fn generate(idl: &Idl, program_id_bytes: Option<[u8; 32]>) -> TokenStream {
    let mod_name = format_ident!("{}", idl.name);

    // Map every named (defined-able) type to the submodule it is generated into,
    // so cross-module references can be emitted as `super::<mod>::<Name>` paths.
    let module_of = build_module_map(idl);

    // Program ID constant (directly in the program module, not in a submodule)
    let id_token = program_id_bytes
        .map(|b| gen_program_id(&b))
        .unwrap_or_default();

    // types submodule: every entry from the IDL `types` section. Field-less
    // enums become "domain" enums (with `from_u8` / `to_u8`); enums with data
    // variants and all structs are emitted as borsh types. Everything derives
    // borsh and gets an inherent `try_from_slice`.
    let mut types_tokens = TokenStream::new();
    for ty in &idl.types {
        if ty.ty.kind == "enum" {
            if ty.ty.variants.iter().all(|v| v.fields.is_empty()) {
                types_tokens.extend(gen_domain_enum(ty));
            } else {
                types_tokens.extend(gen_data_enum(ty, &module_of));
            }
        } else {
            types_tokens.extend(gen_data_struct(ty, &module_of));
        }
    }

    // accounts submodule: account structs
    let accounts_tokens: TokenStream = idl
        .accounts
        .iter()
        .map(|a| gen_account(a, &module_of))
        .collect();

    // instructions submodule: instruction builder functions + their Accounts
    // structs (argument types live in the `types` submodule).
    let mut instructions_tokens = TokenStream::new();
    for ix in &idl.instructions {
        instructions_tokens.extend(gen_instruction(ix, &module_of));
    }

    // errors submodule: error code constants
    let errors_tokens: TokenStream = idl.errors.iter().map(gen_error).collect();

    // Only emit submodules that have content
    let types_mod = wrap_mod("types", types_tokens);
    let accounts_mod = wrap_mod("accounts", accounts_tokens);
    let instructions_mod = wrap_mod("instructions", instructions_tokens);
    let errors_mod = wrap_mod("errors", errors_tokens);

    quote! {
        pub mod #mod_name {
            #id_token
            #types_mod
            #accounts_mod
            #instructions_mod
            #errors_mod
        }
    }
}

fn wrap_mod(name: &str, tokens: TokenStream) -> TokenStream {
    if tokens.is_empty() {
        return quote! {};
    }
    let ident = format_ident!("{}", name);
    quote! { pub mod #ident { #tokens } }
}

/// name -> submodule it is generated into. Accounts live in `accounts`,
/// everything from the IDL `types` section lives in `types`.
fn build_module_map(idl: &Idl) -> HashMap<&str, &'static str> {
    let mut map: HashMap<&str, &'static str> = HashMap::new();
    for a in &idl.accounts {
        map.insert(a.name.as_str(), "accounts");
    }
    for t in &idl.types {
        map.insert(t.name.as_str(), "types");
    }
    map
}

fn gen_program_id(bytes: &[u8; 32]) -> TokenStream {
    let byte_literals: Vec<Literal> = bytes.iter().map(|b| Literal::u8_unsuffixed(*b)).collect();
    quote! {
        #[allow(dead_code)]
        pub const ID: ::shank_parse::__private::Pubkey =
            ::shank_parse::__private::Pubkey::new_from_array([#(#byte_literals),*]);
    }
}

/// `#[derive(...)]` line shared by every borsh-(de)serializable generated type.
/// borsh is reached through the `shank_parse` re-export, so the consuming crate
/// does not need its own borsh dependency.
fn borsh_derive() -> TokenStream {
    quote! {
        #[derive(
            Debug,
            Clone,
            PartialEq,
            ::shank_parse::__private::borsh::BorshSerialize,
            ::shank_parse::__private::borsh::BorshDeserialize,
        )]
        #[borsh(crate = "::shank_parse::__private::borsh")]
        #[allow(dead_code)]
    }
}

/// Generate a field-less ("domain") enum. It derives borsh (a field-less enum
/// serializes as a single u8 variant tag) and gets `from_u8` / `to_u8` helpers
/// plus the shared `try_from_slice`.
fn gen_domain_enum(ty: &IdlTypeDef) -> TokenStream {
    let name = format_ident!("{}", ty.name);
    let derive = borsh_derive();

    let variants: Vec<Ident> = ty
        .ty
        .variants
        .iter()
        .map(|v| format_ident!("{}", v.name))
        .collect();

    let from_u8_arms: Vec<TokenStream> = ty
        .ty
        .variants
        .iter()
        .enumerate()
        .map(|(i, v)| {
            let disc = Literal::u8_unsuffixed(i as u8);
            let vname = format_ident!("{}", v.name);
            quote! { #disc => Some(Self::#vname), }
        })
        .collect();

    let to_u8_arms: Vec<TokenStream> = ty
        .ty
        .variants
        .iter()
        .enumerate()
        .map(|(i, v)| {
            let disc = Literal::u8_unsuffixed(i as u8);
            let vname = format_ident!("{}", v.name);
            quote! { Self::#vname => #disc, }
        })
        .collect();

    let try_from_slice = gen_try_from_slice_impl(&name);

    quote! {
        #derive
        pub enum #name {
            #(#variants,)*
        }

        impl #name {
            #[allow(dead_code)]
            pub fn from_u8(v: u8) -> Option<Self> {
                match v {
                    #(#from_u8_arms)*
                    _ => None,
                }
            }

            #[allow(dead_code)]
            pub fn to_u8(&self) -> u8 {
                match self {
                    #(#to_u8_arms)*
                }
            }
        }

        #try_from_slice
    }
}

/// Generate an enum with data-carrying variants. Each variant's `defined`
/// fields become a tuple variant; field-less variants stay unit. Borsh tags the
/// variant with its declaration-order index as a leading u8.
fn gen_data_enum(ty: &IdlTypeDef, module_of: &HashMap<&str, &'static str>) -> TokenStream {
    let name = format_ident!("{}", ty.name);
    let derive = borsh_derive();

    let variants = ty.ty.variants.iter().map(|v| {
        let vname = format_ident!("{}", v.name);
        if v.fields.is_empty() {
            quote! { #vname }
        } else {
            let types = v.fields.iter().map(|f| defined_path(&f.defined, module_of));
            quote! { #vname(#(#types),*) }
        }
    });

    let try_from_slice = gen_try_from_slice_impl(&name);

    quote! {
        #derive
        pub enum #name {
            #(#variants,)*
        }

        #try_from_slice
    }
}

/// Generate an error constant: `pub const NOT_ADMIN: u32 = 0;`
fn gen_error(err: &IdlError) -> TokenStream {
    let snake = to_snake_case(&err.name);
    let const_name = format_ident!("{}", snake.to_uppercase());
    let code = Literal::u32_unsuffixed(err.code);
    quote! {
        #[allow(dead_code)]
        pub const #const_name: u32 = #code;
    }
}

/// A data struct from the IDL `types` section. Derives borsh and gets the
/// shared `try_from_slice`.
fn gen_data_struct(ty: &IdlTypeDef, module_of: &HashMap<&str, &'static str>) -> TokenStream {
    let name = format_ident!("{}", ty.name);
    let derive = borsh_derive();
    let fields = gen_struct_fields(&ty.ty.fields, module_of);
    let try_from_slice = gen_try_from_slice_impl(&name);

    quote! {
        #derive
        pub struct #name {
            #(#fields,)*
        }

        #try_from_slice
    }
}

fn gen_account(acc: &IdlTypeDef, module_of: &HashMap<&str, &'static str>) -> TokenStream {
    let name = format_ident!("{}", acc.name);
    let derive = borsh_derive();
    let fields = gen_struct_fields(&acc.ty.fields, module_of);

    quote! {
        #derive
        pub struct #name {
            #(#fields,)*
        }

        impl #name {
            /// Deserialize this account from raw account data. Trailing bytes
            /// (e.g. zero padding) after the struct are ignored.
            #[allow(dead_code)]
            pub fn from_account_data(data: &[u8]) -> ::std::io::Result<Self> {
                let mut buf: &[u8] = data;
                <Self as ::shank_parse::__private::borsh::BorshDeserialize>::deserialize_reader(
                    &mut buf,
                )
            }
        }
    }
}

fn gen_instruction(ix: &IdlInstruction, module_of: &HashMap<&str, &'static str>) -> TokenStream {
    let accounts_struct_name = format_ident!("{}Accounts", ix.name);
    let fn_name = format_ident!("{}", to_snake_case(&ix.name));
    let disc_val = Literal::u8_unsuffixed(ix.discriminant.value);
    let derive = borsh_derive();

    // Accounts struct fields
    let acc_fields = ix.accounts.iter().map(|a| {
        let fname = format_ident!("{}", to_snake_case(&a.name));
        quote! { pub #fname: ::shank_parse::__private::Pubkey }
    });

    // AccountMeta list
    let account_metas = ix.accounts.iter().map(|a| {
        let fname = format_ident!("{}", to_snake_case(&a.name));
        let is_signer = a.is_signer;
        if a.is_mut {
            quote! { ::shank_parse::__private::AccountMeta::new(accounts.#fname, #is_signer) }
        } else {
            quote! { ::shank_parse::__private::AccountMeta::new_readonly(accounts.#fname, #is_signer) }
        }
    });

    // Args: function params + borsh serialization (errors are propagated, never
    // unwrapped/expected).
    let mut arg_params = Vec::new();
    let mut arg_ser = Vec::new();
    for arg in &ix.args {
        let arg_name = format_ident!("{}", to_snake_case(&arg.name));
        let rust_ty = rust_type(&arg.ty, module_of);
        // `defined` args are passed by reference (they are usually structs);
        // everything else is passed by value. `receiver` is always `&T`.
        let (param, receiver) =
            if matches!(arg.ty, IdlType::Complex(IdlTypeComplex::Defined { .. })) {
                (quote! { #arg_name: &#rust_ty }, quote! { #arg_name })
            } else {
                (quote! { #arg_name: #rust_ty }, quote! { &#arg_name })
            };
        arg_params.push(param);
        arg_ser.push(quote! {
            ::shank_parse::__private::borsh::BorshSerialize::serialize(#receiver, &mut data)?;
        });
    }

    quote! {
        #derive
        pub struct #accounts_struct_name {
            #(#acc_fields,)*
        }

        #[allow(dead_code)]
        #[allow(clippy::too_many_arguments)]
        pub fn #fn_name(
            program_id: &::shank_parse::__private::Pubkey,
            accounts: &#accounts_struct_name,
            #(#arg_params,)*
        ) -> ::std::result::Result<::shank_parse::__private::Instruction, ::std::io::Error> {
            let mut data: Vec<u8> = ::std::vec![#disc_val];
            #(#arg_ser)*
            let account_metas = ::std::vec![
                #(#account_metas,)*
            ];
            Ok(::shank_parse::__private::Instruction::new_with_bytes(
                *program_id,
                &data,
                account_metas,
            ))
        }
    }
}

/// Inherent `try_from_slice` decoding a value from a borsh-encoded byte slice.
/// Returns `Result<Self, std::io::Error>` — no panics.
fn gen_try_from_slice_impl(name: &Ident) -> TokenStream {
    quote! {
        impl #name {
            #[allow(dead_code)]
            pub fn try_from_slice(data: &[u8]) -> ::std::io::Result<Self> {
                ::shank_parse::__private::borsh::from_slice(data)
            }
        }
    }
}

// ─── Helper generators ───────────────────────────────────────────────────────

fn gen_struct_fields(
    fields: &[IdlField],
    module_of: &HashMap<&str, &'static str>,
) -> Vec<TokenStream> {
    fields
        .iter()
        .map(|f| {
            let fname = format_ident!("{}", to_snake_case(&f.name));
            let ftype = rust_type(&f.ty, module_of);
            quote! { pub #fname: #ftype }
        })
        .collect()
}

/// Resolve a `defined` type name to a fully-qualified `super::<mod>::<Name>`
/// path. Falls back to a bare ident if the name is unknown.
fn defined_path(name: &str, module_of: &HashMap<&str, &'static str>) -> TokenStream {
    let ident = format_ident!("{}", name);
    match module_of.get(name) {
        Some(module) => {
            let module = format_ident!("{}", module);
            quote! { super::#module::#ident }
        }
        None => quote! { #ident },
    }
}

fn rust_type(ty: &IdlType, module_of: &HashMap<&str, &'static str>) -> TokenStream {
    match ty {
        IdlType::Primitive(p) => {
            let ident = format_ident!("{}", p);
            quote! { #ident }
        }
        IdlType::Complex(IdlTypeComplex::Array { array: (elem, size) }) => {
            let elem_ty = rust_type(elem, module_of);
            let size_lit = Literal::usize_unsuffixed(*size);
            quote! { [#elem_ty; #size_lit] }
        }
        IdlType::Complex(IdlTypeComplex::Defined { defined }) => defined_path(defined, module_of),
        IdlType::Complex(IdlTypeComplex::Option { option }) => {
            let inner = rust_type(option, module_of);
            quote! { Option<#inner> }
        }
    }
}

pub fn to_snake_case(s: &str) -> String {
    let mut result = String::new();
    for (i, c) in s.chars().enumerate() {
        if c.is_uppercase() && i > 0 {
            result.push('_');
        }
        for lc in c.to_lowercase() {
            result.push(lc);
        }
    }
    result
}