Skip to main content

obkrnl/uma/
slab.rs

1use super::UmaKeg;
2use crate::lock::Mutex;
3use crate::mem::{RefCnt, too_many_refs};
4use alloc::sync::Arc;
5use core::marker::PhantomPinned;
6use core::ptr::null_mut;
7use core::sync::atomic::{AtomicUsize, Ordering};
8use macros::bitflag;
9
10/// Implementation of `uma_slab`.
11///
12/// Unlike Orbis, we don't support `uma_slab_refcnt`. We use [alloc::sync::Arc] for that job instead
13/// so you need to wrap the item to allocate from the slab with [alloc::sync::Arc] for any zones
14/// that going to create with `UMA_ZONE_REFCNT`.
15///
16/// We use slightly different mechanism here but has the same memory layout (except we don't support
17/// `uma_slab_refcnt` as stated on the above).
18///
19/// # Safety
20/// Adding more fields into this struct without knowing how it work can cause undefined behavior in
21/// some places.
22#[repr(C)]
23pub struct Slab {
24    pub(super) pin: PhantomPinned,
25    pub(super) hdr: SlabHdr, // us_head
26    pub(super) free: [u8],   // us_freelist
27}
28
29impl Slab {
30    pub fn keg(&self) -> &UmaKeg {
31        &self.hdr.keg
32    }
33
34    pub fn flags(&self) -> SlabFlags {
35        self.hdr.flags
36    }
37
38    /// Each allocated item keep a strong reference to the slab, which mean all allocated item need
39    /// to free manually otherwise the slab (and its keg) will be leak.
40    ///
41    /// Unlike Orbis, this method will return null if the slab already full instead of trigger a UB.
42    ///
43    /// See `slab_alloc_item` on the Orbis for a reference.
44    ///
45    /// # Reference offsets
46    /// | Version | Offset |
47    /// |---------|--------|
48    /// |PS4 11.00|0x141FE0|
49    pub fn alloc_item(&self) -> *mut u8 {
50        // Check if full.
51        let mut k = self.hdr.keg.state().lock();
52        let mut s = self.hdr.state.lock();
53
54        if s.free_count == 0 {
55            return null_mut();
56        }
57
58        // Allocate.
59        let f = usize::from(s.first_free);
60
61        s.first_free = self.free[f];
62        s.free_count -= 1;
63        k.free -= 1;
64
65        if s.free_count == 0 {
66            todo!()
67        }
68
69        self.hdr.refs.fetch_add(1, Ordering::Relaxed);
70
71        unsafe { self.hdr.items.add(f * self.hdr.keg.allocated_size()) }
72    }
73}
74
75impl Drop for Slab {
76    #[inline(never)]
77    fn drop(&mut self) {
78        core::sync::atomic::fence(Ordering::Acquire);
79
80        todo!()
81    }
82}
83
84unsafe impl RefCnt for Slab {
85    fn increase_ref(&self) {
86        let p = self.hdr.refs.fetch_add(1, Ordering::Relaxed);
87
88        if p == usize::MAX {
89            too_many_refs();
90        }
91    }
92
93    fn decrease_ref(&self) -> usize {
94        self.hdr.refs.fetch_sub(1, Ordering::Release)
95    }
96}
97
98unsafe impl Send for Slab {}
99unsafe impl Sync for Slab {}
100
101/// Implementation of `uma_slab_head`.
102pub(super) struct SlabHdr {
103    keg: Arc<UmaKeg>, // us_keg
104    items: *mut u8,   // us_data
105    flags: SlabFlags, // us_flags
106    /// This **MUST** be locked after everything else (e.g. keg state and zone state) otherwise it
107    /// will cause a deadlock.
108    state: Mutex<SlabState>,
109    refs: AtomicUsize,
110}
111
112impl SlabHdr {
113    /// # Safety
114    /// - `items` cannot be null.
115    /// - `len` must be a number of elements of the array at `items`.
116    pub unsafe fn new(keg: Arc<UmaKeg>, flags: SlabFlags, items: *mut u8, len: usize) -> Self {
117        Self {
118            keg,
119            items,
120            flags,
121            state: Mutex::new(SlabState {
122                free_count: len,
123                first_free: 0,
124            }),
125            refs: AtomicUsize::new(0),
126        }
127    }
128}
129
130/// Flags for [Slab].
131#[bitflag(u8)]
132pub enum SlabFlags {
133    /// `UMA_SLAB_PRIV`.
134    Private = 0x08,
135    /// `UMA_SLAB_MALLOC`.
136    Malloc = 0x20,
137}
138
139/// Contains mutable data for [SlabHdr].
140struct SlabState {
141    free_count: usize, // us_freecount
142    first_free: u8,    // us_firstfree
143}