Skip to main content

memlink_runtime/
arena.rs

1//! Arena allocator for module memory management.
2
3use std::sync::atomic::{AtomicUsize, Ordering};
4
5pub const DEFAULT_ARENA_CAPACITY: usize = 64 * 1024 * 1024;
6
7pub struct Arena {
8    base: Vec<u8>,
9    capacity: usize,
10    used: AtomicUsize,
11}
12
13unsafe impl Send for Arena {}
14unsafe impl Sync for Arena {}
15
16impl Arena {
17    pub fn new(capacity: usize) -> Self {
18        Arena {
19            base: vec![0u8; capacity],
20            capacity,
21            used: AtomicUsize::new(0),
22        }
23    }
24
25    pub fn with_default_capacity() -> Self {
26        Self::new(DEFAULT_ARENA_CAPACITY)
27    }
28
29    pub fn alloc(&self, size: usize) -> Option<*mut u8> {
30        let mut current = self.used.load(Ordering::Relaxed);
31
32        loop {
33            let new_used = current + size;
34
35            if new_used > self.capacity {
36                return None;
37            }
38
39            match self.used.compare_exchange(
40                current,
41                new_used,
42                Ordering::SeqCst,
43                Ordering::Relaxed,
44            ) {
45                Ok(_) => {
46                    let ptr = self.base.as_ptr() as usize + current;
47                    return Some(ptr as *mut u8);
48                }
49                Err(actual) => {
50                    current = actual;
51                }
52            }
53        }
54    }
55
56    pub fn alloc_aligned(&self, size: usize, align: usize) -> Option<*mut u8> {
57        debug_assert!(align.is_power_of_two(), "Alignment must be a power of 2");
58
59        let mut current = self.used.load(Ordering::Relaxed);
60
61        loop {
62            let aligned = (current + align - 1) & !(align - 1);
63            let new_used = aligned + size;
64
65            if new_used > self.capacity {
66                return None;
67            }
68
69            match self.used.compare_exchange(
70                current,
71                new_used,
72                Ordering::SeqCst,
73                Ordering::Relaxed,
74            ) {
75                Ok(_) => {
76                    let ptr = self.base.as_ptr() as usize + aligned;
77                    return Some(ptr as *mut u8);
78                }
79                Err(actual) => {
80                    current = actual;
81                }
82            }
83        }
84    }
85
86    pub fn reset(&self) {
87        self.used.store(0, Ordering::Release);
88    }
89
90    pub fn capacity(&self) -> usize {
91        self.capacity
92    }
93
94    pub fn used(&self) -> usize {
95        self.used.load(Ordering::Acquire)
96    }
97
98    pub fn remaining(&self) -> usize {
99        self.capacity - self.used.load(Ordering::Relaxed)
100    }
101
102    pub fn usage(&self) -> f32 {
103        self.used() as f32 / self.capacity as f32
104    }
105
106    pub fn as_slice(&self) -> &[u8] {
107        &self.base
108    }
109
110    pub fn as_mut_slice(&mut self) -> &mut [u8] {
111        &mut self.base
112    }
113}
114
115impl Default for Arena {
116    fn default() -> Self {
117        Self::with_default_capacity()
118    }
119}
120
121impl Clone for Arena {
122    fn clone(&self) -> Self {
123        Arena {
124            base: self.base.clone(),
125            capacity: self.capacity,
126            used: AtomicUsize::new(self.used.load(Ordering::Relaxed)),
127        }
128    }
129}
130
131impl std::fmt::Debug for Arena {
132    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
133        f.debug_struct("Arena")
134            .field("capacity", &self.capacity)
135            .field("used", &self.used())
136            .field("remaining", &self.remaining())
137            .field("usage", &self.usage())
138            .finish()
139    }
140}