Skip to main content

scope_local/
scope.rs

1use alloc::alloc::{alloc, dealloc, handle_alloc_error};
2use core::{alloc::Layout, iter::zip, mem::MaybeUninit, ptr::NonNull};
3
4use spin::{LazyLock, Once};
5
6use crate::{
7    boxed::ItemBox,
8    item::{Item, Registry},
9};
10
11/// A scope is a collection of items.
12pub struct Scope {
13    ptr: NonNull<ItemSlot>,
14}
15
16unsafe impl Send for Scope {}
17unsafe impl Sync for Scope {}
18
19impl Scope {
20    fn len() -> usize {
21        Registry.len()
22    }
23
24    fn layout() -> Layout {
25        Layout::array::<ItemSlot>(Self::len()).unwrap()
26    }
27
28    /// Create a new namespace with all resources initialized as their default
29    /// value.
30    pub fn new() -> Self {
31        let layout = Self::layout();
32        let ptr = NonNull::new(unsafe { alloc(layout) })
33            .unwrap_or_else(|| handle_alloc_error(layout))
34            .cast();
35
36        let slice = unsafe {
37            core::slice::from_raw_parts_mut(ptr.cast::<MaybeUninit<_>>().as_ptr(), Registry.len())
38        };
39        for (item, d) in zip(&*Registry, slice) {
40            d.write(ItemSlot::new(item));
41        }
42
43        Self { ptr }
44    }
45
46    pub(crate) fn get(&self, item: &'static Item) -> &ItemBox {
47        let index = item.index();
48        unsafe { self.ptr.add(index).as_ref() }.get()
49    }
50
51    pub(crate) fn get_mut(&mut self, item: &'static Item) -> &mut ItemBox {
52        let index = item.index();
53        unsafe { self.ptr.add(index).as_mut() }.get_mut()
54    }
55}
56
57impl Default for Scope {
58    fn default() -> Self {
59        Self::new()
60    }
61}
62
63impl Drop for Scope {
64    fn drop(&mut self) {
65        let ptr = NonNull::slice_from_raw_parts(self.ptr, Self::len());
66        unsafe {
67            ptr.drop_in_place();
68            dealloc(self.ptr.cast().as_ptr(), Self::layout());
69        }
70    }
71}
72
73struct ItemSlot {
74    item: &'static Item,
75    value: Once<ItemBox>,
76}
77
78impl ItemSlot {
79    fn new(item: &'static Item) -> Self {
80        Self {
81            item,
82            value: Once::new(),
83        }
84    }
85
86    fn get(&self) -> &ItemBox {
87        self.value.call_once(|| ItemBox::new(self.item))
88    }
89
90    fn get_mut(&mut self) -> &mut ItemBox {
91        if !self.value.is_completed() {
92            let item = self.item;
93            self.value.call_once(|| ItemBox::new(item));
94        }
95        self.value
96            .get_mut()
97            .expect("scope-local item must be initialized")
98    }
99}
100
101static GLOBAL_SCOPE: LazyLock<Scope> = LazyLock::new(Scope::new);
102
103#[ax_percpu::def_percpu]
104pub(crate) static ACTIVE_SCOPE_PTR: usize = 0;
105
106/// Currently active scope.
107pub struct ActiveScope;
108
109impl ActiveScope {
110    /// Sets the active scope pointer to the given scope.
111    ///
112    /// # Safety
113    ///
114    /// The caller must ensure that the provided `scope` reference is valid for
115    /// the duration in which it is set as the active scope, and that no data
116    /// races or aliasing violations occur.
117    pub unsafe fn set(scope: &Scope) {
118        ACTIVE_SCOPE_PTR.write_current(scope.ptr.addr().into());
119    }
120
121    /// Set the active scope to the global scope.
122    pub fn set_global() {
123        ACTIVE_SCOPE_PTR.write_current(0);
124    }
125
126    /// Returns true if the active scope is the global scope.
127    pub fn is_global() -> bool {
128        ACTIVE_SCOPE_PTR.read_current() == 0
129    }
130
131    pub(crate) fn get<'a>(item: &'static Item) -> &'a ItemBox {
132        let ptr = ACTIVE_SCOPE_PTR.read_current();
133        let ptr = NonNull::new(ptr as _).unwrap_or(GLOBAL_SCOPE.ptr);
134        let index = item.index();
135        unsafe { ptr.add(index).as_ref() }.get()
136    }
137}