magicsvm 0.2.0

use {
    crate::{
        magic::{
            delegation_action::post_delegation_actions,
            magic_program::{magic_instruction, MagicInstruction},
        },
        MagicSVM,
    },
    dlp_api::{
        args::PostDelegationActions, consts::DELEGATION_PROGRAM_ID, discriminator::DlpDiscriminator,
    },
    ephemeral_rollups_sdk::consts::MAGIC_PROGRAM_ID,
    litesvm::types::TransactionMetadata,
    magicblock_account::AccountSharedData,
    magicblock_magic_program_api::args::BaseActionArgs,
    solana_address::Address,
    solana_message::VersionedMessage,
    solana_transaction::TransactionError,
    std::collections::HashMap,
};

/// A request to delegate an account, together with any post-delegation actions
/// to run once the delegation has been applied.
#[derive(Debug)]
pub struct DelegationEffect {
    /// The account being delegated.
    pub account: Address,
    /// Encrypted instructions to execute after delegation, if any.
    pub actions: Option<PostDelegationActions>,
}

/// A scheduled Magic Action to run on the base layer after commit / undelegate.
#[derive(Debug)]
pub(crate) struct ScheduledBaseAction {
    pub action: BaseActionArgs,
    pub escrow_authority: Address,
}

/// A lifecycle change to an ephemeral account requested by the magic program.
#[derive(Debug)]
pub enum EphemeralAccountEffect {
    /// Create a new ephemeral account, funded by `sponsor`.
    Create {
        sponsor: Address,
        account: Address,
        owner: Address,
        data_len: u32,
    },
    /// Resize an existing ephemeral account, settling the rent difference with
    /// `sponsor`.
    Resize {
        sponsor: Address,
        account: Address,
        new_data_len: u32,
    },
    /// Close an ephemeral account, refunding its rent to `sponsor`.
    Close { sponsor: Address, account: Address },
}

/// The set of MagicBlock side effects collected from a transaction, to be
/// applied to a [`MagicSVM`] after the transaction itself executes.
#[derive(Default, Debug)]
pub struct MagicTransactionEffects {
    /// Accounts to delegate to the ephemeral rollup.
    pub delegated_accounts: Vec<DelegationEffect>,
    /// Accounts whose ephemeral state should be committed to the base layer.
    pub committed_accounts: Vec<Address>,
    /// Accounts to commit and then undelegate.
    pub undelegated_accounts: Vec<Address>,
    /// Base-layer Magic Actions scheduled by commit / intent instructions.
    pub(crate) base_actions: Vec<ScheduledBaseAction>,
    /// Ephemeral account lifecycle changes (create/resize/close).
    pub ephemeral_accounts: Vec<EphemeralAccountEffect>,
}

impl MagicTransactionEffects {
    /// Collects the MagicBlock effects of a base-layer transaction by scanning
    /// its instructions for delegation-program calls (delegate, commit,
    /// undelegate).
    pub fn from_message(message: &VersionedMessage) -> Self {
        let account_keys = message.static_account_keys();
        let mut effects = Self::default();

        for instruction in message.instructions() {
            Self::record_dlp_instruction(
                instruction.program_id_index,
                &instruction.accounts,
                &instruction.data,
                account_keys,
                &mut effects,
            );
        }

        effects
    }

    /// Collects base-layer MagicBlock effects from both top-level instructions
    /// and inner (CPI) instructions recorded in `meta`.
    pub fn from_message_and_metadata(
        message: &VersionedMessage,
        meta: &TransactionMetadata,
    ) -> Self {
        let account_keys = message.static_account_keys();
        let mut effects = Self::from_message(message);

        for inner_instructions in &meta.inner_instructions {
            for inner_instruction in inner_instructions {
                let instruction = &inner_instruction.instruction;
                Self::record_dlp_instruction(
                    instruction.program_id_index,
                    &instruction.accounts,
                    &instruction.data,
                    account_keys,
                    &mut effects,
                );
            }
        }

        effects
    }

    fn from_ephemeral_message(message: &VersionedMessage) -> Self {
        let account_keys = message.static_account_keys();
        let mut effects = Self::default();

        for instruction in message.instructions() {
            Self::record_magic_instruction(
                instruction.program_id_index,
                &instruction.accounts,
                &instruction.data,
                account_keys,
                None,
                &mut effects,
            );
        }

        effects
    }

