Skip to main content

willow_core/
context.rs

1use core::slice;
2use pinocchio::account_info::AccountInfo;
3use pinocchio::program_error::ProgramError;
4use pinocchio::pubkey::Pubkey;
5use crate::account::PinoAccount;
6
7pub struct PinoContext<'a> {
8    pub iter: &'a mut slice::Iter<'a, AccountInfo>,
9    /// Cache for validated accounts to avoid redundant validations
10    validated_accounts: [Option<&'a AccountInfo>; 8],
11}
12
13impl<'a> PinoContext<'a> {
14    pub fn new(iter: &'a mut slice::Iter<'a, AccountInfo>) -> Self {
15        Self {
16            iter,
17            validated_accounts: [None; 8],
18        }
19    }
20
21    /// Pull the next AccountInfo and load as T
22    pub fn next<T: PinoAccount>(&mut self) -> Result<T, ProgramError> {
23        let ai = self.iter.next().ok_or(ProgramError::NotEnoughAccountKeys)?;
24        T::load(ai)
25    }
26
27    /// Pull the next AccountInfo and load_mut as T and return the AccountInfo alongside for save
28    pub fn next_mut<T: PinoAccount>(&mut self) -> Result<(T, &'a AccountInfo), ProgramError> {
29        let ai = self.iter.next().ok_or(ProgramError::NotEnoughAccountKeys)?;
30        let v = T::load_mut(ai)?;
31        Ok((v, ai))
32    }
33
34    /// Get a validated account that has already been checked for common requirements
35    /// 
36    /// This method allows caching of validated accounts to avoid redundant validations
37    /// such as signature checks, owner checks, etc.
38    pub fn get_validated_account(&self, index: usize) -> Option<&'a AccountInfo> {
39        if index < self.validated_accounts.len() {
40            self.validated_accounts[index]
41        } else {
42            None
43        }
44    }
45
46    /// Store a validated account in the cache
47    /// 
48    /// This method should be called after validating an account's requirements
49    /// such as signature checks, owner checks, etc.
50    pub fn set_validated_account(&mut self, index: usize, account: &'a AccountInfo) {
51        if index < self.validated_accounts.len() {
52            self.validated_accounts[index] = Some(account);
53        }
54    }
55
56    /// Validate that an account is a signer and cache it
57    /// 
58    /// This method checks if the account is a signer and caches it for future use
59    /// to avoid redundant signature validations.
60    pub fn validate_signer(&mut self, index: usize) -> Result<&'a AccountInfo, ProgramError> {
61        let account = self.iter.next().ok_or(ProgramError::NotEnoughAccountKeys)?;
62        
63        if !account.is_signer() {
64            return Err(ProgramError::MissingRequiredSignature);
65        }
66        
67        self.set_validated_account(index, account);
68        Ok(account)
69    }
70
71    /// Validate that an account is owned by a specific program and cache it
72    /// 
73    /// This method checks if the account is owned by the specified program ID
74    /// and caches it for future use to avoid redundant owner validations.
75    pub fn validate_owner(&mut self, index: usize, program_id: &Pubkey) -> Result<&'a AccountInfo, ProgramError> {
76        let account = self.iter.next().ok_or(ProgramError::NotEnoughAccountKeys)?;
77        
78        if account.owner() != program_id {
79            return Err(ProgramError::IncorrectProgramId);
80        }
81        
82        self.set_validated_account(index, account);
83        Ok(account)
84    }
85
86    /// Validate that an account is a signer and owned by a specific program, then cache it
87    /// 
88    /// This method combines signature and owner validation and caches the result
89    /// to avoid redundant validations.
90    pub fn validate_signer_and_owner(&mut self, index: usize, program_id: &Pubkey) -> Result<&'a AccountInfo, ProgramError> {
91        let account = self.iter.next().ok_or(ProgramError::NotEnoughAccountKeys)?;
92        
93        if !account.is_signer() {
94            return Err(ProgramError::MissingRequiredSignature);
95        }
96        
97        if account.owner() != program_id {
98            return Err(ProgramError::IncorrectProgramId);
99        }
100        
101        self.set_validated_account(index, account);
102        Ok(account)
103    }
104}