From 4fd8d69133d8e04707bc49397f332a47cc280703 Mon Sep 17 00:00:00 2001
From: Hidenori Matsubayashi <Hidenori.Matsubayashi@gmail.com>
Date: Fri, 16 Aug 2024 13:24:40 +0900
Subject: [PATCH] mm: implement malloc/free for kernel uncached memory more
 correctly

---
 src/mm/mm_local.h                      | 11 +++++
 src/mm/uncached/kmm_uncached_malloc.cc | 58 ++++++++++++++++++--------
 src/mm/uncached/kmm_uncached_malloc.h  |  2 +
 3 files changed, 54 insertions(+), 17 deletions(-)
 create mode 100644 src/mm/mm_local.h

diff --git a/src/mm/mm_local.h b/src/mm/mm_local.h
new file mode 100644
index 0000000..e5fe824
--- /dev/null
+++ b/src/mm/mm_local.h
@@ -0,0 +1,11 @@
+#ifndef EVISOR_MM_MM_LOCAL_H_
+#define EVISOR_MM_MM_LOCAL_H_
+
+#include <cstddef>
+
+struct PageBlock {
+  size_t size;
+  uint8_t ref_count;
+};
+
+#endif  // EVISOR_MM_MM_LOCAL_H_
diff --git a/src/mm/uncached/kmm_uncached_malloc.cc b/src/mm/uncached/kmm_uncached_malloc.cc
index bd8e663..b4bcc6c 100644
--- a/src/mm/uncached/kmm_uncached_malloc.cc
+++ b/src/mm/uncached/kmm_uncached_malloc.cc
@@ -1,41 +1,65 @@
 #include "mm/uncached/kmm_uncached_malloc.h"
 
 #include <cstdbool>
-#include <cstdint>
 
 #include "arch/ld_symbols.h"
+#include "common/assert.h"
 #include "common/logger.h"
+#include "mm/mm_local.h"
 #include "mm/pgtable.h"
 
-// FIXME
-/* 1MiB uncached space */
-#define UNCACHED_PAGES (static_cast<uint64_t>(1024 * 1024 * 1) / PAGE_SIZE)
-
 namespace evisor {
 
 namespace {
-bool kernelUncachedMemMap[UNCACHED_PAGES] = {false};
-}  // namespace
 
-void* kmm_uncached_malloc([[maybe_unused]] size_t size) {
-  if (size < 1) {
-    return nullptr;
+// TODO: use `kUserSize` variable
+const size_t kRegionPageNums = static_cast<size_t>(1) * 1024 * 1024 / PAGE_SIZE;
+PageBlock kernelUncachedMemoryRegionMap_[kRegionPageNums] = {};
+
+bool kmm_is_spaces_available(size_t start, size_t num_pages) {
+  for (size_t i = 0; i < num_pages; i++) {
+    if (kernelUncachedMemoryRegionMap_[start + i].ref_count != 0) {
+      return false;
+    }
   }
+  return true;
+}
+
+}  // namespace
+
+void* kmm_uncached_malloc(size_t size) {
+  ASSERT(size >= 1, "The size value should be 1 or more");
+
+  const size_t aligned_size = __builtin_align_up(size, PAGE_SIZE);
+  const size_t num_pages = aligned_size / PAGE_SIZE;
 
-  // FIXME: need to consider 'size' var
-  for (uint32_t i = 0; i < UNCACHED_PAGES; i++) {
-    if (!kernelUncachedMemMap[i]) {
-      kernelUncachedMemMap[i] = true;
-      return reinterpret_cast<uint64_t*>(kUncachedStart + i * PAGE_SIZE);
+  for (size_t i = 0; i < kRegionPageNums; i++) {
+    const uint64_t paddr = kUncachedStart + i * PAGE_SIZE;
+    if (kmm_is_spaces_available(i, num_pages)) {
+      for (size_t j = 0; j < num_pages; j++) {
+        auto& page = kernelUncachedMemoryRegionMap_[i + j];
+        page.ref_count++;
+        page.size = num_pages;
+      }
+      return reinterpret_cast<uint64_t*>(paddr);
     }
   }
+
   PANIC("No free space!");
   return nullptr;
 }
 
 void kmm_uncached_free(void* va) {
-  kernelUncachedMemMap[(reinterpret_cast<uint64_t>(va) - kUncachedStart) /
-                       PAGE_SIZE] = false;
+  ASSERT(va != nullptr, "Address must not be null");
+
+  const auto page_idx =
+      (reinterpret_cast<uint64_t>(va) - kUncachedStart) / PAGE_SIZE;
+  const size_t page_len = kernelUncachedMemoryRegionMap_[page_idx].size;
+  for (size_t i = 0; i < page_len; i++) {
+    auto& page = kernelUncachedMemoryRegionMap_[page_idx];
+    page.size = 0;
+    page.ref_count--;
+  }
 }
 
 }  // namespace evisor
diff --git a/src/mm/uncached/kmm_uncached_malloc.h b/src/mm/uncached/kmm_uncached_malloc.h
index 878b3de..54b4630 100644
--- a/src/mm/uncached/kmm_uncached_malloc.h
+++ b/src/mm/uncached/kmm_uncached_malloc.h
@@ -5,7 +5,9 @@
 
 namespace evisor {
 
+// Allocate kernel memory in the uncached memory space.
 void* kmm_uncached_malloc(size_t size);
+// Free kernel memory in the uncached memory space.
 void kmm_uncached_free(void* va);
 
 }  // namespace evisor