edit/arena/
scratch.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4use std::ops::Deref;
5
6#[cfg(debug_assertions)]
7use super::debug;
8use super::{Arena, release};
9use crate::apperr;
10use crate::helpers::*;
11
12static mut S_SCRATCH: [release::Arena; 2] =
13    const { [release::Arena::empty(), release::Arena::empty()] };
14
15/// Initialize the scratch arenas with a given capacity.
16/// Call this before using [`scratch_arena`].
17pub fn init(capacity: usize) -> apperr::Result<()> {
18    unsafe {
19        for s in &mut S_SCRATCH[..] {
20            *s = release::Arena::new(capacity)?;
21        }
22    }
23    Ok(())
24}
25
26/// Need an arena for temporary allocations? [`scratch_arena`] got you covered.
27/// Call [`scratch_arena`] and it'll return an [`Arena`] that resets when it goes out of scope.
28///
29/// ---
30///
31/// Most methods make just two kinds of allocations:
32/// * Interior: Temporary data that can be deallocated when the function returns.
33/// * Exterior: Data that is returned to the caller and must remain alive until the caller stops using it.
34///
35/// Such methods only have two lifetimes, for which you consequently also only need two arenas.
36/// ...even if your method calls other methods recursively! This is because the exterior allocations
37/// of a callee are simply interior allocations to the caller, and so on, recursively.
38///
39/// This works as long as the two arenas flip/flop between being used as interior/exterior allocator
40/// along the callstack. To ensure that is the case, we use a recursion counter in debug builds.
41///
42/// This approach was described among others at: <https://nullprogram.com/blog/2023/09/27/>
43///
44/// # Safety
45///
46/// If your function takes an [`Arena`] argument, you **MUST** pass it to `scratch_arena` as `Some(&arena)`.
47pub fn scratch_arena(conflict: Option<&Arena>) -> ScratchArena<'static> {
48    unsafe {
49        #[cfg(debug_assertions)]
50        let conflict = conflict.map(|a| a.delegate_target_unchecked());
51
52        let index = opt_ptr_eq(conflict, Some(&S_SCRATCH[0])) as usize;
53        let arena = &S_SCRATCH[index];
54        ScratchArena::new(arena)
55    }
56}
57
58/// Borrows an [`Arena`] for temporary allocations.
59///
60/// See [`scratch_arena`].
61#[cfg(debug_assertions)]
62pub struct ScratchArena<'a> {
63    arena: debug::Arena,
64    offset: usize,
65    _phantom: std::marker::PhantomData<&'a ()>,
66}
67
68#[cfg(not(debug_assertions))]
69pub struct ScratchArena<'a> {
70    arena: &'a Arena,
71    offset: usize,
72}
73
74#[cfg(debug_assertions)]
75impl<'a> ScratchArena<'a> {
76    fn new(arena: &'a release::Arena) -> Self {
77        let offset = arena.offset();
78        ScratchArena { arena: Arena::delegated(arena), _phantom: std::marker::PhantomData, offset }
79    }
80}
81
82#[cfg(not(debug_assertions))]
83impl<'a> ScratchArena<'a> {
84    fn new(arena: &'a release::Arena) -> Self {
85        let offset = arena.offset();
86        ScratchArena { arena, offset }
87    }
88}
89
90impl Drop for ScratchArena<'_> {
91    fn drop(&mut self) {
92        unsafe { self.arena.reset(self.offset) };
93    }
94}
95
96#[cfg(debug_assertions)]
97impl Deref for ScratchArena<'_> {
98    type Target = debug::Arena;
99
100    fn deref(&self) -> &Self::Target {
101        &self.arena
102    }
103}
104
105#[cfg(not(debug_assertions))]
106impl Deref for ScratchArena<'_> {
107    type Target = Arena;
108
109    fn deref(&self) -> &Self::Target {
110        self.arena
111    }
112}