From a16dcc2aa93ef4f0f121ff2008515bcdcff1e2ea Mon Sep 17 00:00:00 2001
From: Drew Galbraith <dgalbraith33@gmail.com>
Date: Fri, 3 Nov 2023 19:47:39 -0700
Subject: [PATCH] [Zion] Use a binary tree to store address space mappings.

---
 zion/object/address_space.cpp | 34 +++++++++++++++++-----------------
 zion/object/address_space.h   |  7 ++++---
 2 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/zion/object/address_space.cpp b/zion/object/address_space.cpp
index 8cb6244..636db60 100644
--- a/zion/object/address_space.cpp
+++ b/zion/object/address_space.cpp
@@ -39,13 +39,13 @@ uint64_t AddressSpace::GetNextMemMapAddr(uint64_t size) {
 
 void AddressSpace::MapInMemoryObject(
     uint64_t vaddr, const glcr::RefPtr<MemoryObject>& mem_obj) {
-  memory_mappings_.PushBack({.vaddr = vaddr, .mem_obj = mem_obj});
+  memory_mappings_.Insert(vaddr, {.vaddr = vaddr, .mem_obj = mem_obj});
 }
 
 uint64_t AddressSpace::MapInMemoryObject(
     const glcr::RefPtr<MemoryObject>& mem_obj) {
   uint64_t vaddr = GetNextMemMapAddr(mem_obj->size());
-  memory_mappings_.PushBack({.vaddr = vaddr, .mem_obj = mem_obj});
+  memory_mappings_.Insert(vaddr, {.vaddr = vaddr, .mem_obj = mem_obj});
   return vaddr;
 }
 
@@ -62,12 +62,13 @@ bool AddressSpace::HandlePageFault(uint64_t vaddr) {
     return true;
   }
 
-  MemoryMapping* mapping = GetMemoryMappingForAddr(vaddr);
-  if (mapping == nullptr) {
+  auto mapping_or = GetMemoryMappingForAddr(vaddr);
+  if (!mapping_or) {
     return false;
   }
-  uint64_t offset = vaddr - mapping->vaddr;
-  uint64_t physical_addr = mapping->mem_obj->PhysicalPageAtOffset(offset);
+  MemoryMapping& mapping = mapping_or.value();
+  uint64_t offset = vaddr - mapping.vaddr;
+  uint64_t physical_addr = mapping.mem_obj->PhysicalPageAtOffset(offset);
   if (physical_addr == 0) {
     dbgln("WARN: Memory object returned invalid physical addr.");
     return false;
@@ -79,16 +80,15 @@ bool AddressSpace::HandlePageFault(uint64_t vaddr) {
   return true;
 }
 
-AddressSpace::MemoryMapping* AddressSpace::GetMemoryMappingForAddr(
-    uint64_t vaddr) {
-  auto iter = memory_mappings_.begin();
-  while (iter != memory_mappings_.end()) {
-    if ((vaddr >= (*iter).vaddr) &&
-        (vaddr < ((*iter).vaddr + (*iter).mem_obj->size()))) {
-      return &(*iter);
-    }
-    ++iter;
+glcr::Optional<glcr::Ref<AddressSpace::MemoryMapping>>
+AddressSpace::GetMemoryMappingForAddr(uint64_t vaddr) {
+  auto mapping_or = memory_mappings_.Predecessor(vaddr + 1);
+  if (!mapping_or) {
+    return mapping_or;
   }
-
-  return 0;
+  MemoryMapping& mapping = mapping_or.value();
+  if (mapping.vaddr + mapping.mem_obj->size() <= vaddr) {
+    return {};
+  }
+  return mapping_or;
 }
diff --git a/zion/object/address_space.h b/zion/object/address_space.h
index e02ae4d..691fc08 100644
--- a/zion/object/address_space.h
+++ b/zion/object/address_space.h
@@ -1,6 +1,6 @@
 #pragma once
 
-#include <glacier/container/linked_list.h>
+#include <glacier/container/binary_tree.h>
 #include <glacier/memory/ref_ptr.h>
 #include <stdint.h>
 
@@ -92,7 +92,8 @@ class AddressSpace : public KernelObject {
     uint64_t vaddr;
     glcr::RefPtr<MemoryObject> mem_obj;
   };
-  glcr::LinkedList<MemoryMapping> memory_mappings_;
+  glcr::BinaryTree<uint64_t, MemoryMapping> memory_mappings_;
 
-  MemoryMapping* GetMemoryMappingForAddr(uint64_t vaddr);
+  glcr::Optional<glcr::Ref<MemoryMapping>> GetMemoryMappingForAddr(
+      uint64_t vaddr);
 };