Skip to main content

willow_core/
account.rs

1use pinocchio::account_info::AccountInfo;
2use pinocchio::program_error::ProgramError;
3extern crate alloc;
4use pinocchio::pubkey::Pubkey;
5use alloc::string::String;
6use core::slice::Iter;
7
8/// Trait implemented by your on-chain account structs (Vault, UserAccount, ...).
9/// Provides automatic serialization/deserialization capabilities.
10pub trait PinoAccount: Sized {
11    /// Load (immutable) from AccountInfo
12    fn load(ai: &AccountInfo) -> Result<Self, ProgramError> {
13        Self::from_account_info(ai)
14    }
15    /// Load as mutable (same as load if you have internal mutability)
16    fn load_mut(ai: &AccountInfo) -> Result<Self, ProgramError> {
17        Self::from_account_info(ai)
18    }
19    /// Save into AccountInfo
20    fn save(&self, ai: &AccountInfo) -> Result<(), ProgramError> {
21        self.save_to_account_info(ai)
22    }
23    /// Minimum expected length (if available)
24    fn expected_len() -> usize { 0 }
25    
26    /// Unpack account data from a byte slice
27    fn unpack_from_slice(data: &[u8]) -> Result<Self, ProgramError>;
28    /// Pack account data into a byte slice
29    fn pack_into_slice(&self, data: &mut [u8]) -> Result<(), ProgramError>;
30    
31    /// Create account from AccountInfo using unpack_from_slice
32    fn from_account_info(ai: &AccountInfo) -> Result<Self, ProgramError> {
33        let data = ai.try_borrow_data()?;
34        Self::unpack_from_slice(&data)
35    }
36    
37    /// Save account to AccountInfo using pack_into_slice
38    fn save_to_account_info(&self, ai: &AccountInfo) -> Result<(), ProgramError> {
39        let mut data = ai.try_borrow_mut_data()?;
40        self.pack_into_slice(&mut data)
41    }
42}
43
44#[derive(Clone, Debug)]
45pub struct FixedString<const N: usize>(pub String);
46
47#[derive(Clone, Debug)]
48pub struct FixedBytes<const N: usize>(pub [u8; N]);
49
50/// Willow-facing trait alias that mirrors PinoAccount methods.
51/// Use this in user code to avoid the legacy "pino" naming.
52pub trait WillowAccount: PinoAccount {
53    fn load(ai: &AccountInfo) -> Result<Self, ProgramError> {
54        <Self as PinoAccount>::load(ai)
55    }
56    fn load_mut(ai: &AccountInfo) -> Result<Self, ProgramError> {
57        <Self as PinoAccount>::load_mut(ai)
58    }
59    fn save(&self, ai: &AccountInfo) -> Result<(), ProgramError> {
60        <Self as PinoAccount>::save(self, ai)
61    }
62    fn expected_len() -> usize { <Self as PinoAccount>::expected_len() }
63    fn unpack_from_slice(data: &[u8]) -> Result<Self, ProgramError> {
64        <Self as PinoAccount>::unpack_from_slice(data)
65    }
66    fn pack_into_slice(&self, data: &mut [u8]) -> Result<(), ProgramError> {
67        <Self as PinoAccount>::pack_into_slice(self, data)
68    }
69    fn from_account_info(ai: &AccountInfo) -> Result<Self, ProgramError> {
70        <Self as PinoAccount>::from_account_info(ai)
71    }
72    fn save_to_account_info(&self, ai: &AccountInfo) -> Result<(), ProgramError> {
73        <Self as PinoAccount>::save_to_account_info(self, ai)
74    }
75}
76
77impl<T: PinoAccount> WillowAccount for T {}
78
79/// Extension trait for AccountInfo with helper methods
80pub trait AccountInfoExt {
81    /// Check if the account is a signer
82    fn is_signer(&self) -> bool;
83    /// Check if the account is owned by a specific program
84    fn is_owned_by(&self, program_id: &Pubkey) -> bool;
85    /// Get the key of the account
86    fn key(&self) -> &Pubkey;
87}
88
89/// Get the next account from an iterator
90pub fn next_account_info<'a>(
91    iter: &mut Iter<'a, AccountInfo>,
92) -> Result<&'a AccountInfo, ProgramError> {
93    iter.next().ok_or(ProgramError::NotEnoughAccountKeys)
94}
95
96impl AccountInfoExt for AccountInfo {
97    fn is_signer(&self) -> bool {
98        self.is_signer()
99    }
100    
101    fn is_owned_by(&self, program_id: &Pubkey) -> bool {
102        self.owner() == program_id
103    }
104    
105    fn key(&self) -> &Pubkey {
106        self.key()
107    }
108}
109
110/// Helper functions for serialization
111pub struct SerializationHelper;
112
113impl SerializationHelper {
114    /// Pack a u64 value into a byte slice
115    pub fn pack_u64(value: u64, data: &mut [u8], offset: &mut usize) {
116        data[*offset..*offset + 8].copy_from_slice(&value.to_le_bytes());
117        *offset += 8;
118    }
119    
120    /// Unpack a u64 value from a byte slice
121    pub fn unpack_u64(data: &[u8], offset: &mut usize) -> Result<u64, ProgramError> {
122        let value = u64::from_le_bytes(
123            data[*offset..*offset + 8]
124                .try_into()
125                .map_err(|_| ProgramError::InvalidAccountData)?
126        );
127        *offset += 8;
128        Ok(value)
129    }
130    
131    /// Pack a u32 value into a byte slice
132    pub fn pack_u32(value: u32, data: &mut [u8], offset: &mut usize) {
133        data[*offset..*offset + 4].copy_from_slice(&value.to_le_bytes());
134        *offset += 4;
135    }
136    
137    /// Unpack a u32 value from a byte slice
138    pub fn unpack_u32(data: &[u8], offset: &mut usize) -> Result<u32, ProgramError> {
139        let value = u32::from_le_bytes(
140            data[*offset..*offset + 4]
141                .try_into()
142                .map_err(|_| ProgramError::InvalidAccountData)?
143        );
144        *offset += 4;
145        Ok(value)
146    }
147    
148    /// Pack a u16 value into a byte slice
149    pub fn pack_u16(value: u16, data: &mut [u8], offset: &mut usize) {
150        data[*offset..*offset + 2].copy_from_slice(&value.to_le_bytes());
151        *offset += 2;
152    }
153    
154    /// Unpack a u16 value from a byte slice
155    pub fn unpack_u16(data: &[u8], offset: &mut usize) -> Result<u16, ProgramError> {
156        let value = u16::from_le_bytes(
157            data[*offset..*offset + 2]
158                .try_into()
159                .map_err(|_| ProgramError::InvalidAccountData)?
160        );
161        *offset += 2;
162        Ok(value)
163    }
164    
165    /// Pack an i64 value into a byte slice
166    pub fn pack_i64(value: i64, data: &mut [u8], offset: &mut usize) {
167        data[*offset..*offset + 8].copy_from_slice(&value.to_le_bytes());
168        *offset += 8;
169    }
170    
171    /// Unpack an i64 value from a byte slice
172    pub fn unpack_i64(data: &[u8], offset: &mut usize) -> Result<i64, ProgramError> {
173        let value = i64::from_le_bytes(
174            data[*offset..*offset + 8]
175                .try_into()
176                .map_err(|_| ProgramError::InvalidAccountData)?
177        );
178        *offset += 8;
179        Ok(value)
180    }
181    
182    /// Pack a Pubkey into a byte slice
183    pub fn pack_pubkey(value: &Pubkey, data: &mut [u8], offset: &mut usize) {
184        data[*offset..*offset + 32].copy_from_slice(value.as_ref());
185        *offset += 32;
186    }
187
188    /// Pack fixed-length bytes into a byte slice
189    pub fn pack_bytes_fixed<const N: usize>(value: &[u8; N], data: &mut [u8], offset: &mut usize) -> Result<(), ProgramError> {
190        if *offset + N > data.len() {
191            return Err(ProgramError::AccountDataTooSmall);
192        }
193        let end = *offset + N;
194        data[*offset..end].copy_from_slice(value);
195        *offset = end;
196        Ok(())
197    }
198
199    /// Unpack fixed-length bytes from a byte slice
200    pub fn unpack_bytes_fixed<const N: usize>(data: &[u8], offset: &mut usize) -> Result<[u8; N], ProgramError> {
201        if *offset + N > data.len() {
202            return Err(ProgramError::InvalidAccountData);
203        }
204        let out: [u8; N] = data[*offset..*offset + N]
205            .try_into()
206            .map_err(|_| ProgramError::InvalidAccountData)?;
207        *offset += N;
208        Ok(out)
209    }
210    
211    /// Unpack a Pubkey from a byte slice
212    pub fn unpack_pubkey(data: &[u8], offset: &mut usize) -> Result<Pubkey, ProgramError> {
213        let bytes: [u8; 32] = data[*offset..*offset + 32]
214            .try_into()
215            .map_err(|_| ProgramError::InvalidAccountData)?;
216        *offset += 32;
217        Ok(Pubkey::from(bytes))
218    }
219
220    /// Pack a fixed-length string (zero-padded)
221    pub fn pack_string_fixed<const N: usize>(value: &str, data: &mut [u8], offset: &mut usize) -> Result<(), ProgramError> {
222        let bytes = value.as_bytes();
223        if bytes.len() > N {
224            return Err(ProgramError::InvalidInstructionData);
225        }
226        if *offset + N > data.len() {
227            return Err(ProgramError::AccountDataTooSmall);
228        }
229        let end = *offset + N;
230        data[*offset..end].fill(0);
231        data[*offset..*offset + bytes.len()].copy_from_slice(bytes);
232        *offset = end;
233        Ok(())
234    }
235
236    /// Unpack a fixed-length string (trim trailing zeroes)
237    pub fn unpack_string_fixed<const N: usize>(data: &[u8], offset: &mut usize) -> Result<String, ProgramError> {
238        let bytes = Self::unpack_bytes_fixed::<N>(data, offset)?;
239        let end = bytes.iter().position(|b| *b == 0).unwrap_or(bytes.len());
240        let slice = &bytes[..end];
241        core::str::from_utf8(slice)
242            .map(|s| s.to_string())
243            .map_err(|_| ProgramError::InvalidAccountData)
244    }
245    
246    /// Pack a string with length prefix into a byte slice
247    pub fn pack_string(value: &str, data: &mut [u8], offset: &mut usize) -> Result<(), ProgramError> {
248        let bytes = value.as_bytes();
249        let len = bytes.len() as u32;
250        
251        // Check if we have enough space for length + string data
252        if *offset + 4 + bytes.len() > data.len() {
253            return Err(ProgramError::AccountDataTooSmall);
254        }
255        
256        // Write length as 4-byte little-endian
257        data[*offset..*offset + 4].copy_from_slice(&len.to_le_bytes());
258        *offset += 4;
259        
260        // Write string bytes
261        data[*offset..*offset + bytes.len()].copy_from_slice(bytes);
262        *offset += bytes.len();
263        
264        Ok(())
265    }
266    
267    /// Unpack a string with length prefix from a byte slice
268    pub fn unpack_string(data: &[u8], offset: &mut usize) -> Result<String, ProgramError> {
269        // Read length (4 bytes, little-endian)
270        if *offset + 4 > data.len() {
271            return Err(ProgramError::InvalidAccountData);
272        }
273        
274        let len_bytes: [u8; 4] = data[*offset..*offset + 4]
275            .try_into()
276            .map_err(|_| ProgramError::InvalidAccountData)?;
277        let len = u32::from_le_bytes(len_bytes) as usize;
278        *offset += 4;
279        
280        // Check if we have enough data for the string
281        if *offset + len > data.len() {
282            return Err(ProgramError::InvalidAccountData);
283        }
284        
285        // Read string bytes
286        let str_bytes = &data[*offset..*offset + len];
287        *offset += len;
288        
289        // Convert to string
290        let value = core::str::from_utf8(str_bytes)
291            .map_err(|_| ProgramError::InvalidAccountData)?
292            .to_string();
293            
294        Ok(value)
295    }
296    
297    /// Pack a vector with length prefix into a byte slice
298    pub fn pack_vec(value: &[u8], data: &mut [u8], offset: &mut usize) -> Result<(), ProgramError> {
299        let len = value.len() as u32;
300        
301        // Check if we have enough space for length + vector data
302        if *offset + 4 + value.len() > data.len() {
303            return Err(ProgramError::AccountDataTooSmall);
304        }
305        
306        // Write length as 4-byte little-endian
307        data[*offset..*offset + 4].copy_from_slice(&len.to_le_bytes());
308        *offset += 4;
309        
310        // Write vector bytes
311        data[*offset..*offset + value.len()].copy_from_slice(value);
312        *offset += value.len();
313        
314        Ok(())
315    }
316    
317    /// Unpack a vector with length prefix from a byte slice
318    pub fn unpack_vec(data: &[u8], offset: &mut usize) -> Result<Vec<u8>, ProgramError> {
319        // Read length (4 bytes, little-endian)
320        if *offset + 4 > data.len() {
321            return Err(ProgramError::InvalidAccountData);
322        }
323        
324        let len_bytes: [u8; 4] = data[*offset..*offset + 4]
325            .try_into()
326            .map_err(|_| ProgramError::InvalidAccountData)?;
327        let len = u32::from_le_bytes(len_bytes) as usize;
328        *offset += 4;
329        
330        // Check if we have enough data for the vector
331        if *offset + len > data.len() {
332            return Err(ProgramError::InvalidAccountData);
333        }
334        
335        // Read vector bytes
336        let vec_bytes = &data[*offset..*offset + len];
337        *offset += len;
338        
339        // Convert to vector
340        Ok(vec_bytes.to_vec())
341    }
342    
343    /// Pack a u8 value into a byte slice
344    pub fn pack_u8(value: u8, data: &mut [u8], offset: &mut usize) {
345        data[*offset] = value;
346        *offset += 1;
347    }
348    
349    /// Unpack a u8 value from a byte slice
350    pub fn unpack_u8(data: &[u8], offset: &mut usize) -> Result<u8, ProgramError> {
351        let value = data[*offset];
352        *offset += 1;
353        Ok(value)
354    }
355    
356    /// Pack an i8 value into a byte slice
357    pub fn pack_i8(value: i8, data: &mut [u8], offset: &mut usize) {
358        data[*offset] = value as u8;
359        *offset += 1;
360    }
361    
362    /// Unpack an i8 value from a byte slice
363    pub fn unpack_i8(data: &[u8], offset: &mut usize) -> Result<i8, ProgramError> {
364        let value = data[*offset] as i8;
365        *offset += 1;
366        Ok(value)
367    }
368    
369    /// Pack an i16 value into a byte slice
370    pub fn pack_i16(value: i16, data: &mut [u8], offset: &mut usize) {
371        data[*offset..*offset + 2].copy_from_slice(&value.to_le_bytes());
372        *offset += 2;
373    }
374    
375    /// Unpack an i16 value from a byte slice
376    pub fn unpack_i16(data: &[u8], offset: &mut usize) -> Result<i16, ProgramError> {
377        let value = i16::from_le_bytes(
378            data[*offset..*offset + 2]
379                .try_into()
380                .map_err(|_| ProgramError::InvalidAccountData)?
381        );
382        *offset += 2;
383        Ok(value)
384    }
385    
386    /// Pack an i32 value into a byte slice
387    pub fn pack_i32(value: i32, data: &mut [u8], offset: &mut usize) {
388        data[*offset..*offset + 4].copy_from_slice(&value.to_le_bytes());
389        *offset += 4;
390    }
391    
392    /// Unpack an i32 value from a byte slice
393    pub fn unpack_i32(data: &[u8], offset: &mut usize) -> Result<i32, ProgramError> {
394        let value = i32::from_le_bytes(
395            data[*offset..*offset + 4]
396                .try_into()
397                .map_err(|_| ProgramError::InvalidAccountData)?
398        );
399        *offset += 4;
400        Ok(value)
401    }
402    
403    /// Pack a bool value into a byte slice
404    pub fn pack_bool(value: bool, data: &mut [u8], offset: &mut usize) {
405        data[*offset] = if value { 1 } else { 0 };
406        *offset += 1;
407    }
408    
409    /// Unpack a bool value from a byte slice
410    pub fn unpack_bool(data: &[u8], offset: &mut usize) -> Result<bool, ProgramError> {
411        let value = match data[*offset] {
412            0 => false,
413            1 => true,
414            _ => return Err(ProgramError::InvalidAccountData),
415        };
416        *offset += 1;
417        Ok(value)
418    }
419
420    /// Pack a u128 value into a byte slice
421    pub fn pack_u128(value: u128, data: &mut [u8], offset: &mut usize) {
422        data[*offset..*offset + 16].copy_from_slice(&value.to_le_bytes());
423        *offset += 16;
424    }
425
426    /// Unpack a u128 value from a byte slice
427    pub fn unpack_u128(data: &[u8], offset: &mut usize) -> Result<u128, ProgramError> {
428        let value = u128::from_le_bytes(
429            data[*offset..*offset + 16]
430                .try_into()
431                .map_err(|_| ProgramError::InvalidAccountData)?
432        );
433        *offset += 16;
434        Ok(value)
435    }
436
437    /// Pack an i128 value into a byte slice
438    pub fn pack_i128(value: i128, data: &mut [u8], offset: &mut usize) {
439        data[*offset..*offset + 16].copy_from_slice(&value.to_le_bytes());
440        *offset += 16;
441    }
442
443    /// Unpack an i128 value from a byte slice
444    pub fn unpack_i128(data: &[u8], offset: &mut usize) -> Result<i128, ProgramError> {
445        let value = i128::from_le_bytes(
446            data[*offset..*offset + 16]
447                .try_into()
448                .map_err(|_| ProgramError::InvalidAccountData)?
449        );
450        *offset += 16;
451        Ok(value)
452    }
453}