solana_program_runtime/
mem_pool.rs

1use {
2    solana_compute_budget::{
3        compute_budget::{MAX_CALL_DEPTH, MAX_INSTRUCTION_STACK_DEPTH, STACK_FRAME_SIZE},
4        compute_budget_limits::{MAX_HEAP_FRAME_BYTES, MIN_HEAP_FRAME_BYTES},
5    },
6    solana_sbpf::{aligned_memory::AlignedMemory, ebpf::HOST_ALIGN},
7    std::array,
8};
9
10trait Reset {
11    fn reset(&mut self);
12}
13
14struct Pool<T: Reset, const SIZE: usize> {
15    items: [Option<T>; SIZE],
16    next_empty: usize,
17}
18
19impl<T: Reset, const SIZE: usize> Pool<T, SIZE> {
20    fn new(items: [T; SIZE]) -> Self {
21        Self {
22            items: items.map(|i| Some(i)),
23            next_empty: SIZE,
24        }
25    }
26
27    fn len(&self) -> usize {
28        SIZE
29    }
30
31    fn get(&mut self) -> Option<T> {
32        if self.next_empty == 0 {
33            return None;
34        }
35        self.next_empty = self.next_empty.saturating_sub(1);
36        self.items
37            .get_mut(self.next_empty)
38            .and_then(|item| item.take())
39    }
40
41    fn put(&mut self, mut value: T) -> bool {
42        self.items
43            .get_mut(self.next_empty)
44            .map(|item| {
45                value.reset();
46                item.replace(value);
47                self.next_empty = self.next_empty.saturating_add(1);
48                true
49            })
50            .unwrap_or(false)
51    }
52}
53
54impl Reset for AlignedMemory<{ HOST_ALIGN }> {
55    fn reset(&mut self) {
56        self.as_slice_mut().fill(0)
57    }
58}
59
60pub struct VmMemoryPool {
61    stack: Pool<AlignedMemory<{ HOST_ALIGN }>, MAX_INSTRUCTION_STACK_DEPTH>,
62    heap: Pool<AlignedMemory<{ HOST_ALIGN }>, MAX_INSTRUCTION_STACK_DEPTH>,
63}
64
65impl VmMemoryPool {
66    pub fn new() -> Self {
67        Self {
68            stack: Pool::new(array::from_fn(|_| {
69                AlignedMemory::zero_filled(STACK_FRAME_SIZE * MAX_CALL_DEPTH)
70            })),
71            heap: Pool::new(array::from_fn(|_| {
72                AlignedMemory::zero_filled(MAX_HEAP_FRAME_BYTES as usize)
73            })),
74        }
75    }
76
77    pub fn stack_len(&self) -> usize {
78        self.stack.len()
79    }
80
81    pub fn heap_len(&self) -> usize {
82        self.heap.len()
83    }
84
85    pub fn get_stack(&mut self, size: usize) -> AlignedMemory<{ HOST_ALIGN }> {
86        debug_assert!(size == STACK_FRAME_SIZE * MAX_CALL_DEPTH);
87        self.stack
88            .get()
89            .unwrap_or_else(|| AlignedMemory::zero_filled(size))
90    }
91
92    pub fn put_stack(&mut self, stack: AlignedMemory<{ HOST_ALIGN }>) -> bool {
93        self.stack.put(stack)
94    }
95
96    pub fn get_heap(&mut self, heap_size: u32) -> AlignedMemory<{ HOST_ALIGN }> {
97        debug_assert!((MIN_HEAP_FRAME_BYTES..=MAX_HEAP_FRAME_BYTES).contains(&heap_size));
98        self.heap
99            .get()
100            .unwrap_or_else(|| AlignedMemory::zero_filled(MAX_HEAP_FRAME_BYTES as usize))
101    }
102
103    pub fn put_heap(&mut self, heap: AlignedMemory<{ HOST_ALIGN }>) -> bool {
104        let heap_size = heap.len();
105        debug_assert!(
106            heap_size >= MIN_HEAP_FRAME_BYTES as usize
107                && heap_size <= MAX_HEAP_FRAME_BYTES as usize
108        );
109        self.heap.put(heap)
110    }
111}
112
113impl Default for VmMemoryPool {
114    fn default() -> Self {
115        Self::new()
116    }
117}
118
119#[cfg(test)]
120mod test {
121    use super::*;
122
123    #[derive(Debug, Eq, PartialEq)]
124    struct Item(u8, u8);
125    impl Reset for Item {
126        fn reset(&mut self) {
127            self.1 = 0;
128        }
129    }
130
131    #[test]
132    fn test_pool() {
133        let mut pool = Pool::<Item, 2>::new([Item(0, 1), Item(1, 1)]);
134        assert_eq!(pool.get(), Some(Item(1, 1)));
135        assert_eq!(pool.get(), Some(Item(0, 1)));
136        assert_eq!(pool.get(), None);
137        pool.put(Item(1, 1));
138        assert_eq!(pool.get(), Some(Item(1, 0)));
139        pool.put(Item(2, 2));
140        pool.put(Item(3, 3));
141        assert!(!pool.put(Item(4, 4)));
142        assert_eq!(pool.get(), Some(Item(3, 0)));
143        assert_eq!(pool.get(), Some(Item(2, 0)));
144        assert_eq!(pool.get(), None);
145    }
146}
solana_program_runtime/mem_pool.rs

solana_program_runtime/
mem_pool.rs
