From e7b5881903ad81efee63966e7d2c42524d465dfc Mon Sep 17 00:00:00 2001
From: Kimplul <kimi.h.kuparinen@gmail.com>
Date: Wed, 30 Oct 2024 02:18:12 +0200
Subject: tests pass

---
 arch/riscv64/kernel/vmem.c | 18 +++++++++++++++---
 include/arch/vmem.h        |  1 +
 src/pmem.c                 |  8 ++++++--
 src/tcb.c                  |  8 +++-----
 4 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/arch/riscv64/kernel/vmem.c b/arch/riscv64/kernel/vmem.c
index 5a44f53..1648706 100644
--- a/arch/riscv64/kernel/vmem.c
+++ b/arch/riscv64/kernel/vmem.c
@@ -309,9 +309,7 @@ stat_t map_vpage(struct vmem *branch, pm_t paddr, vm_t vaddr, vmflags_t flags,
 	}
 
 	size_t idx = vm_to_index(vaddr, top);
-	if (is_branch(branch->leaf[idx]))
-		/* something has gone terribly wrong? */
-		__destroy_branch(branch->leaf[idx]);
+	assert(!is_branch(branch->leaf[idx]));
 
 	branch->leaf[idx] =
 		(struct vmem *)to_pte((pm_t)__pa(paddr), vp_flags(flags));
@@ -482,6 +480,20 @@ void destroy_vmem(struct vmem *b)
 	free_page(MM_KPAGE, (pm_t)b);
 }
 
+void destroy_rpcmem(struct vmem *b)
+{
+	if (!b)
+		return;
+
+	/* only destroy branches that the rpc vmem is sure to own */
+	for (size_t i = CSTACK_PAGE; i < KERNEL_PAGE; ++i) {
+		if (is_branch(b->leaf[i]))
+			__destroy_branch((struct vmem *)pte_addr(b->leaf[i]));
+	}
+
+	free_page(MM_KPAGE, (pm_t)b);
+}
+
 /**
  * Helper for mapping in the kernel virtual page.
  * Remember that the kernel lives on its own in a 2MiB (riscv64) region at
diff --git a/include/arch/vmem.h b/include/arch/vmem.h
index 79c42b9..0f06cc5 100644
--- a/include/arch/vmem.h
+++ b/include/arch/vmem.h
@@ -172,6 +172,7 @@ void use_vmem(struct vmem *b);
  * @param b Virtual memory to destroy.
  */
 void destroy_vmem(struct vmem *b);
+void destroy_rpcmem(struct vmem *b);
 
 /**
  * Raw clone user virtual memory.
diff --git a/src/pmem.c b/src/pmem.c
index 9c839f1..bb9d326 100644
--- a/src/pmem.c
+++ b/src/pmem.c
@@ -277,7 +277,7 @@ static pm_t __alloc_page(enum mm_order order)
 		__get_bit(bucket, a, &set, &bit);
 
 		bmap = __get_set(bucket, set);
-		bmap->size = order_width(order + 1);
+		assert(bmap->size != 0 && bmap->size <= order_width(order + 1));
 
 		bmap->next = NULL;
 		bmap->prev = NULL;
@@ -491,6 +491,9 @@ static void __mark_area_used(pm_t base, pm_t top)
 		return;
 	}
 
+	top = align_up(top, BASE_PAGE_SIZE);
+	base = align_down(base, BASE_PAGE_SIZE);
+
 	size_t area_left = top - base;
 	pm_t runner = base;
 	while (area_left >= BASE_PAGE_SIZE) {
@@ -501,6 +504,8 @@ static void __mark_area_used(pm_t base, pm_t top)
 
 	if (area_left != 0)
 		mark_used(BASE_PAGE, runner);
+
+	info("marked [%lx - %lx] reserved\n", base, top);
 }
 
 /** Helper for keeping track of which memory regions to avoid placing data into. */
@@ -548,7 +553,6 @@ static void __mark_reserved(pm_t ram_base, pm_t ram_size, size_t avoid_count,
 
 		pm_t top = base + size;
 		__mark_area_used(base, top);
-		info("marked [%lx - %lx] reserved\n", base, top);
 	}
 }
 
diff --git a/src/tcb.c b/src/tcb.c
index f5b1652..1281da2 100644
--- a/src/tcb.c
+++ b/src/tcb.c
@@ -107,7 +107,7 @@ static stat_t __init_free_thread(struct tcb *t)
 
 	if (setup_rpc_stack(t)) {
 		destroy_rpc_stack(t);
-		destroy_vmem(t->rpc.vmem);
+		destroy_rpcmem(t->rpc.vmem);
 		destroy_vmem(t->proc.vmem);
 		return ERR_OOMEM;
 	}
@@ -135,7 +135,7 @@ static stat_t __init_owned_thread(struct tcb *p, struct tcb *t)
 
 	if (setup_rpc_stack(t)) {
 		destroy_rpc_stack(t);
-		destroy_vmem(t->rpc.vmem);
+		destroy_rpcmem(t->rpc.vmem);
 		return ERR_OOMEM;
 	}
 
@@ -252,9 +252,7 @@ stat_t destroy_thread(struct tcb *t)
 
 	/* free memory backing rpc stack */
 	destroy_rpc_stack(t);
-
-	/* free rpc vmem */
-	destroy_vmem(t->rpc.vmem);
+	destroy_rpcmem(t->rpc.vmem);
 
 	unqueue_ipi(t);
 
-- 
cgit v1.3