feat: add RISC-V configuration and memory management enhancements

Co-authored-by: Copilot <copilot@github.com>

6 files changed, 112 insertions, 15 deletions

diff --git a/na/arch/riscv/config.cc b/na/arch/riscv/config.cc
new file mode 100644
index 0000000..139d932
--- /dev/null
+++ b/na/arch/riscv/config.cc
@@ -0,0 +1,10 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later */
+/* Copyright (C) 2026 Richard Qin */
+
+#include <arch/config.h>
+
+namespace arch {
+
+uint64 g_direct_map_base = kDirectMapBaseDefault;
+
+} // namespace arch
diff --git a/na/arch/riscv/include/arch/config.h b/na/arch/riscv/include/arch/config.h
index 86cd8b1..c407b25 100644
--- a/na/arch/riscv/include/arch/config.h
+++ b/na/arch/riscv/include/arch/config.h
@@ -14,11 +14,15 @@
 
 namespace arch {
 
-/* Kernel virtual address offset (higher-half base).
- * SV39: top of 39-bit VA space. VA = PA + kVAOffset. */
-inline constexpr uint64 kVAOffset = 0xFFFFFFC000000000ULL;
+/* Direct-map virtual base (higher-half base).
+ * SV39: top of 39-bit VA space. VA = PA + direct_map_base. */
+inline constexpr uint64 kDirectMapBaseDefault = 0xFFFFFFC000000000ULL;
+extern uint64 g_direct_map_base;
+
+/* Dedicated VA base for the Frame descriptor array (vmemmap). */
+inline constexpr uint64 kVmemmapBase = 0xFFFFFFD000000000ULL;
 
 /* Physical memory base (where RAM starts on RISC-V virt platform) */
 inline constexpr uint64 kPhysBase = 0x80000000ULL;
 
-}  // namespace arch
+} // namespace arch
diff --git a/na/arch/riscv/include/arch/pmap.h b/na/arch/riscv/include/arch/pmap.h
index 3f0b142..4504af7 100644
--- a/na/arch/riscv/include/arch/pmap.h
+++ b/na/arch/riscv/include/arch/pmap.h
@@ -29,6 +29,7 @@ inline constexpr uint64 PTE_D = 1ULL << 7;
 uint64 create();
 void destroy(uint64 root_pa);
 int map(uint64 root_pa, uint64 va, uint64 pa, uint64 perm);
+int map_2mb(uint64 root_pa, uint64 va, uint64 pa, uint64 perm);
 void unmap(uint64 root_pa, uint64 va);
 bool lookup(uint64 root_pa, uint64 va, uint64* pa_out);
 void activate(uint64 root_pa);
diff --git a/na/arch/riscv/mm/pmap.cc b/na/arch/riscv/mm/pmap.cc
index 46db6a0..4ed3bca 100644
--- a/na/arch/riscv/mm/pmap.cc
+++ b/na/arch/riscv/mm/pmap.cc
@@ -18,6 +18,7 @@
 #include <lume/addr.h>
 #include <lume/config.h>
 #include <lume/pmm.h>
+#include <lume/memblock.h>
 #include <arch/mmu.h>
 #include <lume/klog.h>
 
@@ -58,21 +59,33 @@ static inline uint64 pa_to_pte(uint64 pa, uint64 flags)
 /* Convert PA to kernel VA for accessing page table contents */
 static inline uint64 *pa_to_ptr(uint64 pa)
 {
-    return reinterpret_cast<uint64 *>(pa_to_va(pa));
+    return reinterpret_cast<uint64 *>(pa_to_va(phys_addr(pa)).raw);
 }
 
 /* Allocate a zeroed page for a page table level.
+ * During early boot (before PMM), uses memblock.
+ * After PMM init, uses the buddy allocator.
  * Returns PA of the new page, or 0 on failure. */
 static uint64 alloc_pt_page()
 {
-    Frame *f = pmm_alloc_frame();
-    if (!f)
+    uint64 pa;
+    if (memblock_is_active()) {
+        pa = memblock_alloc(kPageSize, kPageSize);
+    } else {
+        Frame *f = pmm_alloc_frame();
+        if (!f) {
+            return 0;
+        }
+        pa = frame_to_pa(f).raw;
+    }
+    if (!pa) {
         return 0;
-    uint64 pa = frame_to_pa(f);
+    }
     /* Zero the page */
     uint64 *ptr = pa_to_ptr(pa);
-    for (int i = 0; i < kPtePerPage; i++)
+    for (int i = 0; i < kPtePerPage; i++) {
         ptr[i] = 0;
+    }
     return pa;
 }
 
@@ -97,13 +110,13 @@ void destroy(uint64 root_pa)
                 if (!(l1[j] & (PTE_R | PTE_W | PTE_X))) {
                     /* L0 page table */
                     uint64 l0_pa = pte_to_pa(l1[j]);
-                    pmm_free_frame(pa_to_frame(l0_pa));
+                    pmm_free_frame(pa_to_frame(phys_addr(l0_pa)));
                 }
             }
