pmm: Implemented nb_alloc() for NBBS.

Signed-off-by: TunaCici <36923239+TunaCici@users.noreply.github.com>

Author: TunaCici

Kernel/Include/Memory/NBBS.h

@@ -31,7 +31,7 @@
 #define BUSY            (OCC | OCC_LEFT | OCC_RIGHT)
 
 #define EXP2(n) (0x1 << (n))
-#define LOG2(n) (64 - __builtin_clzll(n)) // 64 bit
+#define LOG2_LOWER(n) (64 - __builtin_clzll(n) - 1) // 64 bit
 #define CAS(addr, cmp, val) __atomic_compare_exchange (addr, cmp, val, __ATOMIC_SEQ_CST, __ATOMIC_SEQ_CST)
 
 /*
@@ -44,6 +44,14 @@ int  nb_init(uint64_t base_addr, uint64_t size);
 void* nb_alloc(uint64_t size);
 void* nb_free(void *addr);
 
+/*
+ * Private APIs
+ *
+ * TODO: Explain.
+ */
+
+uint32_t try_alloc(uint32_t node);
+
 /*
  * Helpers
  *
@@ -84,3 +92,21 @@ static inline uint8_t is_free(uint8_t val)
 {
         return !(val & BUSY);
 }
+
+/* TODO: Ugly code; refactor */
+static inline uint32_t leftmost(uint32_t node, uint32_t depth)
+{
+        /* Index to level */
+        uint32_t level = LOG2_LOWER(node);
+
+        /* Size (in terms of leaf size) */
+        uint64_t block_size = EXP2(depth - level);
+
+        /* Offset within level */
+        uint64_t offset = node % EXP2(level);
+
+        /* Leftmost leaf */
+        uint32_t leftmost_leaf = EXP2(depth) + offset * block_size;
+
+        return leftmost_leaf;
+}

Kernel/Main.c

@@ -81,10 +81,20 @@ void kmain(boot_sysinfo* boot_params)
         /* 2. Init NBBS */
         klog("[kmain] Initializing NBBS...\n");
 
-        if (nb_init(mem_start, mem_end - mem_start)) {
+        if (nb_init(mem_start, 32 * 1024)) {
                 klog("[kmain] Failed to initialize NBBS ;(\n");
         } else {
-                klog("[kmain] Initialized NBBS!\n");
+                klog("[kmain] Initialized NBBS (%u KiB)!\n", 64u);
+        }
+
+        klog("[kmain] nb_alloc\n");
+
+        uint64_t *bruh = nb_alloc(4096);
+        
+        if (bruh) {
+                klog("[kmain] nb_alloc ok\n");
+        } else {
+                klog("[kmain] nb_alloc fail\n");
         }
 
         /* 2. Init PMM */

Kernel/Memory/NBBS.c

@@ -15,44 +15,46 @@
 
 static volatile uint8_t *tree = 0;
 static volatile uint32_t *index = 0;
+static volatile uint64_t base_address = 0;
 
 static volatile uint32_t depth = 0;
 static volatile uint64_t tree_size = 0; /* bytes */
 static volatile uint64_t index_size = 0; /* bytes */
 
 static volatile uint64_t total_memory = 0;
-static volatile uint64_t max_level = 0; /* [0, depth] */
+static volatile uint64_t max_level = 0;
 static volatile uint64_t min_size = 0;
 static volatile uint64_t max_size = 0;
 
-int nb_init(uint64_t base_addr, uint64_t size)
+int nb_init(uint64_t base, uint64_t size)
 {
-        if (size < min_size) {
+        if (base == 0 | size == 0) {
                 return 1;
         }
 
-        if (base_addr == 0 | size == 0) {
+        if (size < min_size) {
                 return 1;
         }
 
         /* Setup */
+        base_address = base;       
         total_memory = size;
         min_size = PAGE_SIZE;
-        depth = LOG2(total_memory / min_size) - 1;
+        depth = LOG2_LOWER(total_memory / min_size);
         max_level = 0;
-        max_size = (EXP2(depth - max_level) * min_size);
+        max_size = EXP2(depth - max_level) * min_size;
 
         /* Calculate required tree size */
-        uint32_t total_nodes = 1; // we start at idx 1
-        for (uint32_t i = 0; i < depth; i++) {
-                total_nodes += (size / (EXP2(i) * min_size));
+        uint32_t total_nodes = 1; // we have garbage node at idx 0
+        for (uint32_t i = 0; i <= depth; i++) {
+                total_nodes += (total_memory / (EXP2(i) * min_size));
         }
 
         /* Calculate required index size */
-        uint32_t total_elems = (size / min_size);
+        uint32_t total_pages = (total_memory / min_size);
 
-        tree_size = total_nodes * 1;   // each node is 1 byte
-        index_size = total_elems * 4; // each elem is 4 byte
+        tree_size = total_nodes * 1;  // each node is 1 byte
+        index_size = total_pages * 4; // each page id is 4 byte
 
         /* Allocate */
         uint32_t req_pages = (tree_size + PAGE_SIZE - 1) / PAGE_SIZE;
@@ -76,9 +78,54 @@ int nb_init(uint64_t base_addr, uint64_t size)
         return 0;
 }
 
+uint32_t try_alloc(uint32_t node)
+{
+        return 1;
+}
+
 void* nb_alloc(uint64_t size)
 {
+        if (max_size < size) {
+                return 0;
+        }
 
+        if (size < min_size) {
+                size = min_size;
+        }
+
+        uint32_t level = LOG2_LOWER(total_memory / size);
+
+        if (depth < level) {
+                level = depth;
+        }
+
+        /* Range of nodes at target level */
+        uint32_t start_node = EXP2(level);
+        uint32_t end_node = EXP2(level + 1);
+
+        for (uint32_t i = start_node; i < end_node; i++) {
+                if (is_free(tree[i])) {
+                        uint32_t failed_at = try_alloc(i);
+
+                        if (!failed_at) {
+                                /* TODO: Explain what's going on here */
+                                uint32_t leaf = leftmost(i, depth) - EXP2(depth); 
+                                index[leaf] = i;
+
+                                return (void*) (base_address + leaf * min_size);
+                        } else {
+                                /* Skip the entire subtree [of failed] */
+                                uint32_t curr_level = LOG2_LOWER(i);
+                                uint32_t fail_level = LOG2_LOWER(failed_at);
+
+                                uint32_t d = EXP2(curr_level - fail_level);
+                                i = (failed_at + 1) * d;
+                        }
+                }
+        }
+
+        return (void*) 0;
+        
 }
 
 void* nb_free(void *addr)