wp-solana-core 0.1.1

//! Token account input types.
//!
//! Planning logic was migrated to `solana_token_toolkit`. This module retains
//! workspace input/output symbols for compatibility plus the
//! `plan_token_account_instructions` orchestrator that handles the legacy
//! Keypair-replacement flow.

use solana_pubkey::Pubkey;
use solana_sdk::signature::{Keypair, Signer};

use crate::{
    plan_error::PlanError,
    token::{account_state::TokenAccountState, NativeMintWrappingStrategy, WorkspacePlanConfig},
};

/// Strategy for token account creation.
#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)]
pub enum TokenAccountStrategy {
    /// Create token account without checking balance.
    WithoutBalance(Pubkey),
    /// Create token account, asserting that it holds at least `u64` units.
    WithBalance(Pubkey, u64),
}

impl TokenAccountStrategy {
    /// Return the mint pubkey associated with this strategy.
    pub fn mint(&self) -> Pubkey {
        match self {
            TokenAccountStrategy::WithoutBalance(m) => *m,
            TokenAccountStrategy::WithBalance(m, _) => *m,
        }
    }
}

/// Type alias for `solana_token_toolkit::TokenAccountPlan`.
pub type TokenAccountInstructions = solana_token_toolkit::TokenAccountPlan;

/// Pure-planner entry that returns a `TokenAccountInstructions` for the
/// supplied state + strategies + config.
///
/// When `config.wsol_strategy == Keypair`, the caller MUST also supply a
/// `keypair_for_wsol`. The toolkit generates an ephemeral keypair internally;
/// this fn swaps it for the caller-supplied one in all `create_instructions`,
/// `cleanup_instructions`, and `token_account_addresses` so downstream
/// signing uses the deterministic pubkey.
///
/// When `config.wsol_strategy == Seed`, this fn errors (Seed is async-only —
/// callers should use
/// `wp_solana_core::token::async_shell::prepare_seeded_wsol_legacy`).
pub fn plan_token_account_instructions(
    state: &TokenAccountState,
    strategies: &[TokenAccountStrategy],
    config: &WorkspacePlanConfig,
    keypair_for_wsol: Option<Keypair>,
) -> Result<TokenAccountInstructions, PlanError> {
    if config.wsol_strategy == NativeMintWrappingStrategy::Keypair && keypair_for_wsol.is_none() {
        return Err(PlanError::TokenAccountPlanning(
            "NativeMintWrappingStrategy::Keypair requires a keypair_for_wsol arg".into(),
        ));
    }

    let mut plan =
        crate::token::adapter::prepare_token_accounts_from_strategies(state, strategies, config)?;

    if config.wsol_strategy == NativeMintWrappingStrategy::Keypair {
        let keypair = keypair_for_wsol.expect("validated above: Keypair strategy requires keypair");
        apply_keypair_replacement(&mut plan, keypair)?;
    }

    Ok(plan)
}

