solana_program_runtime/
mem_pool.rs

1use {
2    crate::execution_budget::{
3        MAX_CALL_DEPTH, MAX_HEAP_FRAME_BYTES, MAX_INSTRUCTION_STACK_DEPTH_SIMD_0268,
4        MIN_HEAP_FRAME_BYTES,
5    },
6    solana_sbpf::{aligned_memory::AlignedMemory, ebpf::HOST_ALIGN, vm::CallFrame},
7    std::{
8        array,
9        ops::{Deref, DerefMut},
10    },
11};
12
13trait Reset {
14    fn reset(&mut self);
15}
16
17struct Pool<T: Reset, const SIZE: usize> {
18    items: [Option<T>; SIZE],
19    next_empty: usize,
20}
21
22impl<T: Reset, const SIZE: usize> Pool<T, SIZE> {
23    fn new(items: [T; SIZE]) -> Self {
24        Self {
25            items: items.map(|i| Some(i)),
26            next_empty: SIZE,
27        }
28    }
29
30    fn len(&self) -> usize {
31        SIZE
32    }
33
34    fn get(&mut self) -> Option<T> {
35        if self.next_empty == 0 {
36            return None;
37        }
38        self.next_empty = self.next_empty.saturating_sub(1);
39        self.items
40            .get_mut(self.next_empty)
41            .and_then(|item| item.take())
42    }
43
44    fn put(&mut self, mut value: T) -> bool {
45        self.items
46            .get_mut(self.next_empty)
47            .map(|item| {
48                value.reset();
49                item.replace(value);
50                self.next_empty = self.next_empty.saturating_add(1);
51                true
52            })
53            .unwrap_or(false)
54    }
55}
56
57impl Reset for AlignedMemory<{ HOST_ALIGN }> {
58    fn reset(&mut self) {
59        self.as_slice_mut().fill(0)
60    }
61}
62
63pub struct CallFrameBuffer(Box<[CallFrame; MAX_CALL_DEPTH]>);
64
65impl Default for CallFrameBuffer {
66    fn default() -> Self {
67        let mut mem = Box::<[CallFrame; MAX_CALL_DEPTH]>::new_uninit();
68        let ptr = mem.as_mut_ptr().cast::<CallFrame>();
69        for i in 0..MAX_CALL_DEPTH {
70            unsafe { ptr.add(i).write(CallFrame::default()) }
71        }
72        Self(unsafe { mem.assume_init() })
73    }
74}
75
76impl Reset for CallFrameBuffer {
77    fn reset(&mut self) {
78        self.fill(CallFrame::default())
79    }
80}
81
82impl Deref for CallFrameBuffer {
83    type Target = [CallFrame];
84
85    fn deref(&self) -> &Self::Target {
86        self.0.as_slice()
87    }
88}
89
90impl DerefMut for CallFrameBuffer {
91    fn deref_mut(&mut self) -> &mut Self::Target {
92        self.0.as_mut_slice()
93    }
94}
95
96pub struct VmMemoryPool {
97    stack: Pool<AlignedMemory<{ HOST_ALIGN }>, MAX_INSTRUCTION_STACK_DEPTH_SIMD_0268>,
98    heap: Pool<AlignedMemory<{ HOST_ALIGN }>, MAX_INSTRUCTION_STACK_DEPTH_SIMD_0268>,
99    call_frame: Pool<CallFrameBuffer, MAX_INSTRUCTION_STACK_DEPTH_SIMD_0268>,
100}
101
102impl VmMemoryPool {
103    pub fn new() -> Self {
104        Self {
105            stack: Pool::new(array::from_fn(|_| {
106                #[allow(clippy::arithmetic_side_effects)]
107                AlignedMemory::zero_filled(solana_sbpf::vm::get_stack_frame_size() * MAX_CALL_DEPTH)
108            })),
109            heap: Pool::new(array::from_fn(|_| {
110                AlignedMemory::zero_filled(MAX_HEAP_FRAME_BYTES as usize)
111            })),
112            call_frame: Pool::new(array::from_fn(|_| CallFrameBuffer::default())),
113        }
114    }
115
116    pub fn stack_len(&self) -> usize {
117        self.stack.len()
118    }
119
120    pub fn heap_len(&self) -> usize {
121        self.heap.len()
122    }
123
124    #[allow(clippy::arithmetic_side_effects)]
125    pub fn get_stack(&mut self, size: usize) -> AlignedMemory<{ HOST_ALIGN }> {
126        debug_assert!(size == solana_sbpf::vm::get_stack_frame_size() * MAX_CALL_DEPTH);
127        self.stack
128            .get()
129            .unwrap_or_else(|| AlignedMemory::zero_filled(size))
130    }
131
132    pub fn put_stack(&mut self, stack: AlignedMemory<{ HOST_ALIGN }>) -> bool {
133        self.stack.put(stack)
134    }
135
136    pub fn get_heap(&mut self, heap_size: u32) -> AlignedMemory<{ HOST_ALIGN }> {
137        debug_assert!((MIN_HEAP_FRAME_BYTES..=MAX_HEAP_FRAME_BYTES).contains(&heap_size));
138        self.heap
139            .get()
140            .unwrap_or_else(|| AlignedMemory::zero_filled(MAX_HEAP_FRAME_BYTES as usize))
141    }
142
143    pub fn put_heap(&mut self, heap: AlignedMemory<{ HOST_ALIGN }>) -> bool {
144        let heap_size = heap.len();
145        debug_assert!(
146            heap_size >= MIN_HEAP_FRAME_BYTES as usize
147                && heap_size <= MAX_HEAP_FRAME_BYTES as usize
148        );
149        self.heap.put(heap)
150    }
151
152    pub fn get_call_frames(&mut self) -> CallFrameBuffer {
153        self.call_frame.get().unwrap_or_default()
154    }
155
156    pub fn put_call_frames(&mut self, call_frame: CallFrameBuffer) -> bool {
157        self.call_frame.put(call_frame)
158    }
159}
160
161#[cfg(test)]
162mod test {
163    use super::*;
164
165    #[derive(Debug, Eq, PartialEq)]
166    struct Item(u8, u8);
167    impl Reset for Item {
168        fn reset(&mut self) {
169            self.1 = 0;
170        }
171    }
172
173    #[test]
174    fn test_pool() {
175        let mut pool = Pool::<Item, 2>::new([Item(0, 1), Item(1, 1)]);
176        assert_eq!(pool.get(), Some(Item(1, 1)));
177        assert_eq!(pool.get(), Some(Item(0, 1)));
178        assert_eq!(pool.get(), None);
179        pool.put(Item(1, 1));
180        assert_eq!(pool.get(), Some(Item(1, 0)));
181        pool.put(Item(2, 2));
182        pool.put(Item(3, 3));
183        assert!(!pool.put(Item(4, 4)));
184        assert_eq!(pool.get(), Some(Item(3, 0)));
185        assert_eq!(pool.get(), Some(Item(2, 0)));
186        assert_eq!(pool.get(), None);
187    }
188}
solana_program_runtime/mem_pool.rs

solana_program_runtime/
mem_pool.rs
