solana_program_runtime/
mem_pool.rs1use {
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, len: usize);
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, len: usize) -> bool {
45 self.items
46 .get_mut(self.next_empty)
47 .map(|item| {
48 value.reset(len);
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, len: usize) {
59 let slice = self.as_slice_mut();
60 let len = len.min(slice.len());
61 if let Some(head) = slice.get_mut(..len) {
62 head.fill(0);
63 }
64 }
65}
66
67pub struct CallFrameBuffer(Box<[CallFrame; MAX_CALL_DEPTH]>);
68
69impl Default for CallFrameBuffer {
70 fn default() -> Self {
71 let mut mem = Box::<[CallFrame; MAX_CALL_DEPTH]>::new_uninit();
72 let ptr = mem.as_mut_ptr().cast::<CallFrame>();
73 for i in 0..MAX_CALL_DEPTH {
74 unsafe { ptr.add(i).write(CallFrame::default()) }
75 }
76 Self(unsafe { mem.assume_init() })
77 }
78}
79
80impl Reset for CallFrameBuffer {
81 fn reset(&mut self, _len: usize) {
82 self.fill(CallFrame::default())
83 }
84}
85
86impl Deref for CallFrameBuffer {
87 type Target = [CallFrame];
88
89 fn deref(&self) -> &Self::Target {
90 self.0.as_slice()
91 }
92}
93
94impl DerefMut for CallFrameBuffer {
95 fn deref_mut(&mut self) -> &mut Self::Target {
96 self.0.as_mut_slice()
97 }
98}
99
100pub struct VmMemoryPool {
101 stack: Pool<AlignedMemory<{ HOST_ALIGN }>, MAX_INSTRUCTION_STACK_DEPTH_SIMD_0268>,
102 heap: Pool<AlignedMemory<{ HOST_ALIGN }>, MAX_INSTRUCTION_STACK_DEPTH_SIMD_0268>,
103 call_frame: Pool<CallFrameBuffer, MAX_INSTRUCTION_STACK_DEPTH_SIMD_0268>,
104}
105
106impl VmMemoryPool {
107 pub fn new() -> Self {
108 Self {
109 stack: Pool::new(array::from_fn(|_| {
110 #[allow(clippy::arithmetic_side_effects)]
111 AlignedMemory::zero_filled(solana_sbpf::vm::get_stack_frame_size() * MAX_CALL_DEPTH)
112 })),
113 heap: Pool::new(array::from_fn(|_| {
114 AlignedMemory::zero_filled(MAX_HEAP_FRAME_BYTES as usize)
115 })),
116 call_frame: Pool::new(array::from_fn(|_| CallFrameBuffer::default())),
117 }
118 }
119
120 pub fn stack_len(&self) -> usize {
121 self.stack.len()
122 }
123
124 pub fn heap_len(&self) -> usize {
125 self.heap.len()
126 }
127
128 #[allow(clippy::arithmetic_side_effects)]
129 pub fn get_stack(&mut self, size: usize) -> AlignedMemory<{ HOST_ALIGN }> {
130 debug_assert!(size == solana_sbpf::vm::get_stack_frame_size() * MAX_CALL_DEPTH);
131 self.stack
132 .get()
133 .unwrap_or_else(|| AlignedMemory::zero_filled(size))
134 }
135
136 pub fn put_stack(&mut self, stack: AlignedMemory<{ HOST_ALIGN }>) -> bool {
137 let len = stack.len();
138 self.stack.put(stack, len)
139 }
140
141 pub fn get_heap(&mut self, heap_size: u32) -> AlignedMemory<{ HOST_ALIGN }> {
142 debug_assert!((MIN_HEAP_FRAME_BYTES..=MAX_HEAP_FRAME_BYTES).contains(&heap_size));
143 self.heap
144 .get()
145 .unwrap_or_else(|| AlignedMemory::zero_filled(MAX_HEAP_FRAME_BYTES as usize))
146 }
147
148 pub fn put_heap(&mut self, heap: AlignedMemory<{ HOST_ALIGN }>, mapped_len: usize) -> bool {
149 let heap_size = heap.len();
150 debug_assert!(
151 heap_size >= MIN_HEAP_FRAME_BYTES as usize
152 && heap_size <= MAX_HEAP_FRAME_BYTES as usize
153 );
154 debug_assert!(mapped_len <= heap_size);
155 self.heap.put(heap, mapped_len.min(heap_size))
156 }
157
158 pub fn get_call_frames(&mut self) -> CallFrameBuffer {
159 self.call_frame.get().unwrap_or_default()
160 }
161
162 pub fn put_call_frames(&mut self, call_frame: CallFrameBuffer) -> bool {
163 self.call_frame.put(call_frame, 0)
164 }
165}
166
167#[cfg(test)]
168mod test {
169 use super::*;
170
171 #[derive(Debug, Eq, PartialEq)]
172 struct Item(u8, u8);
173 impl Reset for Item {
174 fn reset(&mut self, _len: usize) {
175 self.1 = 0;
176 }
177 }
178
179 #[test]
180 fn test_heap_shrink_then_grow_stays_zeroed() {
181 let mut pool = VmMemoryPool::new();
182 let big = MAX_HEAP_FRAME_BYTES;
183 let small = MIN_HEAP_FRAME_BYTES;
184
185 let mut heap = pool.get_heap(big);
186 heap.as_slice_mut().fill(0xaa);
187 assert!(pool.put_heap(heap, big as usize));
188
189 let mut heap = pool.get_heap(small);
190 heap.as_slice_mut()
191 .get_mut(..small as usize)
192 .unwrap()
193 .fill(0xbb);
194 assert!(pool.put_heap(heap, small as usize));
195
196 let heap = pool.get_heap(big);
197 assert!(heap.as_slice().iter().all(|byte| *byte == 0));
198 }
199
200 #[test]
201 fn test_pool() {
202 let mut pool = Pool::<Item, 2>::new([Item(0, 1), Item(1, 1)]);
203 assert_eq!(pool.get(), Some(Item(1, 1)));
204 assert_eq!(pool.get(), Some(Item(0, 1)));
205 assert_eq!(pool.get(), None);
206 pool.put(Item(1, 1), 0);
207 assert_eq!(pool.get(), Some(Item(1, 0)));
208 pool.put(Item(2, 2), 0);
209 pool.put(Item(3, 3), 0);
210 assert!(!pool.put(Item(4, 4), 0));
211 assert_eq!(pool.get(), Some(Item(3, 0)));
212 assert_eq!(pool.get(), Some(Item(2, 0)));
213 assert_eq!(pool.get(), None);
214 }
215}