/// Replace the toolkit-generated ephemeral wSOL keypair with the
/// caller-supplied one across all instructions, signers, and the
/// `token_account_addresses` map.
///
/// Toolkit v0.2 generates exactly one keypair signer when
/// `WrapSolStrategy::Keypair` is used (one wSOL account per plan).
/// 0 or >1 signers indicates either a misconfigured spec (no native SOL)
/// or a future toolkit contract change; both fail loudly here rather
/// than silently dropping signers.
pub(super) fn apply_keypair_replacement(
    plan: &mut TokenAccountInstructions,
    keypair: Keypair,
) -> Result<(), PlanError> {
    match plan.additional_signers.len() {
        0 => {
            return Err(PlanError::TokenAccountPlanning(
                "Keypair wrapping requested but toolkit produced no additional_signers".into(),
            ));
        }
        1 => {}
        n => {
            return Err(PlanError::TokenAccountPlanning(format!(
                "Keypair wrapping expected exactly 1 toolkit-generated signer, got {n}; workspace \
                 replacement logic only handles the single-wSOL case"
            )));
        }
    }
    let generated_pubkey = plan.additional_signers[0].pubkey();
    let replacement_pubkey = keypair.pubkey();
    for ix in &mut plan.create_instructions {
        for account in &mut ix.accounts {
            if account.pubkey == generated_pubkey {
                account.pubkey = replacement_pubkey;
            }
        }
    }
    for ix in &mut plan.cleanup_instructions {
        for account in &mut ix.accounts {
            if account.pubkey == generated_pubkey {
                account.pubkey = replacement_pubkey;
            }
        }
    }
    for address in plan.token_account_addresses.values_mut() {
        if *address == generated_pubkey {
            *address = replacement_pubkey;
        }
    }
    plan.additional_signers = vec![keypair];
    Ok(())
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use solana_instruction::{AccountMeta, Instruction};
    use solana_sdk::signature::{Keypair, Signer};
    use spl_token_interface::native_mint;

    use super::*;

    fn empty_plan() -> TokenAccountInstructions {
        TokenAccountInstructions {
            create_instructions: vec![],
            cleanup_instructions: vec![],
            additional_signers: vec![],
            token_account_addresses: HashMap::new(),
        }
    }

    fn ix_with_account(pk: Pubkey) -> Instruction {
        Instruction {
            program_id: Pubkey::new_unique(),
            // These tests assert pubkey replacement only; the synthetic
            // account does not need signer semantics.
            accounts: vec![AccountMeta::new(pk, false)],
            data: vec![],
        }
    }

    #[test]
    fn keypair_strategy_without_keypair_arg_errors_via_public_api() {
        let owner = Pubkey::new_unique();
        let state = TokenAccountState { owner, mints: HashMap::new() };
        let config = WorkspacePlanConfig {
            wsol_strategy: NativeMintWrappingStrategy::Keypair,
            ..Default::default()
        };
        let strategies = [TokenAccountStrategy::WithBalance(native_mint::ID, 1_000_000)];

        let err = plan_token_account_instructions(&state, &strategies, &config, None)
            .expect_err("expected Err when Keypair strategy is set without keypair_for_wsol");

        let msg = match err {
            PlanError::TokenAccountPlanning(s) => s,
            other => panic!("expected TokenAccountPlanning, got {other:?}"),
        };
        assert!(
            msg.contains("Keypair") && msg.contains("keypair_for_wsol"),
            "unexpected error message: {msg}"
        );
    }

    #[test]
    fn apply_keypair_replacement_errors_on_zero_signers() {
        let mut plan = empty_plan();
        let user_keypair = Keypair::new();

        let err = apply_keypair_replacement(&mut plan, user_keypair)
            .expect_err("expected Err when plan has zero additional_signers");

        let msg = match err {
            PlanError::TokenAccountPlanning(s) => s,
            other => panic!("expected TokenAccountPlanning, got {other:?}"),
        };
        assert!(msg.contains("no additional_signers"), "unexpected error message: {msg}");
    }

    #[test]
    fn apply_keypair_replacement_errors_on_multi_signer() {
        let mut plan = empty_plan();
        plan.additional_signers = vec![Keypair::new(), Keypair::new()];
        let user_keypair = Keypair::new();

        let err = apply_keypair_replacement(&mut plan, user_keypair)
            .expect_err("expected Err when plan has more than one additional_signer");

        let msg = match err {
            PlanError::TokenAccountPlanning(s) => s,
            other => panic!("expected TokenAccountPlanning, got {other:?}"),
        };
        assert!(
            msg.contains("got 2") && msg.contains("single-wSOL"),
            "unexpected error message: {msg}"
        );
    }

    #[test]
    fn apply_keypair_replacement_swaps_pubkeys_in_all_locations() {
        let toolkit_keypair = Keypair::new();
        let toolkit_pubkey = toolkit_keypair.pubkey();
        let unrelated_mint = Pubkey::new_unique();
        let unrelated_pubkey = Pubkey::new_unique();

        let mut plan = empty_plan();
        plan.create_instructions =
            vec![ix_with_account(toolkit_pubkey), ix_with_account(unrelated_pubkey)];
        plan.cleanup_instructions = vec![ix_with_account(toolkit_pubkey)];
        plan.additional_signers = vec![toolkit_keypair];
        plan.token_account_addresses.insert(native_mint::ID, toolkit_pubkey);
        plan.token_account_addresses.insert(unrelated_mint, unrelated_pubkey);

        let user_keypair = Keypair::new();
        let user_pubkey = user_keypair.pubkey();

        apply_keypair_replacement(&mut plan, user_keypair).expect("happy path should succeed");

        assert_eq!(plan.create_instructions[0].accounts[0].pubkey, user_pubkey);
        assert_eq!(plan.create_instructions[1].accounts[0].pubkey, unrelated_pubkey);
        assert_eq!(plan.cleanup_instructions[0].accounts[0].pubkey, user_pubkey);
        assert_eq!(plan.token_account_addresses[&native_mint::ID], user_pubkey);
        assert_eq!(plan.token_account_addresses[&unrelated_mint], unrelated_pubkey);
        assert_eq!(plan.additional_signers.len(), 1);
        assert_eq!(plan.additional_signers[0].pubkey(), user_pubkey);
    }
}