Move Process & Thread to the object folder.

This should prevent me from actually creating 2 shared ptrs of
a single kernel object with their separate ref counts.

lib/mammoth/src/debug.cpp

@@ -1,5 +1,5 @@
 #include "include/mammoth/debug.h"
 
-#include "zcall.h"
+#include <zcall.h>
 
 void dbgln(const char* str) { ZDebug(str); }

lib/mammoth/src/thread.cpp

@@ -1,6 +1,6 @@
 #include "include/mammoth/thread.h"
 
-#include "zcall.h"
+#include <zcall.h>
 
 namespace {
 

sys/test.cpp

@@ -1,6 +1,5 @@
-
-#include "zcall.h"
-#include "zerrors.h"
+#include <zcall.h>
+#include <zerrors.h>
 
 constexpr uint64_t prog2 = 0x00000020'00000000;
 

sys/test2.cpp

@@ -1,5 +1,5 @@
-#include "mammoth/debug.h"
-#include "mammoth/thread.h"
+#include <mammoth/debug.h>
+#include <mammoth/thread.h>
 
 #define CHECK(expr)       \
   {                       \

zion/CMakeLists.txt

@@ -15,12 +15,12 @@ add_executable(zion
   memory/physical_memory.cpp
   memory/user_stack_manager.cpp
   memory/virtual_memory.cpp
+  object/process.cpp
+  object/thread.cpp
   scheduler/context_switch.s
   scheduler/jump_user_space.s
-  scheduler/process.cpp
   scheduler/process_manager.cpp
   scheduler/scheduler.cpp
-  scheduler/thread.cpp
   syscall/syscall.cpp
   syscall/syscall_enter.s
   zion.cpp)

zion/capability/capability.cpp

@@ -1,11 +1,14 @@
 #include "capability/capability.h"
 
+#include "object/process.h"
+#include "object/thread.h"
+
 template <>
 Process& Capability::obj<Process>() {
   if (type_ != PROCESS) {
     panic("Accessing %u cap as object.", type_);
   }
-  return *static_cast<Process*>(obj_);
+  return *static_cast<Process*>(obj_.get());
 }
 
 template <>
@@ -13,5 +16,5 @@ Thread& Capability::obj<Thread>() {
   if (type_ != THREAD) {
     panic("Accessing %u cap as object.", type_);
   }
-  return *static_cast<Thread*>(obj_);
+  return *static_cast<Thread*>(obj_.get());
 }

zion/capability/capability.h

@@ -2,7 +2,8 @@
 
 #include <stdint.h>
 
-#include "debug/debug.h"
+#include "lib/ref_ptr.h"
+#include "object/kernel_object.h"
 
 class Process;
 class Thread;
@@ -14,9 +15,15 @@ class Capability {
     PROCESS,
     THREAD,
   };
-  Capability(void* obj, Type type, uint64_t id, uint64_t permissions)
+  Capability(const RefPtr<KernelObject>& obj, Type type, uint64_t id,
+             uint64_t permissions)
       : obj_(obj), type_(type), id_(id), permissions_(permissions) {}
 
+  template <typename T>
+  Capability(const RefPtr<T>& obj, Type type, uint64_t id, uint64_t permissions)
+      : Capability(StaticCastRefPtr<KernelObject>(obj), type, id, permissions) {
+  }
+
   template <typename T>
   T& obj();
 
@@ -30,8 +37,7 @@ class Capability {
   }
 
  private:
-  // FIXME: This should somehow be a shared ptr to keep the object alive.
-  void* obj_;
+  RefPtr<KernelObject> obj_;
   Type type_;
   uint64_t id_;
   uint64_t permissions_;

zion/lib/ref_counted.h

@@ -0,0 +1,41 @@
+#pragma once
+
+#include <stdint.h>
+
+#include "debug/debug.h"
+
+template <typename T>
+class RefCounted {
+ public:
+  RefCounted() {}
+  ~RefCounted() { dbgln("RefCounted object destroyed"); }
+  void Adopt() {
+    if (ref_count_ != -1) {
+      panic("Adopting owned ptr");
+    } else {
+      ref_count_ = 1;
+    }
+  }
+
+  void Acquire() {
+    if (ref_count_ == -1) {
+      panic("Acquiring unowned ptr");
+    }
+    ref_count_++;
+  }
+  bool Release() {
+    if (ref_count_ == -1 || ref_count_ == 0) {
+      panic("Releasing unowned ptr");
+    }
+    return (--ref_count_) == 0;
+  }
+
+ private:
+  // FIXME: This should be an atomic type.
+  uint64_t ref_count_ = -1;
+  // Disallow copy and move.
+  RefCounted(RefCounted&) = delete;
+  RefCounted(RefCounted&&) = delete;
+  RefCounted& operator=(RefCounted&) = delete;
+  RefCounted& operator=(RefCounted&&) = delete;
+};

zion/lib/ref_ptr.h

@@ -0,0 +1,89 @@
+#pragma once
+
+template <typename T>
+class RefPtr;
+
+template <typename T>
+RefPtr<T> AdoptPtr(T* ptr);
+
+template <typename T, typename U>
+RefPtr<T> StaticCastRefPtr(const RefPtr<U>& ref);
+
+template <typename T>
+class RefPtr {
+ public:
+  RefPtr() : ptr_(nullptr) {}
+  RefPtr(decltype(nullptr)) : ptr_(nullptr) {}
+  RefPtr(const RefPtr& other) : ptr_(other.ptr_) {
+    if (ptr_) {
+      ptr_->Acquire();
+    }
+  }
+  RefPtr& operator=(const RefPtr& other) {
+    T* old = ptr_;
+    ptr_ = other.ptr_;
+    if (ptr_) {
+      ptr_->Acquire();
+    }
+    if (old && old->Release()) {
+      delete old;
+    }
+
+    return *this;
+  }
+
+  RefPtr(RefPtr&& other) : ptr_(other.ptr_) { other.ptr_ = nullptr; }
+  RefPtr& operator=(RefPtr&& other) {
+    // Swap
+    T* ptr = ptr_;
+    ptr_ = other.ptr_;
+    other.ptr_ = ptr;
+    return *this;
+  }
+
+  enum DontAdoptTag {
+    DontAdopt,
+  };
+  RefPtr(T* ptr, DontAdoptTag) : ptr_(ptr) { ptr->Acquire(); }
+
+  T* get() const { return ptr_; };
+  T& operator*() const { return *ptr_; }
+  T* operator->() const { return ptr_; }
+  operator bool() const { return ptr_ != nullptr; }
+
+  bool operator==(decltype(nullptr)) const { return (ptr_ == nullptr); }
+  bool operator!=(decltype(nullptr)) const { return (ptr_ != nullptr); }
+
+  bool operator==(const RefPtr<T>& other) const { return (ptr_ == other.ptr_); }
+  bool operator!=(const RefPtr<T>& other) const { return (ptr_ != other.ptr_); }
+
+ private:
+  T* ptr_;
+
+  friend RefPtr<T> AdoptPtr<T>(T* ptr);
+  RefPtr(T* ptr) : ptr_(ptr) { ptr->Adopt(); }
+};
+
+template <typename T>
+class MakeRefCountedFriend final {
+ public:
+  template <typename... Args>
+  static RefPtr<T> Make(Args&&... args) {
+    return AdoptPtr(new T(args...));
+  }
+};
+
+template <typename T, typename... Args>
+RefPtr<T> MakeRefCounted(Args&&... args) {
+  return MakeRefCountedFriend<T>::Make(args...);
+}
+
+template <typename T>
+RefPtr<T> AdoptPtr(T* ptr) {
+  return RefPtr(ptr);
+}
+
+template <typename T, typename U>
+RefPtr<T> StaticCastRefPtr(const RefPtr<U>& ref) {
+  return RefPtr(static_cast<T*>(ref.get()), RefPtr<T>::DontAdopt);
+}

zion/loader/elf_loader.h

@@ -2,7 +2,7 @@
 
 #include <stdint.h>
 
-#include "scheduler/process.h"
+#include "object/process.h"
 
 // Loads the elf program and returns its entry point.
 uint64_t LoadElfProgram(Process& dest_proc, uint64_t base, uint64_t length);

zion/loader/init_loader.cpp

@@ -2,11 +2,12 @@
 
 #include "boot/boot_info.h"
 #include "debug/debug.h"
+#include "lib/ref_ptr.h"
 #include "loader/elf_loader.h"
 #include "memory/paging_util.h"
-#include "scheduler/process.h"
+#include "object/process.h"
+#include "object/thread.h"
 #include "scheduler/process_manager.h"
-#include "scheduler/thread.h"
 
 namespace {
 
@@ -49,7 +50,7 @@ void LoadInitProgram() {
   const limine_file& init_prog = GetInitProgram("/sys/test");
   const limine_file& prog2 = GetInitProgram("/sys/test2");
 
-  SharedPtr<Process> proc = MakeShared<Process>();
+  RefPtr<Process> proc = Process::Create();
   gProcMan->InsertProcess(proc);
 
   uint64_t entry = LoadElfProgram(

zion/memory/paging_util.h

@@ -2,7 +2,7 @@
 
 #include <stdint.h>
 
-#include "scheduler/process.h"
+#include "object/process.h"
 
 void InitializePml4(uint64_t pml4_physical_addr);
 

zion/object/kernel_object.h

@@ -0,0 +1,5 @@
+#pragma once
+
+#include "lib/ref_counted.h"
+
+class KernelObject : public RefCounted<KernelObject> {};

zion/object/process.cpp

@@ -1,11 +1,11 @@
-#include "scheduler/process.h"
+#include "object/process.h"
 
 #include "debug/debug.h"
 #include "include/zcall.h"
 #include "memory/paging_util.h"
 #include "memory/physical_memory.h"
+#include "object/thread.h"
 #include "scheduler/scheduler.h"
-#include "scheduler/thread.h"
 
 namespace {
 
@@ -13,26 +13,30 @@ static uint64_t gNextId = 1;
 
 }
 
-SharedPtr<Process> Process::RootProcess() {
-  SharedPtr<Process> proc(new Process(0));
+RefPtr<Process> Process::RootProcess() {
+  RefPtr<Process> proc = MakeRefCounted<Process>(0);
   proc->threads_.PushBack(Thread::RootThread(*proc));
   proc->next_thread_id_ = 1;
 
   return proc;
 }
-
-Process::Process() : id_(gNextId++), state_(RUNNING) {
-  caps_.PushBack(new Capability(this, Capability::PROCESS, Z_INIT_PROC_SELF,
-                                ZC_PROC_SPAWN_PROC | ZC_PROC_SPAWN_THREAD));
+RefPtr<Process> Process::Create() {
+  auto proc = MakeRefCounted<Process>();
+  proc->caps_.PushBack(
+      new Capability(proc, Capability::PROCESS, Z_INIT_PROC_SELF,
+                     ZC_PROC_SPAWN_PROC | ZC_PROC_SPAWN_THREAD));
+  return proc;
 }
 
-SharedPtr<Thread> Process::CreateThread() {
-  SharedPtr<Thread> thread{new Thread(*this, next_thread_id_++, 0)};
+Process::Process() : id_(gNextId++), state_(RUNNING) {}
+
+RefPtr<Thread> Process::CreateThread() {
+  RefPtr<Thread> thread = MakeRefCounted<Thread>(*this, next_thread_id_++);
   threads_.PushBack(thread);
   return thread;
 }
 
-SharedPtr<Thread> Process::GetThread(uint64_t tid) {
+RefPtr<Thread> Process::GetThread(uint64_t tid) {
   auto iter = threads_.begin();
   while (iter != threads_.end()) {
     if (iter->tid() == tid) {
@@ -64,12 +68,12 @@ SharedPtr<Capability> Process::GetCapability(uint64_t cid) {
     ++iter;
   }
   dbgln("Bad cap access");
+  dbgln("Num caps: %u", caps_.size());
   return {};
 }
 
-uint64_t Process::AddCapability(SharedPtr<Thread>& thread) {
+uint64_t Process::AddCapability(const RefPtr<Thread>& thread) {
   uint64_t cap_id = next_cap_id_++;
-  caps_.PushBack(
-      new Capability(thread.ptr(), Capability::THREAD, cap_id, ZC_WRITE));
+  caps_.PushBack(new Capability(thread, Capability::THREAD, cap_id, ZC_WRITE));
   return cap_id;
 }

zion/object/process.h

@@ -4,13 +4,14 @@
 
 #include "capability/capability.h"
 #include "lib/linked_list.h"
+#include "lib/ref_ptr.h"
 #include "lib/shared_ptr.h"
 #include "memory/virtual_memory.h"
 
 // Forward decl due to cyclic dependency.
 class Thread;
 
-class Process {
+class Process : public KernelObject {
  public:
   enum State {
     UNSPECIFIED,
@@ -18,17 +19,17 @@ class Process {
     RUNNING,
     FINISHED,
   };
-  static SharedPtr<Process> RootProcess();
-  Process();
+  static RefPtr<Process> RootProcess();
+  static RefPtr<Process> Create();
 
   uint64_t id() const { return id_; }
   VirtualMemory& vmm() { return vmm_; }
 
-  SharedPtr<Thread> CreateThread();
-  SharedPtr<Thread> GetThread(uint64_t tid);
+  RefPtr<Thread> CreateThread();
+  RefPtr<Thread> GetThread(uint64_t tid);
 
   SharedPtr<Capability> GetCapability(uint64_t cid);
-  uint64_t AddCapability(SharedPtr<Thread>& t);
+  uint64_t AddCapability(const RefPtr<Thread>& t);
   // Checks the state of all child threads and transitions to
   // finished if all have finished.
   void CheckState();
@@ -36,6 +37,8 @@ class Process {
   State GetState() { return state_; }
 
  private:
+  friend class MakeRefCountedFriend<Process>;
+  Process();
   Process(uint64_t id) : id_(id), vmm_(VirtualMemory::ForRoot()) {}
   uint64_t id_;
   VirtualMemory vmm_;
@@ -44,6 +47,6 @@ class Process {
   uint64_t next_thread_id_ = 0;
   uint64_t next_cap_id_ = 0x100;
 
-  LinkedList<SharedPtr<Thread>> threads_;
+  LinkedList<RefPtr<Thread>> threads_;
   LinkedList<SharedPtr<Capability>> caps_;
 };

zion/object/thread.cpp

@@ -1,10 +1,10 @@
-#include "scheduler/thread.h"
+#include "object/thread.h"
 
 #include "common/gdt.h"
 #include "debug/debug.h"
 #include "loader/elf_loader.h"
 #include "memory/paging_util.h"
-#include "scheduler/process.h"
+#include "object/process.h"
 #include "scheduler/scheduler.h"
 
 namespace {
@@ -20,12 +20,15 @@ extern "C" void thread_init() {
 
 }  // namespace
 
-SharedPtr<Thread> Thread::RootThread(Process& root_proc) {
-  return new Thread(root_proc);
+RefPtr<Thread> Thread::RootThread(Process& root_proc) {
+  return MakeRefCounted<Thread>(root_proc);
 }
 
-Thread::Thread(Process& proc, uint64_t tid, uint64_t entry)
-    : process_(proc), id_(tid), rip_(entry) {
+RefPtr<Thread> Thread::Create(Process& proc, uint64_t tid) {
+  return MakeRefCounted<Thread>(proc, tid);
+}
+
+Thread::Thread(Process& proc, uint64_t tid) : process_(proc), id_(tid) {
   uint64_t* stack_ptr = proc.vmm().AllocateKernelStack();
   // 0: rip
   *(stack_ptr) = reinterpret_cast<uint64_t>(thread_init);

zion/object/thread.h

@@ -2,12 +2,13 @@
 
 #include <stdint.h>
 
-#include "lib/shared_ptr.h"
+#include "lib/ref_ptr.h"
+#include "object/kernel_object.h"
 
 // Forward decl due to cyclic dependency.
 class Process;
 
-class Thread {
+class Thread : public KernelObject {
  public:
   enum State {
     UNSPECIFIED,
@@ -16,9 +17,8 @@ class Thread {
     RUNNABLE,
     FINISHED,
   };
-  static SharedPtr<Thread> RootThread(Process& root_proc);
-
-  Thread(Process& proc, uint64_t tid, uint64_t entry);
+  static RefPtr<Thread> RootThread(Process& root_proc);
+  static RefPtr<Thread> Create(Process& proc, uint64_t tid);
 
   uint64_t tid() const { return id_; };
   uint64_t pid() const;
@@ -40,6 +40,8 @@ class Thread {
   void Exit();
 
  private:
+  friend class MakeRefCountedFriend<Thread>;
+  Thread(Process& proc, uint64_t tid);
   // Special constructor for the root thread only.
   Thread(Process& proc) : process_(proc), id_(0) {}
   Process& process_;

zion/scheduler/process_manager.cpp

@@ -7,7 +7,7 @@ void ProcessManager::Init() {
   gProcMan->InsertProcess(Process::RootProcess());
 }
 
-void ProcessManager::InsertProcess(const SharedPtr<Process>& proc) {
+void ProcessManager::InsertProcess(const RefPtr<Process>& proc) {
   proc_list_.PushBack(proc);
 }
 

zion/scheduler/process_manager.h

@@ -1,8 +1,8 @@
 #pragma once
 
 #include "lib/linked_list.h"
-#include "lib/shared_ptr.h"
-#include "scheduler/process.h"
+#include "lib/ref_ptr.h"
+#include "object/process.h"
 
 class ProcessManager {
  public:
@@ -10,14 +10,14 @@ class ProcessManager {
   // and stores it in the global variable.
   static void Init();
 
-  void InsertProcess(const SharedPtr<Process>& proc);
+  void InsertProcess(const RefPtr<Process>& proc);
   Process& FromId(uint64_t id);
 
   void DumpProcessStates();
 
  private:
   // TODO: This should be a hashmap.
-  LinkedList<SharedPtr<Process>> proc_list_;
+  LinkedList<RefPtr<Process>> proc_list_;
 };
 
 extern ProcessManager* gProcMan;

zion/scheduler/scheduler.cpp

@@ -48,7 +48,7 @@ void Scheduler::Preempt() {
     return;
   }
 
-  SharedPtr<Thread> prev = current_thread_;
+  RefPtr<Thread> prev = current_thread_;
   prev->SetState(Thread::RUNNABLE);
   current_thread_ = runnable_threads_.CycleFront(prev);
 
@@ -62,7 +62,7 @@ void Scheduler::Yield() {
   }
   asm volatile("cli");
 
-  SharedPtr<Thread> prev = current_thread_;
+  RefPtr<Thread> prev = current_thread_;
   if (prev == sleep_thread_) {
     if (runnable_threads_.size() == 0) {
       panic("Sleep thread yielded without next.");

zion/scheduler/scheduler.h

@@ -1,7 +1,7 @@
 #pragma once
 
-#include "scheduler/process.h"
-#include "scheduler/thread.h"
+#include "object/process.h"
+#include "object/thread.h"
 
 class Scheduler {
  public:
@@ -15,7 +15,7 @@ class Scheduler {
   Process& CurrentProcess() { return current_thread_->process(); }
   Thread& CurrentThread() { return *current_thread_; }
 
-  void Enqueue(const SharedPtr<Thread> thread) {
+  void Enqueue(const RefPtr<Thread>& thread) {
     runnable_threads_.PushBack(thread);
   }
 
@@ -25,10 +25,10 @@ class Scheduler {
  private:
   bool enabled_ = false;
 
-  SharedPtr<Thread> current_thread_;
-  LinkedList<SharedPtr<Thread>> runnable_threads_;
+  RefPtr<Thread> current_thread_;
+  LinkedList<RefPtr<Thread>> runnable_threads_;
 
-  SharedPtr<Thread> sleep_thread_;
+  RefPtr<Thread> sleep_thread_;
 
   Scheduler();
   void SwapToCurrent(Thread& prev);

zion/syscall/syscall.cpp

@@ -6,7 +6,7 @@
 #include "include/zcall.h"
 #include "include/zerrors.h"
 #include "loader/elf_loader.h"
-#include "scheduler/process.h"
+#include "object/process.h"
 #include "scheduler/process_manager.h"
 #include "scheduler/scheduler.h"
 
@@ -71,7 +71,7 @@ uint64_t ProcessSpawnElf(ZProcessSpawnElfReq* req) {
     return ZE_DENIED;
   }
   dbgln("Proc spawn: %u:%u", req->elf_base, req->elf_size);
-  SharedPtr<Process> proc = MakeShared<Process>();
+  RefPtr<Process> proc = Process::Create();
   gProcMan->InsertProcess(proc);
   uint64_t entry = LoadElfProgram(*proc, req->elf_base, req->elf_size);
   proc->CreateThread()->Start(entry, 0, 0);