From d3024211a76bdaf1d220f686bb47447c21b69b38 Mon Sep 17 00:00:00 2001
From: Drew Galbraith <dgalbraith33@gmail.com>
Date: Thu, 18 May 2023 13:56:54 -0700
Subject: [PATCH] Cycle through multiple tasks in multiple processes

---
 zion/memory/paging_util.cpp  | 19 ++++++++++++++++++-
 zion/scheduler/process.cpp   | 25 +++++++++++++++++++------
 zion/scheduler/scheduler.cpp |  4 +++-
 zion/scheduler/thread.cpp    | 12 +++++++++---
 zion/scheduler/thread.h      |  3 +++
 zion/zion.cpp                |  5 ++++-
 6 files changed, 56 insertions(+), 12 deletions(-)

diff --git a/zion/memory/paging_util.cpp b/zion/memory/paging_util.cpp
index ccad031..89faaf6 100644
--- a/zion/memory/paging_util.cpp
+++ b/zion/memory/paging_util.cpp
@@ -98,20 +98,37 @@ void MapPage(uint64_t virt, uint64_t phys) {
   *PageTableEntry(virt) = PageAlign(phys) | PRESENT_BIT | READ_WRITE_BIT;
   ZeroOutPage(reinterpret_cast<uint64_t*>(virt));
 }
+
+uint64_t Pml4Index(uint64_t addr) { return (addr >> PML_OFFSET) & 0x1FF; }
 }  // namespace
 
 void InitPaging() {
   uint64_t pml4_addr = 0;
   asm volatile("mov %%cr3, %0;" : "=r"(pml4_addr));
-  InitializePml4(pml4_addr);
+  uint64_t* pml4_virtual =
+      reinterpret_cast<uint64_t*>(boot::GetHigherHalfDirectMap() + pml4_addr);
+
+  uint64_t recursive_entry = pml4_addr | PRESENT_BIT | READ_WRITE_BIT;
+  pml4_virtual[0x1FE] = recursive_entry;
 }
 
 void InitializePml4(uint64_t pml4_physical_addr) {
   uint64_t* pml4_virtual = reinterpret_cast<uint64_t*>(
       boot::GetHigherHalfDirectMap() + pml4_physical_addr);
 
+  // Map the recursive entry.
   uint64_t recursive_entry = pml4_physical_addr | PRESENT_BIT | READ_WRITE_BIT;
   pml4_virtual[0x1FE] = recursive_entry;
+
+  // Map the kernel entry.
+  // This should contain the heap at 0xFFFFFFFF'40000000
+  uint64_t kernel_addr = 0xFFFFFFFF'80000000;
+  pml4_virtual[Pml4Index(kernel_addr)] = *Pml4Entry(kernel_addr);
+
+  // Map the HHDM.
+  // This is necessary to access values off of the kernel stack.
+  uint64_t hhdm = boot::GetHigherHalfDirectMap();
+  pml4_virtual[Pml4Index(hhdm)] = *Pml4Entry(hhdm);
 }
 
 void AllocatePage(uint64_t addr) {
diff --git a/zion/scheduler/process.cpp b/zion/scheduler/process.cpp
index fd1fae8..94b8e72 100644
--- a/zion/scheduler/process.cpp
+++ b/zion/scheduler/process.cpp
@@ -1,6 +1,8 @@
 #include "scheduler/process.h"
 
 #include "debug/debug.h"
+#include "memory/paging_util.h"
+#include "memory/physical_memory.h"
 #include "scheduler/scheduler.h"
 #include "scheduler/thread.h"
 
@@ -23,17 +25,28 @@ Process* Process::RootProcess() {
   return proc;
 }
 
+Process::Process() : id_(gNextId++) {
+  cr3_ = phys_mem::AllocatePage();
+  InitializePml4(cr3_);
+  CreateThread();
+}
+
 void Process::CreateThread() {
   Thread* thread = new Thread(this, next_thread_id_++);
-
-  ThreadEntry* entry = thread_list_front_;
-  while (entry->next != nullptr) {
-    entry = entry->next;
-  }
-  entry->next = new ThreadEntry{
+  ThreadEntry* tentry = new ThreadEntry{
       .thread = thread,
       .next = nullptr,
   };
+
+  if (thread_list_front_ == nullptr) {
+    thread_list_front_ = tentry;
+  } else {
+    ThreadEntry* entry = thread_list_front_;
+    while (entry->next != nullptr) {
+      entry = entry->next;
+    }
+    entry->next = tentry;
+  }
   sched::EnqueueThread(thread);
 }
 
diff --git a/zion/scheduler/scheduler.cpp b/zion/scheduler/scheduler.cpp
index 614d653..8d8164b 100644
--- a/zion/scheduler/scheduler.cpp
+++ b/zion/scheduler/scheduler.cpp
@@ -77,7 +77,9 @@ class Scheduler {
     Thread* prev = current_thread_;
     current_thread_ = current_thread_->next_thread_;
     prev->next_thread_ = nullptr;
-    Enqueue(prev);
+    if (prev->pid() != 0) {
+      Enqueue(prev);
+    }
     context_switch(prev->Rsp0Ptr(), current_thread_->Rsp0Ptr());
 
     asm volatile("sti");
diff --git a/zion/scheduler/thread.cpp b/zion/scheduler/thread.cpp
index e42543d..c4d893b 100644
--- a/zion/scheduler/thread.cpp
+++ b/zion/scheduler/thread.cpp
@@ -8,9 +8,8 @@ namespace {
 
 extern "C" void thread_init() {
   asm("sti");
-  dbgln("New Thread!");
-  sched::Yield();
-  panic("End of thread.");
+  sched::CurrentThread().Init();
+  panic("Reached end of thread.");
 }
 
 }  // namespace
@@ -32,3 +31,10 @@ Thread::Thread(Process* proc, uint64_t tid) : process_(proc), id_(tid) {
 }
 
 uint64_t Thread::pid() { return process_->id(); }
+
+void Thread::Init() {
+  while (true) {
+    dbgln("[%u.%u]", pid(), id_);
+    sched::Yield();
+  }
+}
diff --git a/zion/scheduler/thread.h b/zion/scheduler/thread.h
index bbb72ba..33dc09c 100644
--- a/zion/scheduler/thread.h
+++ b/zion/scheduler/thread.h
@@ -18,6 +18,9 @@ class Thread {
 
   uint64_t* Rsp0Ptr() { return &rsp0_; }
 
+  // Called the first time the thread starts up.
+  void Init();
+
   // FIXME: Probably make this private.
   Thread* next_thread_;
 
diff --git a/zion/zion.cpp b/zion/zion.cpp
index e961c82..d72256f 100644
--- a/zion/zion.cpp
+++ b/zion/zion.cpp
@@ -18,7 +18,10 @@ extern "C" void zion() {
   phys_mem::InitPhysicalMemoryManager();
 
   sched::InitScheduler();
-  sched::CurrentProcess().CreateThread();
+  Process p1;
+  p1.CreateThread();
+  Process p2;
+  p2.CreateThread();
   sched::EnableScheduler();
   sched::Yield();