Skip to main content

obkrnl/malloc/
vm.rs

1use crate::config::{PAGE_MASK, PAGE_SHIFT, PAGE_SIZE};
2use crate::context::{CpuLocal, current_thread, uma};
3use crate::uma::{Alloc, SlabFlags, UmaFlags, UmaZone};
4use crate::vm::{PageObj, Vm, kaddr_to_phys};
5use alloc::string::ToString;
6use alloc::sync::Arc;
7use alloc::vec::Vec;
8use core::alloc::Layout;
9use core::cell::RefCell;
10use core::num::NonZero;
11use core::ptr::null_mut;
12
13/// Kernel heap that allocate a memory from a virtual memory management system. This struct is a
14/// merge of `malloc_type` and `malloc_type_internal` structure.
15pub struct VmHeap {
16    vm: &'static Vm,
17    zones: [Vec<Arc<UmaZone>>; PAGE_SHIFT + 1], // kmemsize + kmemzones
18    stats: CpuLocal<RefCell<Stats>>,            // mti_stats
19}
20
21impl VmHeap {
22    const KMEM_ZSHIFT: usize = 4;
23    const KMEM_ZBASE: usize = 16;
24    const KMEM_ZMASK: usize = Self::KMEM_ZBASE - 1;
25    const KMEM_ZSIZE: usize = PAGE_SIZE.get() >> Self::KMEM_ZSHIFT;
26
27    /// See `kmeminit` on the Orbis for a reference.
28    ///
29    /// # Reference offsets
30    /// | Version | Offset |
31    /// |---------|--------|
32    /// |PS4 11.00|0x1A4B80|
33    pub fn new(vm: &'static Vm) -> Self {
34        let uma = uma().unwrap();
35        let zones = core::array::from_fn(|align| {
36            let mut zones = Vec::with_capacity(Self::KMEM_ZSIZE + 1);
37            let mut last = 0;
38            let align = align
39                .try_into()
40                .ok()
41                .and_then(|align| 1usize.checked_shl(align))
42                .unwrap();
43
44            for i in Self::KMEM_ZSHIFT.. {
45                // Stop if size larger than page size.
46                let size = NonZero::new(1usize << i).unwrap();
47
48                if size > PAGE_SIZE {
49                    break;
50                }
51
52                // Create zone.
53                let zone = Arc::new(uma.into_owned().create_zone(
54                    size.to_string(),
55                    size,
56                    Some(align - 1),
57                    None,
58                    None,
59                    None,
60                    UmaFlags::Malloc,
61                ));
62
63                while last <= size.get() {
64                    zones.push(zone.clone());
65                    last += Self::KMEM_ZBASE;
66                }
67            }
68
69            zones
70        });
71
72        Self {
73            vm,
74            zones,
75            stats: CpuLocal::new(|_| RefCell::default()),
76        }
77    }
78
79    /// Returns null on failure.
80    ///
81    /// See `malloc` on the Orbis for a reference.
82    ///
83    /// # Safety
84    /// `layout` must be nonzero.
85    ///
86    /// # Reference offsets
87    /// | Version | Offset |
88    /// |---------|--------|
89    /// |PS4 11.00|0x1A4220|
90    pub unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
91        // Our implementation imply M_WAITOK.
92        let td = current_thread();
93
94        if !td.can_sleep() {
95            panic!("heap allocation in a non-sleeping context is not supported");
96        }
97
98        // Determine how to allocate.
99        let lock = td.disable_vm_heap();
100        let size = layout.size();
101        let mem = if size <= PAGE_SIZE.get() {
102            // Get zone to allocate from.
103            let zone = match self.zone_for_layout(layout) {
104                Some(v) => v,
105                None => return null_mut(),
106            };
107
108            // Allocate a memory from UMA zone.
109            let mem = zone.alloc(Alloc::Wait | Alloc::Zero);
110            let stats = self.stats.lock();
111            let mut stats = stats.borrow_mut();
112            let size = if mem.is_null() { 0 } else { zone.size().get() };
113
114            if size != 0 {
115                stats.alloc_bytes = stats.alloc_bytes.strict_add(size.try_into().unwrap());
116                stats.alloc_count += 1;
117            }
118
119            // TODO: How to update mts_size here since our zone table also indexed by alignment?
120            mem
121        } else {
122            todo!()
123        };
124
125        drop(lock);
126
127        mem
128    }
129
130    /// See `free` on the Orbis for a reference.
131    ///
132    /// # Safety
133    /// `ptr` must be obtained with [Self::alloc()] and `layout` must be the same one that was
134    /// passed to that method.
135    ///
136    /// # Reference offsets
137    /// | Version | Offset |
138    /// |---------|--------|
139    /// |PS4 11.00|0x1A43E0|
140    pub unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
141        let page = (ptr as usize) & !PAGE_MASK.get();
142        let page = unsafe { kaddr_to_phys(page) };
143        let page = self.vm.phys_to_page(page).unwrap(); // Orbis assume the pointer is not null.
144        let ps = page.state.lock();
145        let obj = ps.object.as_ref().unwrap(); // Orbis panic when this is null.
146        let PageObj::Slab(slab) = obj;
147        let size = if slab.flags().has_any(SlabFlags::Malloc) {
148            todo!()
149        } else {
150            // TODO: Should we drop the lock on page state before doing this?
151            let zone = self.zone_for_layout(layout).unwrap(); // Layout is the same as allocation.
152
153            unsafe { zone.free(ptr) };
154
155            slab.keg().size()
156        };
157
158        drop(ps);
159
160        // Update stats.
161        let stats = self.stats.lock();
162        let mut stats = stats.borrow_mut();
163
164        stats.freed_bytes = stats.freed_bytes.strict_add(size.get().try_into().unwrap());
165        stats.freed_count += 1;
166    }
167
168    /// Returns [None] if align is not supported.
169    ///
170    /// # Panics
171    /// If size greater than [PAGE_SIZE].
172    fn zone_for_layout(&self, layout: Layout) -> Option<&UmaZone> {
173        // Check if align supported.
174        let align = layout.align().trailing_zeros() as usize;
175        let zones = self.zones.get(align)?;
176
177        // Round size.
178        let size = layout.size();
179        let size = if (size & Self::KMEM_ZMASK) != 0 {
180            // TODO: Refactor this for readability.
181            (size + Self::KMEM_ZBASE) & !Self::KMEM_ZMASK
182        } else {
183            size
184        };
185
186        Some(&zones[size >> Self::KMEM_ZSHIFT])
187    }
188}
189
190/// Implementation of `malloc_type_stats` structure.
191#[derive(Default)]
192struct Stats {
193    alloc_bytes: u64, // mts_memalloced
194    alloc_count: u64, // mts_numallocs
195    freed_bytes: u64, // mts_memfreed
196    freed_count: u64, // mts_numfrees
197}