rkyv/util/alloc/
arena.rs

1use crate::ser::allocator::Arena;
2
3#[cfg(feature = "std")]
4mod detail {
5    use core::cell::Cell;
6
7    use crate::ser::allocator::Arena;
8
9    thread_local! {
10        static THREAD_ARENA: Cell<Option<Arena>> = const { Cell::new(None) };
11    }
12
13    pub fn with_arena<T>(f: impl FnOnce(&mut Arena) -> T) -> T {
14        THREAD_ARENA.with(|thread_arena| {
15            let mut arena = thread_arena.take().unwrap_or_default();
16
17            let result = f(&mut arena);
18            let capacity = arena.shrink();
19
20            if let Some(other) = thread_arena.take() {
21                if other.capacity() > capacity {
22                    arena = other;
23                }
24            }
25            thread_arena.set(Some(arena));
26
27            result
28        })
29    }
30
31    #[inline]
32    pub fn clear_arena() {
33        THREAD_ARENA.take();
34    }
35}
36
37#[cfg(all(not(feature = "std"), target_has_atomic = "ptr",))]
38mod detail {
39    use core::{
40        ptr::{self, NonNull},
41        sync::atomic::{AtomicPtr, Ordering},
42    };
43
44    use crate::ser::allocator::Arena;
45
46    static GLOBAL_ARENA: AtomicPtr<()> = AtomicPtr::new(ptr::null_mut());
47
48    pub fn with_arena<T>(f: impl FnOnce(&mut Arena) -> T) -> T {
49        let ptr = GLOBAL_ARENA.swap(ptr::null_mut(), Ordering::AcqRel);
50
51        let mut arena = if let Some(raw) = NonNull::new(ptr) {
52            unsafe { Arena::from_raw(raw) }
53        } else {
54            Arena::new()
55        };
56
57        let result = f(&mut arena);
58        arena.shrink();
59
60        let raw = arena.into_raw();
61
62        let swap = GLOBAL_ARENA.compare_exchange(
63            ptr::null_mut(),
64            raw.as_ptr(),
65            Ordering::AcqRel,
66            Ordering::Relaxed,
67        );
68        if swap.is_err() {
69            // Another arena was swapped in while we were executing `f`. We need
70            // to free the current arena.
71            unsafe {
72                drop(Arena::from_raw(raw));
73            }
74        }
75
76        result
77    }
78
79    #[inline]
80    pub fn clear_arena() {
81        let ptr = GLOBAL_ARENA.swap(ptr::null_mut(), Ordering::AcqRel);
82
83        if let Some(raw) = NonNull::new(ptr) {
84            unsafe {
85                drop(Arena::from_raw(raw));
86            }
87        }
88    }
89}
90
91#[cfg(all(not(feature = "std"), not(target_has_atomic = "ptr"),))]
92mod detail {
93    use crate::ser::allocator::Arena;
94
95    pub fn with_arena<T>(f: impl FnOnce(&mut Arena) -> T) -> T {
96        let mut arena = Arena::new();
97        f(&mut arena)
98    }
99
100    #[inline]
101    pub fn clear_arena() {}
102}
103
104/// Calls the given function with the builtin arena allocator.
105///
106/// When the `std` feature is enabled, the builtin arena allocator is a
107/// thread-local variable, with one allocator per thread. When atomic pointers
108/// are supported, it is a global static and all threads share the same arena.
109/// Otherwise, this will create and drop a new arena each time it is called.
110pub fn with_arena<T>(f: impl FnOnce(&mut Arena) -> T) -> T {
111    detail::with_arena(f)
112}
113
114/// Clears the builtin arena allocator.
115///
116/// When the `std` feature is enabled, this only clears the allocator for the
117/// current thread. When atomic pointers are supported, this will clear the
118/// allocator for all threads. Otherwise, this function does nothing.
119#[inline]
120pub fn clear_arena() {
121    detail::clear_arena()
122}
rkyv/util/alloc/arena.rs

rkyv/util/alloc/
arena.rs