-            pmm_free_frame(pa_to_frame(l1_pa));
+            pmm_free_frame(pa_to_frame(phys_addr(l1_pa)));
         }
     }
-    pmm_free_frame(pa_to_frame(root_pa));
+    pmm_free_frame(pa_to_frame(phys_addr(root_pa)));
 }
 
 int map(uint64 root_pa, uint64 va, uint64 pa, uint64 perm)
@@ -128,6 +141,33 @@ int map(uint64 root_pa, uint64 va, uint64 pa, uint64 perm)
     return 0;
 }
 
+int map_2mb(uint64 root_pa, uint64 va, uint64 pa, uint64 perm)
+{
+    /* VA and PA must be 2MB-aligned */
+    constexpr uint64 k2MB = 2ULL * 1024 * 1024;
+    if ((va & (k2MB - 1)) != 0 || (pa & (k2MB - 1)) != 0) {
+        return -22; /* -EINVAL */
+    }
+
+    uint64 *table = pa_to_ptr(root_pa);
+
+    /* Walk L2: create intermediate table if needed */
+    uint64 idx2 = vpn(va, 2);
+    if (!(table[idx2] & PTE_V)) {
+        uint64 child_pa = alloc_pt_page();
+        if (!child_pa) {
+            return -12; /* -ENOMEM */
+        }
+        table[idx2] = pa_to_pte(child_pa, PTE_V);
+    }
+    table = pa_to_ptr(pte_to_pa(table[idx2]));
+
+    /* L1: install leaf PTE (2MB mega page) */
+    uint64 idx1 = vpn(va, 1);
+    table[idx1] = pa_to_pte(pa, perm | PTE_V | PTE_A | PTE_D);
+    return 0;
+}
+
 void unmap(uint64 root_pa, uint64 va)
 {
     uint64 *table = pa_to_ptr(root_pa);
diff --git a/na/arch/riscv/module.mk b/na/arch/riscv/module.mk
index 154ae77..9635fea 100644
--- a/na/arch/riscv/module.mk
+++ b/na/arch/riscv/module.mk
@@ -1,5 +1,8 @@
 # arch/riscv/module.mk
 SRCS_CC += arch/riscv/mm/pmap.cc
+SRCS_CC += arch/riscv/config.cc
+SRCS_CC += arch/riscv/trap/trap_md.cc
+SRCS_CC += arch/riscv/lib/random.cc
 SRCS_S  += arch/riscv/boot/entry.S
 SRCS_S  += arch/riscv/trap/swtch.S
 SRCS_S  += arch/riscv/trap/trap_vector.S
diff --git a/na/arch/riscv/trap/swtch.S b/na/arch/riscv/trap/swtch.S
index 4393505..8be7904 100644
--- a/na/arch/riscv/trap/swtch.S
+++ b/na/arch/riscv/trap/swtch.S
@@ -1,11 +1,28 @@
 /* SPDX-License-Identifier: GPL-2.0-or-later */
 /* Copyright (C) 2026 Richard Qin */
 
+/*
+ * Context Switch & Kernel Thread Trampoline (RISC-V 64)
+ *
+ * context_switch: saves/restores callee-saved registers.
+ * kthread_entry_trampoline: entry point for newly created kernel threads.
+ *   After context_switch loads the new Context, execution lands here via
+ *   Context.ra = kthread_entry_trampoline. The entry function and argument
+ *   are in s0 and s1 respectively, set by sched_init/create_tcb.
+ *
+ * The trampoline calls the MI-layer kthread_entry(entry, arg).
+ *
+ * Reference: docs/specs/scheduler.md §6
+ */
+
     .section .text
+
+/* ========================================================================
+ * void context_switch(Context* old_ctx, Context* next_ctx);
+ *   a0 = old_ctx_ptr, a1 = next_ctx_ptr
+ * ======================================================================== */
     .globl context_switch
     .align 2
-# void context_switch(Context* old_ctx, Context* next_ctx);
-# a0 = old_ctx_ptr, a1 = next_ctx_ptr
 context_switch:
     sd ra, 0(a0)
     sd sp, 8(a0)
@@ -37,4 +54,26 @@ context_switch:
     ld s10, 96(a1)
     ld s11, 104(a1)
 
-    ret
\ No newline at end of file
+    ret
+
+/* ========================================================================
+ * kthread_entry_trampoline
+ *
+ * Reached via context_switch when a new kthread is first scheduled.
+ * Context was set up by create_tcb():
+ *   s0 = entry function pointer
+ *   s1 = argument pointer
+ *
+ * Calls the MI-layer C++ function:
+ *   void kthread_entry(void (*entry)(void*), void* arg);
+ * ======================================================================== */
+    .globl kthread_entry_trampoline
+    .align 2
+kthread_entry_trampoline:
+    mv a0, s0           /* a0 = entry function */
+    mv a1, s1           /* a1 = argument */
+    call kthread_entry  /* MI-layer C++ function */
+
+    /* kthread_entry should never return, but if it does: */
+    wfi
+    j kthread_entry_trampoline + 12  /* loop on wfi */
\ No newline at end of file