From 93c695b93d0cbf09b11db3ab6694059fad1c1ee5 Mon Sep 17 00:00:00 2001
From: Putta Khunchalee <ultimaweapon@outlook.com>
Date: Mon, 13 Jan 2025 02:12:10 +0700
Subject: [PATCH] Fixes incorrect keg_small_init implementation

---
 kernel/src/uma/keg.rs | 32 ++++++++++++++++++++++++++++----
 kernel/src/uma/mod.rs | 12 ++++++++++--
 2 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/kernel/src/uma/keg.rs b/kernel/src/uma/keg.rs
index 50e86b27b..a4e85cd6f 100644
--- a/kernel/src/uma/keg.rs
+++ b/kernel/src/uma/keg.rs
@@ -2,7 +2,9 @@ use super::slab::RcFree;
 use super::UmaFlags;
 use crate::config::PAGE_SIZE;
 use crate::uma::slab::{Free, SlabHdr};
+use crate::uma::Uma;
 use core::alloc::Layout;
+use core::cmp::max;
 use core::num::NonZero;
 
 /// Implementation of `uma_keg` structure.
@@ -15,7 +17,7 @@ impl UmaKeg {
     /// | Version | Offset |
     /// |---------|--------|
     /// |PS4 11.00|0x13CF40|
-    pub(super) fn new(size: NonZero<usize>, _: usize, mut flags: UmaFlags) -> Self {
+    pub(super) fn new(size: NonZero<usize>, align: usize, mut flags: UmaFlags) -> Self {
         if flags.has(UmaFlags::Vm) {
             todo!()
         }
@@ -41,10 +43,12 @@ impl UmaKeg {
 
             min = min.pad_to_align();
 
+            // Get UMA_FRITM_SZ and UMA_FRITMREF_SZ.
+            let free_item = min.size() - off;
+            let available = PAGE_SIZE.get() - min.size();
+
             // TODO: Not sure why we need space at least for 2 free item?
-            if (size.get() + (min.size() - off)) > (PAGE_SIZE.get() - min.size()) {
-                todo!()
-            } else {
+            if (size.get() + free_item) > available {
                 // TODO: Set uk_ppera, uk_ipers and uk_rsize.
                 if !flags.has(UmaFlags::Internal) {
                     flags |= UmaFlags::Offpage;
@@ -53,6 +57,26 @@ impl UmaKeg {
                         flags |= UmaFlags::Hash;
                     }
                 }
+            } else {
+                // Get uk_rsize.
+                let rsize = max(size, Uma::SMALLEST_UNIT);
+                let rsize = if (align & rsize.get()) == 0 {
+                    rsize.get()
+                } else {
+                    // Size is not multiple of alignment, align up.
+                    align + 1 + (!align & rsize.get())
+                };
+
+                // Get uk_ipers.
+                let ipers = available / (rsize + free_item);
+
+                // TODO: Verify if this valid for PAGE_SIZE < 0x4000.
+                if !flags.has(UmaFlags::Internal | UmaFlags::Cacheonly)
+                    && (available % (rsize + free_item)) >= Uma::MAX_WASTE.get()
+                    && (PAGE_SIZE.get() / rsize) > ipers
+                {
+                    todo!()
+                }
             }
         }
 
diff --git a/kernel/src/uma/mod.rs b/kernel/src/uma/mod.rs
index 09cd81f7d..0a7bf2534 100644
--- a/kernel/src/uma/mod.rs
+++ b/kernel/src/uma/mod.rs
@@ -1,5 +1,6 @@
 pub use self::zone::*;
 
+use crate::config::PAGE_SIZE;
 use alloc::string::String;
 use alloc::sync::Arc;
 use core::num::NonZero;
@@ -14,8 +15,13 @@ mod zone;
 pub struct Uma {}
 
 impl Uma {
-    /// See `uma_startup` on the Orbis for a reference. Beware that our implementation cannot access
-    /// the CPU context due to this function can be called before context activation.
+    /// `UMA_SMALLEST_UNIT`.
+    const SMALLEST_UNIT: NonZero<usize> = NonZero::new(PAGE_SIZE.get() / 256).unwrap();
+
+    /// `UMA_MAX_WASTE`.
+    const MAX_WASTE: NonZero<usize> = NonZero::new(PAGE_SIZE.get() / 10).unwrap();
+
+    /// See `uma_startup` on the Orbis for a reference.
     ///
     /// # Reference offsets
     /// | Version | Offset |
@@ -67,4 +73,6 @@ pub enum UmaFlags {
     VToSlab = 0x2000,
     /// `UMA_ZFLAG_INTERNAL`.
     Internal = 0x20000000,
+    /// `UMA_ZFLAG_CACHEONLY`.
+    Cacheonly = 0x80000000,
 }