rustpython_vm/
datastack.rs

1/// Thread data stack for interpreter frames (`_PyStackChunk` /
2/// `tstate->datastack_*`).
3///
4/// A linked list of chunks providing bump allocation for frame-local data
5/// (localsplus arrays).  Normal function calls allocate via `push()`
6/// (pointer bump).  Generators and coroutines use heap-allocated storage.
7use alloc::alloc::{alloc, dealloc};
8use core::alloc::Layout;
9use core::ptr;
10
11/// Minimum chunk size in bytes (`_PY_DATA_STACK_CHUNK_SIZE`).
12/// Smaller on WASM (4 KB) to reduce initial memory footprint; 16 KB otherwise.
13const MIN_CHUNK_SIZE: usize = if cfg!(target_arch = "wasm32") {
14    4 * 1024
15} else {
16    16 * 1024
17};
18
19/// Extra headroom (in bytes) to avoid allocating a new chunk for the next
20/// frame right after growing.
21const MINIMUM_OVERHEAD: usize = 1000 * core::mem::size_of::<usize>();
22
23/// Alignment for all data stack allocations.
24const ALIGN: usize = 16;
25
26/// Header for a data stack chunk.
27///
28/// The usable data region starts right after this header (aligned to [`ALIGN`]).
29#[repr(C)]
30struct DataStackChunk {
31    /// Previous chunk in the linked list (NULL for the root chunk).
32    previous: *mut Self,
33    /// Total allocation size in bytes (including this header).
34    size: usize,
35    /// Saved `top` offset when a newer chunk was pushed.  Used to restore
36    /// `DataStack::top` when popping back to this chunk.
37    saved_top: usize,
38}
39
40impl DataStackChunk {
41    /// Pointer to the first usable byte after the header (aligned).
42    #[inline(always)]
43    fn data_start(&self) -> *mut u8 {
44        let header_end = (self as *const Self as usize) + core::mem::size_of::<Self>();
45        let aligned = (header_end + ALIGN - 1) & !(ALIGN - 1);
46        aligned as *mut u8
47    }
48
49    /// Pointer past the last usable byte.
50    #[inline(always)]
51    fn data_limit(&self) -> *mut u8 {
52        unsafe { (self as *const Self as *mut u8).add(self.size) }
53    }
54}
55
56/// Per-thread data stack for bump-allocating frame-local data.
57pub struct DataStack {
58    /// Current chunk.
59    chunk: *mut DataStackChunk,
60    /// Current allocation position within the current chunk.
61    top: *mut u8,
62    /// End of usable space in the current chunk.
63    limit: *mut u8,
64    /// Most recently popped full-frame allocation whose localsplus slots were
65    /// cleared before the pop. An exact LIFO reuse can skip zero-filling them.
66    reusable_frame: Option<(*mut u8, usize)>,
67}
68
69impl DataStack {
70    /// Create a new data stack with an initial root chunk.
71    #[must_use]
72    pub fn new() -> Self {
73        let chunk = Self::alloc_chunk(MIN_CHUNK_SIZE, ptr::null_mut());
74        let top = unsafe { (*chunk).data_start() };
75        let limit = unsafe { (*chunk).data_limit() };
76        // Skip one ALIGN-sized slot in the root chunk so that `pop()` never
77        // frees it (`push_chunk` convention).
78        let top = unsafe { top.add(ALIGN) };
79        Self {
80            chunk,
81            top,
82            limit,
83            reusable_frame: None,
84        }
85    }
86
87    /// Check if the current chunk has at least `size` bytes available.
88    #[inline(always)]
89    #[must_use]
90    pub fn has_space(&self, size: usize) -> bool {
91        let aligned_size = (size + ALIGN - 1) & !(ALIGN - 1);
92        (self.limit as usize).saturating_sub(self.top as usize) >= aligned_size
93    }
94
95    /// Allocate `size` bytes from the data stack.
96    ///
97    /// Returns a pointer to the allocated region (aligned to `ALIGN`).
98    /// The caller must call `pop()` with the returned pointer when done
99    /// (LIFO order).
100    #[inline(always)]
101    pub fn push(&mut self, size: usize) -> *mut u8 {
102        self.reusable_frame = None;
103        self.push_inner(size)
104    }
105
106    /// Allocate a full interpreter frame and report whether it exactly reuses
107    /// a just-cleared frame block.
108    #[inline(always)]
109    pub fn push_frame(&mut self, size: usize) -> (*mut u8, bool) {
110        let reusable_frame = self.reusable_frame.take();
111        let ptr = self.push_inner(size);
112        // Exact sizes, not aligned ones: the caller reads "reused" as "every
113        // slot of this frame was cleared by the last one", and two frames whose
114        // sizes differ by less than ALIGN share an aligned size while the
115        // larger one's tail slots were never touched, let alone cleared.
116        let reused = reusable_frame.is_some_and(|(base, old_size)| base == ptr && old_size == size);
117        (ptr, reused)
118    }
119
120    #[inline(always)]
121    fn push_inner(&mut self, size: usize) -> *mut u8 {
122        let aligned_size = (size + ALIGN - 1) & !(ALIGN - 1);
123        unsafe {
124            if self.top.add(aligned_size) <= self.limit {
125                let ptr = self.top;
126                self.top = self.top.add(aligned_size);
127                ptr
128            } else {
129                self.push_slow(aligned_size)
130            }
131        }
132    }
133
134    /// Slow path: allocate a new chunk and push from it.
135    #[cold]
136    #[inline(never)]
137    fn push_slow(&mut self, aligned_size: usize) -> *mut u8 {
138        let mut chunk_size = MIN_CHUNK_SIZE;
139        let needed = aligned_size
140            .checked_add(MINIMUM_OVERHEAD)
141            .and_then(|v| v.checked_add(core::mem::size_of::<DataStackChunk>()))
142            .and_then(|v| v.checked_add(ALIGN))
143            .expect("DataStack chunk size overflow");
144        while chunk_size < needed {
145            chunk_size = chunk_size
146                .checked_mul(2)
147                .expect("DataStack chunk size overflow");
148        }
149        // Save current position in old chunk.
150        unsafe {
151            (*self.chunk).saved_top = self.top as usize - self.chunk as usize;
152        }
153        let new_chunk = Self::alloc_chunk(chunk_size, self.chunk);
154        self.chunk = new_chunk;
155        let start = unsafe { (*new_chunk).data_start() };
156        self.limit = unsafe { (*new_chunk).data_limit() };
157        self.top = unsafe { start.add(aligned_size) };
158        start
159    }
160
161    /// Pop a previous allocation.  `base` must be the pointer returned by
162    /// `push()`.  Calls must be in LIFO order.
163    ///
164    /// # Safety
165    /// `base` must be a valid pointer returned by `push()` on this data stack,
166    /// and all allocations made after it must already have been popped.
167    #[inline(always)]
168    pub unsafe fn pop(&mut self, base: *mut u8) {
169        self.reusable_frame = None;
170        unsafe { self.pop_inner(base) };
171    }
172
173    /// Pop a full frame whose localsplus slots have already been cleared.
174    ///
175    /// # Safety
176    /// `base` and `size` must describe the most recent allocation returned by
177    /// `push_frame`, every later allocation must already be popped, and all
178    /// localsplus slots in the frame must have been cleared.
179    #[inline(always)]
180    pub unsafe fn pop_frame(&mut self, base: *mut u8, size: usize) {
181        unsafe { self.pop_inner(base) };
182        self.reusable_frame = Some((base, size));
183    }
184
185    #[inline(always)]
186    unsafe fn pop_inner(&mut self, base: *mut u8) {
187        debug_assert!(!base.is_null());
188        if self.is_in_current_chunk(base) {
189            // Common case: base is within the current chunk.
190            self.top = base;
191        } else {
192            // base is in a previous chunk — free the current chunk.
193            unsafe { self.pop_slow(base) };
194        }
195    }
196
197    /// Check if `ptr` falls within the current chunk's data area.
198    /// Both bounds are checked to handle non-monotonic allocation addresses
199    /// (e.g. on Windows where newer chunks may be at lower addresses).
200    #[inline(always)]
201    fn is_in_current_chunk(&self, ptr: *mut u8) -> bool {
202        let chunk_start = unsafe { (*self.chunk).data_start() };
203        ptr >= chunk_start && ptr <= self.limit
204    }
205
206    /// Slow path: pop back to a previous chunk.
207    #[cold]
208    #[inline(never)]
209    unsafe fn pop_slow(&mut self, base: *mut u8) {
210        loop {
211            let old_chunk = self.chunk;
212            let prev = unsafe { (*old_chunk).previous };
213            debug_assert!(!prev.is_null(), "tried to pop past the root chunk");
214            unsafe { Self::free_chunk(old_chunk) };
215            self.chunk = prev;
216            self.limit = unsafe { (*prev).data_limit() };
217            if self.is_in_current_chunk(base) {
218                self.top = base;
219                return;
220            }
221        }
222    }
223
224    /// Allocate a new chunk.
225    fn alloc_chunk(size: usize, previous: *mut DataStackChunk) -> *mut DataStackChunk {
226        let layout = Layout::from_size_align(size, ALIGN).expect("invalid chunk layout");
227        let ptr = unsafe { alloc(layout) };
228        if ptr.is_null() {
229            alloc::alloc::handle_alloc_error(layout);
230        }
231        let chunk = ptr as *mut DataStackChunk;
232        unsafe {
233            (*chunk).previous = previous;
234            (*chunk).size = size;
235            (*chunk).saved_top = 0;
236        }
237        chunk
238    }
239
240    /// Free a chunk.
241    unsafe fn free_chunk(chunk: *mut DataStackChunk) {
242        let size = unsafe { (*chunk).size };
243        let layout = Layout::from_size_align(size, ALIGN).expect("invalid chunk layout");
244        unsafe { dealloc(chunk as *mut u8, layout) };
245    }
246}
247
248// SAFETY: DataStack is per-thread and not shared.  The raw pointers
249// it contains point to memory exclusively owned by this DataStack.
250unsafe impl Send for DataStack {}
251
252impl Default for DataStack {
253    fn default() -> Self {
254        Self::new()
255    }
256}
257
258impl Drop for DataStack {
259    fn drop(&mut self) {
260        let mut chunk = self.chunk;
261        while !chunk.is_null() {
262            let prev = unsafe { (*chunk).previous };
263            unsafe { Self::free_chunk(chunk) };
264            chunk = prev;
265        }
266    }
267}
268
269#[cfg(test)]
270mod tests {
271    use super::*;
272
273    #[test]
274    fn basic_push_pop() {
275        let mut ds = DataStack::new();
276        let p1 = ds.push(64);
277        assert!(!p1.is_null());
278        let p2 = ds.push(128);
279        assert!(!p2.is_null());
280        assert!(p2 > p1);
281        unsafe {
282            ds.pop(p2);
283            ds.pop(p1);
284        }
285    }
286
287    #[test]
288    fn cross_chunk_push_pop() {
289        let mut ds = DataStack::new();
290        // Push enough to force a new chunk
291        let mut ptrs = Vec::new();
292        for _ in 0..100 {
293            ptrs.push(ds.push(1024));
294        }
295        // Pop all in reverse
296        for p in ptrs.into_iter().rev() {
297            unsafe { ds.pop(p) };
298        }
299    }
300
301    #[test]
302    fn alignment() {
303        let mut ds = DataStack::new();
304        for size in [1, 7, 15, 16, 17, 31, 32, 33, 64, 100] {
305            let p = ds.push(size);
306            assert_eq!(p as usize % ALIGN, 0, "alignment violated for size {size}");
307            unsafe { ds.pop(p) };
308        }
309    }
310}
rustpython_vm/datastack.rs

rustpython_vm/
datastack.rs
