Implements zone_fetch_slab (#1274)

kernel/src/malloc/mod.rs

@@ -1,4 +1,4 @@
-use self::stage2::VmHeap;
+use self::vm::VmHeap;
 use crate::context::current_thread;
 use crate::lock::Mutex;
 use alloc::boxed::Box;
@@ -8,7 +8,7 @@ use core::hint::unreachable_unchecked;
 use core::ptr::{null_mut, NonNull};
 use talc::{ClaimOnOom, Span, Talc};
 
-mod stage2;
+mod vm;
 
 /// Implementation of [`GlobalAlloc`] for objects belong to kernel space.
 ///
@@ -42,19 +42,19 @@ impl KernelHeap {
     /// This must be called by main CPU and can be called only once.
     pub unsafe fn activate_stage2(&self) {
         // Setup VM  heap using primitive heap.
-        let stage2 = Box::new(VmHeap::new());
+        let vm = Box::new(VmHeap::new());
 
         // What we are doing here is highly unsafe. Do not edit the code after this unless you know
         // what you are doing!
         let stage = self.stage.get();
-        let stage1 = match stage.read() {
+        let primitive = match stage.read() {
             Stage::One(v) => Mutex::new(v.into_inner()),
             Stage::Two(_, _) => unreachable_unchecked(),
         };
 
         // Switch to stage 2 WITHOUT dropping the value contained in Stage::One. What we did here is
         // moving the value from Stage::One to Stage::Two.
-        stage.write(Stage::Two(stage2, stage1));
+        stage.write(Stage::Two(vm, primitive));
     }
 }
 

kernel/src/uma/keg.rs

@@ -1,8 +1,6 @@
-use super::slab::RcFree;
-use super::UmaFlags;
+use super::slab::{Free, RcFree, Slab};
+use super::{Alloc, Uma, UmaFlags, UmaZone};
 use crate::config::{PAGE_MASK, PAGE_SHIFT, PAGE_SIZE};
-use crate::uma::slab::{Free, SlabHdr};
-use crate::uma::Uma;
 use core::alloc::Layout;
 use core::cmp::{max, min};
 use core::num::NonZero;
@@ -11,6 +9,7 @@ use core::num::NonZero;
 pub struct UmaKeg {
     size: NonZero<usize>, // uk_size
     ipers: usize,         // uk_ipers
+    recurse: u32,         // uk_recurse
     flags: UmaFlags,      // uk_flags
 }
 
@@ -38,7 +37,7 @@ impl UmaKeg {
         }
 
         // Get header layout.
-        let hdr = Layout::new::<SlabHdr>();
+        let hdr = Layout::new::<Slab<()>>();
         let (mut hdr, off) = if flags.has(UmaFlags::RefCnt) {
             hdr.extend(Layout::new::<RcFree>()).unwrap()
         } else {
@@ -154,18 +153,37 @@ impl UmaKeg {
 
         // TODO: Add uk_zones.
         // TODO: Add uma_kegs.
-        Self { size, ipers, flags }
+        Self {
+            size,
+            ipers,
+            recurse: 0,
+            flags,
+        }
     }
 
     pub fn size(&self) -> NonZero<usize> {
         self.size
     }
 
+    pub fn item_per_slab(&self) -> usize {
+        self.ipers
+    }
+
+    pub fn recurse(&self) -> u32 {
+        self.recurse
+    }
+
     pub fn flags(&self) -> UmaFlags {
         self.flags
     }
 
-    pub fn item_per_slab(&self) -> usize {
-        self.ipers
+    /// See `keg_fetch_slab` on the Orbis for a reference.
+    ///
+    /// # Reference offsets
+    /// | Version | Offset |
+    /// |---------|--------|
+    /// |PS4 11.00|0x141E20|
+    pub fn fetch_slab(&self, _: &UmaZone, _: Alloc) -> Option<()> {
+        todo!()
     }
 }

kernel/src/uma/mod.rs

@@ -146,6 +146,8 @@ pub enum UmaFlags {
 /// Implementation of `malloc` flags.
 #[bitflag(u32)]
 pub enum Alloc {
+    /// `M_NOWAIT`.
+    NoWait = 0x1,
     /// `M_WAITOK`.
     Wait = 0x2,
     /// `M_ZERO`.

kernel/src/uma/slab.rs

@@ -2,7 +2,9 @@
 ///
 /// We use slightly different mechanism here but has the same memory layout.
 #[repr(C)]
-pub struct SlabHdr {}
+pub struct Slab<I: ?Sized> {
+    free: I, // us_freelist
+}
 
 /// Item in the slab to represents `uma_slab` structure.
 #[repr(C)]

kernel/src/uma/zone.rs

@@ -3,6 +3,7 @@ use super::keg::UmaKeg;
 use super::{Alloc, Uma, UmaBox, UmaFlags};
 use crate::context::{current_thread, CpuLocal};
 use crate::lock::{Gutex, GutexGroup, GutexWrite};
+use alloc::collections::linked_list::LinkedList;
 use alloc::collections::VecDeque;
 use alloc::string::String;
 use alloc::sync::Arc;
@@ -21,6 +22,8 @@ pub struct UmaZone {
     bucket_zones: Arc<Vec<UmaZone>>,
     ty: ZoneType,
     size: NonZero<usize>,                                           // uz_size
+    kegs: LinkedList<UmaKeg>,                                       // uz_kegs + uz_klink
+    slab: fn(&Self, Option<&UmaKeg>, Alloc) -> Option<()>,          // uz_slab
     caches: CpuLocal<RefCell<UmaCache>>,                            // uz_cpu
     full_buckets: Gutex<VecDeque<UmaBox<UmaBucket<[BucketItem]>>>>, // uz_full_bucket
     free_buckets: Gutex<VecDeque<UmaBox<UmaBucket<[BucketItem]>>>>, // uz_free_bucket
@@ -51,7 +54,7 @@ impl UmaZone {
         flags: UmaFlags,
     ) -> Self {
         let name = name.into();
-        let (keg, flags) = if flags.has(UmaFlags::Secondary) {
+        let (keg, mut flags) = if flags.has(UmaFlags::Secondary) {
             todo!()
         } else {
             // We use a different approach here to make it idiomatic to Rust. On Orbis it will
@@ -103,20 +106,32 @@ impl UmaZone {
 
         // Construct uma_zone.
         let gg = GutexGroup::new();
+        let inherit = UmaFlags::Offpage
+            | UmaFlags::Malloc
+            | UmaFlags::Hash
+            | UmaFlags::RefCnt
+            | UmaFlags::VToSlab
+            | UmaFlags::Bucket
+            | UmaFlags::Internal
+            | UmaFlags::CacheOnly;
+
+        flags |= keg.flags() & inherit;
 
         Self {
             bucket_enable,
             bucket_keys,
             bucket_zones,
             ty,
             size: keg.size(),
+            kegs: LinkedList::from([keg]),
+            slab: Self::fetch_slab,
             caches: CpuLocal::new(|_| RefCell::default()),
             full_buckets: gg.clone().spawn_default(),
             free_buckets: gg.clone().spawn_default(),
             alloc_count: gg.clone().spawn_default(),
             free_count: gg.clone().spawn_default(),
             count: gg.spawn(count),
-            flags: flags | (keg.flags() & UmaFlags::from(0xa2002518)), // TODO: Use named flags.
+            flags,
         }
     }
 
@@ -279,9 +294,40 @@ impl UmaZone {
     /// | Version | Offset |
     /// |---------|--------|
     /// |PS4 11.00|0x13DD50|
-    fn alloc_item(&self, _: Alloc) -> *mut u8 {
+    fn alloc_item(&self, flags: Alloc) -> *mut u8 {
+        // Get a slab.
+        let slab = (self.slab)(self, None, flags);
+
+        if slab.is_some() {
+            todo!()
+        }
+
         todo!()
     }
+
+    /// See `zone_fetch_slab` on the Orbis for a reference.
+    ///
+    /// # Reference offsets
+    /// | Version | Offset |
+    /// |---------|--------|
+    /// |PS4 11.00|0x141DB0|
+    fn fetch_slab(&self, keg: Option<&UmaKeg>, flags: Alloc) -> Option<()> {
+        let keg = keg.unwrap_or(self.kegs.front().unwrap());
+
+        if !keg.flags().has(UmaFlags::Bucket) || keg.recurse() == 0 {
+            loop {
+                if let Some(v) = keg.fetch_slab(self, flags) {
+                    return Some(v);
+                }
+
+                if flags.has(Alloc::NoWait | Alloc::NoVm) {
+                    break;
+                }
+            }
+        }
+
+        None
+    }
 }
 
 /// Type of [`UmaZone`].