[Zion] Move Memory Mappings to a dedicated tree impl.

lib/glacier/container/binary_tree.h

@@ -107,6 +107,9 @@ void BinaryTree<K, V>::Delete(K key) {
 template <typename K, typename V>
 Optional<Ref<V>> BinaryTree<K, V>::Predecessor(K key) {
   auto current = FindOrInsertionParent(key);
+  if (current.empty()) {
+    return {};
+  }
 
   // The case where the current is the insertion parent and
   // the predecessor is unique. If the key was going to be
@@ -139,6 +142,9 @@ Optional<Ref<V>> BinaryTree<K, V>::Predecessor(K key) {
 template <typename K, typename V>
 Optional<Ref<V>> BinaryTree<K, V>::Successor(K key) {
   auto current = FindOrInsertionParent(key);
+  if (current.empty()) {
+    return {};
+  }
 
   // The case where the current is the insertion parent and
   // the predecessor is unique. If the key was going to be
@@ -171,6 +177,9 @@ Optional<Ref<V>> BinaryTree<K, V>::Successor(K key) {
 template <typename K, typename V>
 Optional<Ref<V>> BinaryTree<K, V>::Find(K key) {
   auto current = FindOrInsertionParent(key);
+  if (current.empty()) {
+    return {};
+  }
   if (current->key == key) {
     return Optional<Ref<V>>(current->value);
   }

lib/glacier/memory/reference.h

@@ -10,6 +10,7 @@ class Ref {
   Ref(Ref&& other) = default;
 
   operator T&() const { return ref_; }
+  T& get() const { return ref_; }
 
  private:
   T& ref_;

zion/lib/memory_mapping_tree.cpp

@@ -0,0 +1,57 @@
+#include "lib/memory_mapping_tree.h"
+
+#include "debug/debug.h"
+
+glcr::ErrorCode MemoryMappingTree::AddInMemoryObject(
+    uint64_t vaddr, const glcr::RefPtr<MemoryObject>& object) {
+  // TODO: This implementation is inefficient as it traverses the tree a lot, we
+  // should have some solution with iterators to avoid this.
+  auto predecessor_or = mapping_tree_.Predecessor(vaddr);
+  if (predecessor_or && predecessor_or.value().get().vaddr_limit > vaddr) {
+    return glcr::ALREADY_EXISTS;
+  }
+  if (mapping_tree_.Find(vaddr)) {
+    return glcr::ALREADY_EXISTS;
+  }
+  auto successor_or = mapping_tree_.Successor(vaddr);
+  if (successor_or &&
+      successor_or.value().get().vaddr_base < vaddr + object->size()) {
+    return glcr::ALREADY_EXISTS;
+  }
+
+  mapping_tree_.Insert(vaddr, MemoryMapping{
+                                  .vaddr_base = vaddr,
+                                  .vaddr_limit = vaddr + object->size(),
+                                  .mem_object = object,
+                              });
+
+  return glcr::OK;
+}
+
+glcr::ErrorCode FreeMemoryRange(uint64_t vaddr_base, uint64_t vaddr_limit) {
+  dbgln("Unhandled free memory range!");
+  return glcr::OK;
+}
+
+glcr::ErrorOr<uint64_t> MemoryMappingTree::GetPhysicalPageAtVaddr(
+    uint64_t vaddr) {
+  auto mapping_or = GetMemoryMappingForAddr(vaddr);
+  if (!mapping_or) {
+    return glcr::NOT_FOUND;
+  }
+  MemoryMapping& mapping = mapping_or.value();
+  return mapping.mem_object->PhysicalPageAtOffset(vaddr - mapping.vaddr_base);
+}
+
+glcr::Optional<glcr::Ref<MemoryMappingTree::MemoryMapping>>
+MemoryMappingTree::GetMemoryMappingForAddr(uint64_t vaddr) {
+  auto mapping_or = mapping_tree_.Predecessor(vaddr + 1);
+  if (!mapping_or) {
+    return mapping_or;
+  }
+  MemoryMapping& mapping = mapping_or.value();
+  if (mapping.vaddr_base + mapping.mem_object->size() <= vaddr) {
+    return {};
+  }
+  return mapping_or;
+}

zion/lib/memory_mapping_tree.h

@@ -0,0 +1,41 @@
+#pragma once
+
+#include <glacier/container/binary_tree.h>
+
+#include "object/memory_object.h"
+
+/* AddressRangeTree stores memory objects referred to by
+ * ranges and ensures those ranges do not overlap.
+ */
+class MemoryMappingTree {
+ public:
+  MemoryMappingTree() = default;
+
+  MemoryMappingTree(const MemoryMappingTree&) = delete;
+  MemoryMappingTree(MemoryMappingTree&&) = delete;
+
+  glcr::ErrorCode AddInMemoryObject(uint64_t vaddr,
+                                    const glcr::RefPtr<MemoryObject>& object);
+
+  glcr::ErrorCode FreeMemoryRange(uint64_t vaddr_base, uint64_t vaddr_limit);
+
+  glcr::ErrorOr<uint64_t> GetPhysicalPageAtVaddr(uint64_t vaddr);
+
+ private:
+  struct MemoryMapping {
+    uint64_t vaddr_base;
+    uint64_t vaddr_limit;
+    glcr::RefPtr<MemoryObject> mem_object;
+  };
+
+  // TODO: Consider adding a red-black tree implementation here.
+  // As is this tree functions about as well as a linked list
+  // because mappings are likely to be added in near-perfect ascedning order.
+  // Also worth considering creating a special tree implementation for
+  // just this purpose, or maybe a BinaryTree implementation that accepts
+  // ranges rather than a single key.
+  glcr::BinaryTree<uint64_t, MemoryMapping> mapping_tree_;
+
+  glcr::Optional<glcr::Ref<MemoryMapping>> GetMemoryMappingForAddr(
+      uint64_t vaddr);
+};

zion/loader/init_loader.cpp

@@ -76,7 +76,8 @@ uint64_t LoadElfProgram(Process& dest_proc, uint64_t base, uint64_t offset) {
 #endif
     auto mem_obj = glcr::MakeRefCounted<MemoryObject>(program.memsz);
     mem_obj->CopyBytesToObject(base + program.offset, program.filesz);
-    dest_proc.vmas()->MapInMemoryObject(program.vaddr, mem_obj);
+    PANIC_ON_ERR(dest_proc.vmas()->MapInMemoryObject(program.vaddr, mem_obj),
+                 "Couldn't map in init program.");
   }
   return header->entry;
 }

zion/object/address_space.cpp

@@ -35,15 +35,15 @@ uint64_t AddressSpace::GetNextMemMapAddr(uint64_t size) {
   return addr;
 }
 
-void AddressSpace::MapInMemoryObject(
+glcr::ErrorCode AddressSpace::MapInMemoryObject(
     uint64_t vaddr, const glcr::RefPtr<MemoryObject>& mem_obj) {
-  memory_mappings_.Insert(vaddr, {.vaddr = vaddr, .mem_obj = mem_obj});
+  return mapping_tree_.AddInMemoryObject(vaddr, mem_obj);
 }
 
-uint64_t AddressSpace::MapInMemoryObject(
+glcr::ErrorOr<uint64_t> AddressSpace::MapInMemoryObject(
     const glcr::RefPtr<MemoryObject>& mem_obj) {
   uint64_t vaddr = GetNextMemMapAddr(mem_obj->size());
-  memory_mappings_.Insert(vaddr, {.vaddr = vaddr, .mem_obj = mem_obj});
+  RET_ERR(mapping_tree_.AddInMemoryObject(vaddr, mem_obj));
   return vaddr;
 }
 
@@ -55,38 +55,23 @@ bool AddressSpace::HandlePageFault(uint64_t vaddr) {
 #if K_VMAS_DEBUG
   dbgln("[VMAS] Page Fault!");
 #endif
+  if (vaddr < kPageSize) {
+    // Invalid page access.
+    return false;
+  }
+
   if (user_stacks_.IsValidStack(vaddr)) {
     MapPage(cr3_, vaddr, phys_mem::AllocatePage());
     return true;
   }
 
-  auto mapping_or = GetMemoryMappingForAddr(vaddr);
-  if (!mapping_or) {
-    return false;
-  }
-  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.");
+  auto offset_or = mapping_tree_.GetPhysicalPageAtVaddr(vaddr);
+  if (!offset_or.ok()) {
     return false;
   }
 #if K_VMAS_DEBUG
   dbgln("[VMAS] Mapping P({x}) at V({x})", physical_addr, vaddr);
 #endif
-  MapPage(cr3_, vaddr, physical_addr);
+  MapPage(cr3_, vaddr, offset_or.value());
   return true;
 }
-
-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;
-  }
-  MemoryMapping& mapping = mapping_or.value();
-  if (mapping.vaddr + mapping.mem_obj->size() <= vaddr) {
-    return {};
-  }
-  return mapping_or;
-}

zion/object/address_space.h

@@ -5,6 +5,7 @@
 #include <stdint.h>
 
 #include "include/ztypes.h"
+#include "lib/memory_mapping_tree.h"
 #include "memory/user_stack_manager.h"
 #include "object/memory_object.h"
 
@@ -69,16 +70,17 @@ class AddressSpace : public KernelObject {
 
   // Maps in a memory object at a specific address.
   // Note this is unsafe for now as it may clobber other mappings.
-  void MapInMemoryObject(uint64_t vaddr,
-                         const glcr::RefPtr<MemoryObject>& mem_obj);
+  [[nodiscard]] glcr::ErrorCode MapInMemoryObject(
+      uint64_t vaddr, const glcr::RefPtr<MemoryObject>& mem_obj);
 
-  uint64_t MapInMemoryObject(const glcr::RefPtr<MemoryObject>& mem_obj);
+  [[nodiscard]] glcr::ErrorOr<uint64_t> MapInMemoryObject(
+      const glcr::RefPtr<MemoryObject>& mem_obj);
 
   // Kernel Mappings.
   uint64_t AllocateKernelStack();
 
   // Returns true if the page fault has been resolved.
-  bool HandlePageFault(uint64_t vaddr);
+  [[nodiscard]] bool HandlePageFault(uint64_t vaddr);
 
  private:
   friend class glcr::MakeRefCountedFriend<AddressSpace>;
@@ -88,19 +90,5 @@ class AddressSpace : public KernelObject {
   UserStackManager user_stacks_;
   uint64_t next_memmap_addr_ = 0x20'00000000;
 
-  struct MemoryMapping {
-    uint64_t vaddr;
-    glcr::RefPtr<MemoryObject> mem_obj;
-  };
-
-  // TODO: Consider adding a red-black tree implementation here.
-  // As is this tree functions about as well as a linked list
-  // because mappings are likely to be added in near-perfect ascedning order.
-  // Also worth considering creating a special tree implementation for
-  // just this purpose, or maybe a BinaryTree implementation that accepts
-  // ranges rather than a single key.
-  glcr::BinaryTree<uint64_t, MemoryMapping> memory_mappings_;
-
-  glcr::Optional<glcr::Ref<MemoryMapping>> GetMemoryMappingForAddr(
-      uint64_t vaddr);
+  MemoryMappingTree mapping_tree_;
 };