spl_stake_pool/
big_vec.rs

1//! Big vector type, used with vectors that can't be deserialized on-chain
2#![allow(clippy::arithmetic_side_effects)] // checked math involves too many compute units
3
4use {
5    arrayref::array_ref, borsh::BorshDeserialize, bytemuck::Pod,
6    solana_program_error::ProgramError, solana_program_memory::sol_memmove, std::mem,
7};
8
9/// Contains easy to use utilities for a big vector of Borsh-compatible types,
10/// to avoid managing the entire struct on-chain and blow through stack limits.
11pub struct BigVec<'data> {
12    /// Underlying data buffer, pieces of which are serialized
13    pub data: &'data mut [u8],
14}
15
16const VEC_SIZE_BYTES: usize = 4;
17
18impl BigVec<'_> {
19    /// Get the length of the vector
20    pub fn len(&self) -> u32 {
21        let vec_len = array_ref![self.data, 0, VEC_SIZE_BYTES];
22        u32::from_le_bytes(*vec_len)
23    }
24
25    /// Find out if the vector has no contents
26    pub fn is_empty(&self) -> bool {
27        self.len() == 0
28    }
29
30    /// Retain all elements that match the provided function, discard all others
31    pub fn retain<T: Pod, F: Fn(&[u8]) -> bool>(
32        &mut self,
33        predicate: F,
34    ) -> Result<(), ProgramError> {
35        let mut vec_len = self.len();
36        let mut removals_found = 0;
37        let mut dst_start_index = 0;
38
39        let data_start_index = VEC_SIZE_BYTES;
40        let data_end_index =
41            data_start_index.saturating_add((vec_len as usize).saturating_mul(mem::size_of::<T>()));
42        for start_index in (data_start_index..data_end_index).step_by(mem::size_of::<T>()) {
43            let end_index = start_index + mem::size_of::<T>();
44            let slice = &self.data[start_index..end_index];
45            if !predicate(slice) {
46                let gap = removals_found * mem::size_of::<T>();
47                if removals_found > 0 {
48                    // In case the compute budget is ever bumped up, allowing us
49                    // to use this safe code instead:
50                    // self.data.copy_within(dst_start_index + gap..start_index, dst_start_index);
51                    unsafe {
52                        sol_memmove(
53                            self.data[dst_start_index..start_index - gap].as_mut_ptr(),
54                            self.data[dst_start_index + gap..start_index].as_mut_ptr(),
55                            start_index - gap - dst_start_index,
56                        );
57                    }
58                }
59                dst_start_index = start_index - gap;
60                removals_found += 1;
61                vec_len -= 1;
62            }
63        }
64
65        // final memmove
66        if removals_found > 0 {
67            let gap = removals_found * mem::size_of::<T>();
68            // In case the compute budget is ever bumped up, allowing us
69            // to use this safe code instead:
70            //    self.data.copy_within(
71            //        dst_start_index + gap..data_end_index,
72            //        dst_start_index,
73            //    );
74            unsafe {
75                sol_memmove(
76                    self.data[dst_start_index..data_end_index - gap].as_mut_ptr(),
77                    self.data[dst_start_index + gap..data_end_index].as_mut_ptr(),
78                    data_end_index - gap - dst_start_index,
79                );
80            }
81        }
82
83        let vec_len_ref = &mut self.data[0..VEC_SIZE_BYTES];
84        borsh::to_writer(vec_len_ref, &vec_len)?;
85
86        Ok(())
87    }
88
89    /// Extracts a slice of the data types
90    pub fn deserialize_mut_slice<T: Pod>(
91        &mut self,
92        skip: usize,
93        len: usize,
94    ) -> Result<&mut [T], ProgramError> {
95        let vec_len = self.len();
96        let last_item_index = skip
97            .checked_add(len)
98            .ok_or(ProgramError::AccountDataTooSmall)?;
99        if last_item_index > vec_len as usize {
100            return Err(ProgramError::AccountDataTooSmall);
101        }
102
103        let start_index = VEC_SIZE_BYTES.saturating_add(skip.saturating_mul(mem::size_of::<T>()));
104        let end_index = start_index.saturating_add(len.saturating_mul(mem::size_of::<T>()));
105        bytemuck::try_cast_slice_mut(&mut self.data[start_index..end_index])
106            .map_err(|_| ProgramError::InvalidAccountData)
107    }
108
109    /// Extracts a slice of the data types
110    pub fn deserialize_slice<T: Pod>(&self, skip: usize, len: usize) -> Result<&[T], ProgramError> {
111        let vec_len = self.len();
112        let last_item_index = skip
113            .checked_add(len)
114            .ok_or(ProgramError::AccountDataTooSmall)?;
115        if last_item_index > vec_len as usize {
116            return Err(ProgramError::AccountDataTooSmall);
117        }
118
119        let start_index = VEC_SIZE_BYTES.saturating_add(skip.saturating_mul(mem::size_of::<T>()));
120        let end_index = start_index.saturating_add(len.saturating_mul(mem::size_of::<T>()));
121        bytemuck::try_cast_slice(&self.data[start_index..end_index])
122            .map_err(|_| ProgramError::InvalidAccountData)
123    }
124
125    /// Add new element to the end
126    pub fn push<T: Pod>(&mut self, element: T) -> Result<(), ProgramError> {
127        let vec_len_ref = &mut self.data[0..VEC_SIZE_BYTES];
128        let mut vec_len = u32::try_from_slice(vec_len_ref)?;
129
130        let start_index = VEC_SIZE_BYTES + vec_len as usize * mem::size_of::<T>();
131        let end_index = start_index + mem::size_of::<T>();
132
133        vec_len += 1;
134        borsh::to_writer(vec_len_ref, &vec_len)?;
135
136        if self.data.len() < end_index {
137            return Err(ProgramError::AccountDataTooSmall);
138        }
139        let element_ref = bytemuck::try_from_bytes_mut(
140            &mut self.data[start_index..start_index + mem::size_of::<T>()],
141        )
142        .map_err(|_| ProgramError::InvalidAccountData)?;
143        *element_ref = element;
144        Ok(())
145    }
146
147    /// Find matching data in the array
148    pub fn find<T: Pod, F: Fn(&[u8]) -> bool>(&self, predicate: F) -> Option<&T> {
149        let len = self.len() as usize;
150        let mut current = 0;
151        let mut current_index = VEC_SIZE_BYTES;
152        while current != len {
153            let end_index = current_index + mem::size_of::<T>();
154            let current_slice = &self.data[current_index..end_index];
155            if predicate(current_slice) {
156                return Some(bytemuck::from_bytes(current_slice));
157            }
158            current_index = end_index;
159            current += 1;
160        }
161        None
162    }
163
164    /// Find matching data in the array
165    pub fn find_mut<T: Pod, F: Fn(&[u8]) -> bool>(&mut self, predicate: F) -> Option<&mut T> {
166        let len = self.len() as usize;
167        let mut current = 0;
168        let mut current_index = VEC_SIZE_BYTES;
169        while current != len {
170            let end_index = current_index + mem::size_of::<T>();
171            let current_slice = &self.data[current_index..end_index];
172            if predicate(current_slice) {
173                return Some(bytemuck::from_bytes_mut(
174                    &mut self.data[current_index..end_index],
175                ));
176            }
177            current_index = end_index;
178            current += 1;
179        }
180        None
181    }
182}
183
184#[cfg(test)]
185mod tests {
186    use {super::*, bytemuck::Zeroable};
187
188    #[repr(C)]
189    #[derive(Debug, Copy, Clone, PartialEq, Pod, Zeroable)]
190    struct TestStruct {
191        value: [u8; 8],
192    }
193
194    impl TestStruct {
195        fn new(value: u8) -> Self {
196            let value = [value, 0, 0, 0, 0, 0, 0, 0];
197            Self { value }
198        }
199    }
200
201    fn from_slice<'data>(data: &'data mut [u8], vec: &[u8]) -> BigVec<'data> {
202        let mut big_vec = BigVec { data };
203        for element in vec {
204            big_vec.push(TestStruct::new(*element)).unwrap();
205        }
206        big_vec
207    }
208
209    fn check_big_vec_eq(big_vec: &BigVec, slice: &[u8]) {
210        assert!(big_vec
211            .deserialize_slice::<TestStruct>(0, big_vec.len() as usize)
212            .unwrap()
213            .iter()
214            .map(|x| &x.value[0])
215            .zip(slice.iter())
216            .all(|(a, b)| a == b));
217    }
218
219    #[test]
220    fn push() {
221        let mut data = [0u8; 4 + 8 * 3];
222        let mut v = BigVec { data: &mut data };
223        v.push(TestStruct::new(1)).unwrap();
224        check_big_vec_eq(&v, &[1]);
225        v.push(TestStruct::new(2)).unwrap();
226        check_big_vec_eq(&v, &[1, 2]);
227        v.push(TestStruct::new(3)).unwrap();
228        check_big_vec_eq(&v, &[1, 2, 3]);
229        assert_eq!(
230            v.push(TestStruct::new(4)).unwrap_err(),
231            ProgramError::AccountDataTooSmall
232        );
233    }
234
235    #[test]
236    fn retain() {
237        fn mod_2_predicate(data: &[u8]) -> bool {
238            u64::try_from_slice(data).unwrap() % 2 == 0
239        }
240
241        let mut data = [0u8; 4 + 8 * 4];
242        let mut v = from_slice(&mut data, &[1, 2, 3, 4]);
243        v.retain::<TestStruct, _>(mod_2_predicate).unwrap();
244        check_big_vec_eq(&v, &[2, 4]);
245    }
246
247    fn find_predicate(a: &[u8], b: u8) -> bool {
248        if a.len() != 8 {
249            false
250        } else {
251            a[0] == b
252        }
253    }
254
255    #[test]
256    fn find() {
257        let mut data = [0u8; 4 + 8 * 4];
258        let v = from_slice(&mut data, &[1, 2, 3, 4]);
259        assert_eq!(
260            v.find::<TestStruct, _>(|x| find_predicate(x, 1)),
261            Some(&TestStruct::new(1))
262        );
263        assert_eq!(
264            v.find::<TestStruct, _>(|x| find_predicate(x, 4)),
265            Some(&TestStruct::new(4))
266        );
267        assert_eq!(v.find::<TestStruct, _>(|x| find_predicate(x, 5)), None);
268    }
269
270    #[test]
271    fn find_mut() {
272        let mut data = [0u8; 4 + 8 * 4];
273        let mut v = from_slice(&mut data, &[1, 2, 3, 4]);
274        let test_struct = v
275            .find_mut::<TestStruct, _>(|x| find_predicate(x, 1))
276            .unwrap();
277        test_struct.value = [0; 8];
278        check_big_vec_eq(&v, &[0, 2, 3, 4]);
279        assert_eq!(v.find_mut::<TestStruct, _>(|x| find_predicate(x, 5)), None);
280    }
281
282    #[test]
283    fn deserialize_mut_slice() {
284        let mut data = [0u8; 4 + 8 * 4];
285        let mut v = from_slice(&mut data, &[1, 2, 3, 4]);
286        let slice = v.deserialize_mut_slice::<TestStruct>(1, 2).unwrap();
287        slice[0].value[0] = 10;
288        slice[1].value[0] = 11;
289        check_big_vec_eq(&v, &[1, 10, 11, 4]);
290        assert_eq!(
291            v.deserialize_mut_slice::<TestStruct>(1, 4).unwrap_err(),
292            ProgramError::AccountDataTooSmall
293        );
294        assert_eq!(
295            v.deserialize_mut_slice::<TestStruct>(4, 1).unwrap_err(),
296            ProgramError::AccountDataTooSmall
297        );
298    }
299}
spl_stake_pool/big_vec.rs

spl_stake_pool/
big_vec.rs