    /// Collects the MagicBlock effects of an ephemeral transaction, scanning
    /// both its top-level instructions and the inner (CPI) instructions recorded
    /// in `meta` for magic-program calls.
    pub fn from_ephemeral_message_and_metadata(
        message: &VersionedMessage,
        meta: &TransactionMetadata,
    ) -> Self {
        let account_keys = message.static_account_keys();
        let mut effects = Self::from_ephemeral_message(message);

        for (outer_index, inner_instructions) in meta.inner_instructions.iter().enumerate() {
            for (inner_index, inner_instruction) in inner_instructions.iter().enumerate() {
                let instruction = &inner_instruction.instruction;
                let caller_program_id = caller_program_id_for_inner_instruction(
                    message,
                    inner_instructions,
                    outer_index,
                    inner_index,
                );
                Self::record_magic_instruction(
                    instruction.program_id_index,
                    &instruction.accounts,
                    &instruction.data,
                    account_keys,
                    caller_program_id,
                    &mut effects,
                );
            }
        }

        effects
    }

    /// Applies the collected ephemeral account lifecycle effects
    /// (create/resize/close) to `svm`.
    ///
    /// `pre_accounts` holds the ephemeral account states captured before the
    /// transaction ran, used to recover accounts that the transaction itself
    /// removed. Returns the first error encountered, if any.
    pub fn apply_ephemeral_account(
        &self,
        svm: &mut MagicSVM,
        pre_accounts: &HashMap<Address, AccountSharedData>,
    ) -> Result<(), TransactionError> {
        for effect in &self.ephemeral_accounts {
            match *effect {
                EphemeralAccountEffect::Create {
                    sponsor,
                    account,
                    owner,
                    data_len,
                } => svm.create_ephemeral_account(sponsor, account, owner, data_len)?,
                EphemeralAccountEffect::Resize {
                    sponsor,
                    account,
                    new_data_len,
                } => svm.resize_ephemeral_account(sponsor, account, new_data_len, pre_accounts)?,
                EphemeralAccountEffect::Close { sponsor, account } => {
                    svm.close_ephemeral_account(sponsor, account, pre_accounts)?
                }
            }
        }
        Ok(())
    }

    /// Applies the collected delegation, commit and undelegation effects to the
    /// base layer of `svm`, running any post-delegation actions and keeping the
    /// two ledgers in sync. Returns the first error encountered, if any.
    pub fn apply_base(
        self,
        svm: &mut MagicSVM,
        fee_payer: Address,
    ) -> Result<(), TransactionError> {
        for delegation in self.delegated_accounts {
            if svm.delegate_account(delegation.account).is_ok() {
                svm.run_post_delegation_actions(delegation.actions, fee_payer)?;
            }
        }
        for account in self.committed_accounts {
            svm.commit_account(account);
        }
        for account in self.undelegated_accounts {
            svm.commit_account(account);
            svm.undelegate_account(account);
        }
        svm.run_post_commit_actions(&self.base_actions)?;
        Ok(())
    }

    fn record_dlp_instruction(
        program_id_index: u8,
        instruction_accounts: &[u8],
        instruction_data: &[u8],
        account_keys: &[Address],
        effects: &mut Self,
    ) {
        let Some(program_id) = account_keys.get(usize::from(program_id_index)) else {
            return;
        };
        if *program_id != DELEGATION_PROGRAM_ID {
            return;
        }
        let Some(discriminator) = instruction_discriminator(instruction_data) else {
            return;
        };
        match discriminator {
            DlpDiscriminator::Delegate
            | DlpDiscriminator::DelegateWithAnyValidator
            | DlpDiscriminator::DelegateWithActions => {
                if let Some(account) = instruction_accounts
                    .get(1)
                    .and_then(|index| account_keys.get(usize::from(*index)))
                {
                    effects.delegated_accounts.push(DelegationEffect {
                        account: *account,
                        actions: post_delegation_actions(discriminator, instruction_data),
                    });
                }
            }
            DlpDiscriminator::CommitState
            | DlpDiscriminator::Finalize
            | DlpDiscriminator::CommitStateFromBuffer
            | DlpDiscriminator::CommitDiff
            | DlpDiscriminator::CommitDiffFromBuffer
            | DlpDiscriminator::CommitFinalize
            | DlpDiscriminator::CommitFinalizeFromBuffer => {
                if let Some(account) = instruction_accounts
                    .get(1)
                    .and_then(|index| account_keys.get(usize::from(*index)))
                {
                    effects.committed_accounts.push(*account);
                }
            }
            DlpDiscriminator::Undelegate | DlpDiscriminator::UndelegateConfinedAccount => {
                if let Some(account) = instruction_accounts
                    .get(1)
                    .and_then(|index| account_keys.get(usize::from(*index)))
                {
                    effects.undelegated_accounts.push(*account);
                }
            }
            _ => {}
        }
    }

    fn record_magic_instruction(
        program_id_index: u8,
        instruction_accounts: &[u8],
        instruction_data: &[u8],
        account_keys: &[Address],
        caller_program_id: Option<Address>,
        effects: &mut Self,
    ) {
        let Some(program_id) = account_keys.get(usize::from(program_id_index)) else {
            return;
        };
        if *program_id != MAGIC_PROGRAM_ID {
            return;
        }
        let Ok(magic_ix) = magic_instruction(instruction_data) else {
            return;
        };

        let accounts = instruction_accounts
            .iter()
            .skip(2)
            .filter_map(|index| account_keys.get(usize::from(*index)).copied());

        match magic_ix {
            MagicInstruction::ScheduleCommit
            | MagicInstruction::ScheduleCommitFinalize {
                request_undelegation: false,
            } => effects.committed_accounts.extend(accounts),
            MagicInstruction::ScheduleCommitAndUndelegate
            | MagicInstruction::ScheduleCommitFinalize {
                request_undelegation: true,
            } => effects.undelegated_accounts.extend(accounts),
            MagicInstruction::ScheduleBaseIntent {
                committed_accounts,
                undelegated_accounts,
                base_actions,
            }
            | MagicInstruction::ScheduleIntentBundle {
                committed_accounts,
                undelegated_accounts,
                base_actions,
            } => record_scheduled_intent(
                instruction_accounts,
                account_keys,
                committed_accounts,
                undelegated_accounts,
                base_actions,
                effects,
            ),
            MagicInstruction::CreateEphemeralAccount { data_len } => {
                let Some(owner) = caller_program_id else {
                    return;
                };
                let Some((sponsor, account)) =
                    ephemeral_sponsor_and_account(instruction_accounts, account_keys)
                else {
                    return;
                };
                effects
                    .ephemeral_accounts
                    .push(EphemeralAccountEffect::Create {
                        sponsor,
                        account,
                        owner,
                        data_len,
                    });
            }
            MagicInstruction::ResizeEphemeralAccount { new_data_len } => {
                let Some((sponsor, account)) =
                    ephemeral_sponsor_and_account(instruction_accounts, account_keys)
                else {
                    return;
                };
                effects
                    .ephemeral_accounts
                    .push(EphemeralAccountEffect::Resize {
                        sponsor,
                        account,
                        new_data_len,
                    });
            }
            MagicInstruction::CloseEphemeralAccount => {
                let Some((sponsor, account)) =
                    ephemeral_sponsor_and_account(instruction_accounts, account_keys)
                else {
                    return;
                };
                effects
                    .ephemeral_accounts
                    .push(EphemeralAccountEffect::Close { sponsor, account });
            }
            MagicInstruction::Noop => {}
        }
    }
}

fn record_scheduled_intent(
    instruction_accounts: &[u8],
    account_keys: &[Address],
    committed_accounts: Vec<u8>,
    undelegated_accounts: Vec<u8>,
    base_actions: Vec<BaseActionArgs>,
    effects: &mut MagicTransactionEffects,
) {
    effects
        .committed_accounts
        .extend(committed_accounts.into_iter().filter_map(|index| {
            instruction_accounts
                .get(usize::from(index))
                .and_then(|account_index| account_keys.get(usize::from(*account_index)))
                .copied()
        }));
    effects
        .undelegated_accounts
        .extend(undelegated_accounts.into_iter().filter_map(|index| {
            instruction_accounts
                .get(usize::from(index))
                .and_then(|account_index| account_keys.get(usize::from(*account_index)))
                .copied()
        }));
    effects
        .base_actions
        .extend(base_actions.into_iter().filter_map(|action| {
            let escrow_authority = instruction_accounts
                .get(usize::from(action.escrow_authority))
                .and_then(|account_index| account_keys.get(usize::from(*account_index)))
                .copied()?;
            Some(ScheduledBaseAction {
                action,
                escrow_authority,
            })
        }));
}

fn ephemeral_sponsor_and_account(
    instruction_accounts: &[u8],
    account_keys: &[Address],
) -> Option<(Address, Address)> {
    Some((
        *account_keys.get(usize::from(*instruction_accounts.first()?))?,
        *account_keys.get(usize::from(*instruction_accounts.get(1)?))?,
    ))
}

fn caller_program_id_for_inner_instruction(
    message: &VersionedMessage,
    inner_instructions: &[solana_message::inner_instruction::InnerInstruction],
    outer_index: usize,
    inner_index: usize,
) -> Option<Address> {
    let account_keys = message.static_account_keys();
    let current_stack_height = inner_instructions.get(inner_index)?.stack_height;
    inner_instructions[..inner_index]
        .iter()
        .rev()
        .find(|inner_instruction| inner_instruction.stack_height < current_stack_height)
        .and_then(|inner_instruction| {
            account_keys.get(usize::from(inner_instruction.instruction.program_id_index))
        })
        .or_else(|| {
            message
                .instructions()
                .get(outer_index)
                .and_then(|ix| account_keys.get(usize::from(ix.program_id_index)))
        })
        .copied()
}

fn instruction_discriminator(data: &[u8]) -> Option<DlpDiscriminator> {
    let bytes = data.get(..8)?;
    let discriminator = u64::from_le_bytes(bytes.try_into().ok()?);
    u8::try_from(discriminator).ok()?.try_into().ok()
}