Some minor changes to virtual memory handling

+ Should start to think about how virtual memory should be handled
  efficiently, I'm sort of flailing around and not actually doing
  anything beneficial. The mapping system I've so far implemented
  _works_, but it's not good and I'm not happy with it. I might come
  back and fix/improve it, but I think I'll go ahead and do other things for
  a while, might give me a better overhead view of what the virtual
  memory system should do.

+ TODO: should probably document the memory layout I'm going with at the
  moment.

10 files changed, 260 insertions, 39 deletions

diff --git a/TODO.txt b/TODO.txt
index 0d65372..29252d0 100644
--- a/TODO.txt
+++ b/TODO.txt
@@ -16,3 +16,5 @@ init from vmem when jumping to vmem
 
 + should I consider passing more pointers around rather than return structures?
 + should probably come up with some kind of general error handling scheme?
+
++ Fix unmap/map_vmem when called from virtual memory
diff --git a/arch/riscv/common/vmem.c b/arch/riscv/common/vmem.c
index 7b2d66b..6b022e1 100644
--- a/arch/riscv/common/vmem.c
+++ b/arch/riscv/common/vmem.c
@@ -1,33 +1,91 @@
 #include <apos/string.h>
 #include <apos/pmem.h>
 #include <apos/vmem.h>
+#include <apos/mem.h>
+#include <apos/debug.h>
 #include <pages.h>
 
 #define pte_ppn(pte) (((pm_t)(pte)) >> 10)
 #define pte_flags(pte) (((pm_t)(pte)) & 0xff)
 #define to_pte(p, f) ((pm_to_pnum(p) << 10) + (f))
 #define pte_addr(pte) (pnum_to_paddr(pte_ppn(pte)))
+#define vm_to_index(a, o) (pm_to_index(a, o))
+
+#if defined(KERNEL)
+
+static vm_t *tmp_pte = 0;
+void arch_init_vmem(struct vm_branch_t *branch, vm_t pte)
+{
+	(void)branch;
+	tmp_pte = (vm_t *)pte;
+}
+
+/* this seems like an awful idea, should probably try to come up with a smarter
+ * approach? flushing the cache every time is probably pretty slow and also
+ * this code doesn't take into account the cpu ASID, although that's probably
+ * not relevant here */
+
+/* what I probably should do is something like Linux, where the kernel has a
+ * direct (or linear, I suppose?) so that all direct memory accesses are also
+ * accessible from virtual memory. Not entirely sure how I should do that,
+ * wouldn't there be situations where the RAM is larger than the allocated
+ * memory region?
+ *
+ * Or maybe some kind of walker, where I maintain a virtual page table that
+ * mirrors the physical page table? Not sure yet, but this seems to work for
+ * _now_.
+ */
+#define QUOTE2(x) #x
+#define QUOTE(x) QUOTE2(x)
+static void *set_vptr(uint8_t flags, pm_t a)
+{
+	*tmp_pte = to_pte(a, flags);
+	__asm__ ("sfence.vma %0, %1"
+			:: "rk" (0), "rK" (TMP_PTE)
+			: "memory");
+
+	return (void *)TMP_PTE;
+}
+
+#else
+struct vm_branch_t *arch_get_tmp_pte(struct vm_branch_t *branch)
+{
+	enum mm_order_t top = __mm_max_order;
+	while(top){
+		size_t idx = vm_to_index(TMP_PTE, top);
+		branch = (struct vm_branch_t *)pte_addr(branch->leaf[idx]);
+		top--;
+	}
+
+	return branch;
+}
+
+#define set_vptr(flags, a) (a)
+#endif
 
-/* probably not actually this simple, right? */
 void map_vmem(struct vm_branch_t *branch,
 		pm_t paddr, vm_t vaddr, uint8_t flags, enum mm_order_t order)
 {
 	enum mm_order_t top = __mm_max_order;
 	while (top != order) {
-		size_t idx = pm_to_index(vaddr, top);
+		size_t idx = vm_to_index(vaddr, top);
 
 		if (!branch->leaf[idx]) {
 			pm_t new_leaf = alloc_page(MM_KPAGE, 0);
 			branch->leaf[idx] =
 				(struct vm_branch_t *)to_pte(new_leaf, VM_V);
-			memset((void *)new_leaf, 0, sizeof(struct vm_branch_t));
+
+			void *leaf_ptr = (void *)set_vptr(VM_R | VM_W | VM_V, new_leaf);
+			memset(leaf_ptr, 0, sizeof(struct vm_branch_t));
 		}
 
-		branch = (struct vm_branch_t *)pte_addr(branch->leaf[idx]);
+		pm_t pte = (pm_t)branch->leaf[idx];
+		pm_t branch_pptr = (pm_t)pte_addr(pte);
+		branch = (struct vm_branch_t *)set_vptr(VM_R | VM_W | VM_V, branch_pptr);
 		top--;
 	}
 
-	size_t idx = pm_to_index(vaddr, top);
+	size_t idx = vm_to_index(vaddr, top);
 	branch->leaf[idx] = (struct vm_branch_t *)to_pte(paddr, flags);
 }
 
diff --git a/arch/riscv/config.h b/arch/riscv/config.h
index 75c898f..fa3c178 100644
--- a/arch/riscv/config.h
+++ b/arch/riscv/config.h
@@ -10,3 +10,6 @@
 #define PM_STACK_BASE	(PM_KERN + SZ_2M)
 #define PM_STACK_TOP	(PM_STACK_BASE + SZ_2M - 2)
 #define VM_KERN		(-SZ_1G)
+#define TMP_PTE		(-SZ_2G)
+#define ROOT_PTE	(0)
+#define ROOT_REGION	(SZ_4K)
diff --git a/arch/riscv/init/init.c b/arch/riscv/init/init.c
index 48fa661..9b7d371 100644
--- a/arch/riscv/init/init.c
+++ b/arch/riscv/init/init.c
@@ -17,6 +17,12 @@ struct pm_layout_t {
 	pm_t top;
 };
 
+struct pm_orders_t {
+	size_t max_order;
+	size_t bits[10];
+	size_t page_shift;
+};
+
 static struct pm_layout_t get_memlayout(void *fdt)
 {
 	struct cell_info_t ci = get_reginfo(fdt, "/memory");
@@ -125,16 +131,16 @@ static void mark_reserved_mem(void *fdt)
 	mark_area_used(base, top);
 }
 
-static void init_pmem(void *fdt)
+static struct pm_orders_t init_pmem(void *fdt)
 {
 	enum mm_mode_t mmode = get_mmode(fdt);
 
 	size_t max_order = 0;
-	size_t order_width = 9;
+	size_t order_bits = 9;
 	switch(mmode){
 		case Sv32:
 			max_order = 1;
-			order_width = 10;
+			order_bits = 10;
 			break;
 
 		case Sv39:
@@ -146,17 +152,21 @@ static void init_pmem(void *fdt)
 			break;
 	};
 
-	size_t widths[10] = {0};
+	size_t bits[10] = {0};
 	for(size_t i = 0; i <= max_order; ++i)
-		widths[i] = order_width;
+		bits[i] = order_bits;
+
+	init_mem(max_order, bits, 12);
 
-	init_mem(max_order, widths, 12);
+	struct pm_orders_t ret = {max_order, {0}, 12};
+	for(size_t i = 0; i <= __mm_max_order; ++i)
+		ret.bits[i] = bits[i];
+
+	return ret;
 }
 
-static void setup_pmem(void *fdt)
+static struct pm_layout_t setup_pmem(void *fdt)
 {
-	init_pmem(fdt);
-
 	struct pm_layout_t pmem = get_memlayout(fdt);
 
 	pm_t initrd_top = get_initrdtop(fdt);
@@ -199,6 +209,8 @@ static void setup_pmem(void *fdt)
 
 	/* mark reserved mem */
 	mark_reserved_mem(fdt);
+
+	return (struct pm_layout_t){.base = pmap_base, .top = actual_size + pmap_base};
 }
 
 pm_t move_kernel()
@@ -229,10 +241,27 @@ struct vm_branch_t *prepare_vmem()
 	/* map stack */
 	map_vmem(branch, PM_STACK_BASE, PM_STACK_BASE, VM_R | VM_W | VM_V, MM_MPAGE);
 
+	/* map root pte */
+	map_vmem(branch, (pm_t)branch, ROOT_PTE, VM_R | VM_W | VM_V, MM_KPAGE);
+
+	map_vmem(branch, 0x10000000, 0x10000000, VM_R | VM_W | VM_V, MM_KPAGE);
+
 	/* TODO: map more stuff? */
 	return branch;
 }
 
+static vm_t prepare_tmp_pte(struct vm_branch_t *branch)
+{
+	/* make sure the TMP_PTE is mapped to *something* */
+	map_vmem(branch, 0, TMP_PTE, VM_R | VM_W | VM_V, MM_KPAGE);
+
+	struct vm_branch_t *tmp_pte = arch_get_tmp_pte(branch);
+
+	vm_t addr = TMP_PTE + MM_KPAGE_SIZE;
+	map_vmem(branch, (vm_t)tmp_pte, addr, VM_R | VM_W | VM_V, MM_KPAGE);
+	return addr;
+}
+
 void start_vmem(void *fdt, struct vm_branch_t *branch)
 {
 	/* TODO: get ASID from CPU id */
@@ -250,7 +279,10 @@ void start_vmem(void *fdt, struct vm_branch_t *branch)
 	/* Sv57 && Sv64 in the future? */
 }
 
-struct init_data_t populate_initdata(void *fdt, struct vm_branch_t *branch)
+struct init_data_t populate_initdata(void *fdt, struct pm_orders_t o,
+		struct pm_layout_t p,
+		struct vm_branch_t *b,
+		vm_t v)
 {
 	extern char __init_start, __init_end;
 
@@ -261,13 +293,29 @@ struct init_data_t populate_initdata(void *fdt, struct vm_branch_t *branch)
 	d.initrd_base = get_initrdbase(fdt);
 	d.initrd_top = get_initrdtop(fdt);
 
+	d.pmap_base = p.base;
+	d.pmap_top = p.top;
+
 	d.fdt_base = get_fdtbase(fdt);
 	d.fdt_top = get_fdttop(fdt);
 
 	d.stack_base = PM_STACK_BASE;
 	d.stack_top = PM_STACK_TOP;
 
-	d.kernel_vm_base = branch;
+	d.kernel_vm_base = b;
+	d.tmp_pte = v;
+
+	d.max_order = o.max_order;
+	for(size_t i = 0; i <= __mm_max_order; ++i)
+		d.bits[i] = o.bits[i];
+	d.page_shift = o.page_shift;
+
+	/* initialize pmap in vmem */
+	pm_t addr = p.base;
+	while(addr < p.top){
+		map_vmem(b, addr, addr, VM_R | VM_W | VM_V, MM_MPAGE);
+		addr += MM_MPAGE_SIZE;
+	}
 
 	return d;
 }
@@ -276,11 +324,13 @@ void init(void *fdt)
 {
 	init_debug(fdt);
 	dbg_fdt(fdt);
-	setup_pmem(fdt);
 
-	struct vm_branch_t *branch = prepare_vmem();
-	struct init_data_t d = populate_initdata(fdt, branch);
-	start_vmem(fdt, branch);
+	struct pm_orders_t o = init_pmem(fdt);
+	struct pm_layout_t p = setup_pmem(fdt);
+	struct vm_branch_t *b = prepare_vmem();
+	vm_t v = prepare_tmp_pte(b);
+	struct init_data_t d = populate_initdata(fdt, o, p, b, v);
+	start_vmem(fdt, b);
 
 	/* update_pmap(TODO: figure out where to place pmap in vmem); */
 	void (*main)(struct init_data_t) = (void (*)(struct init_data_t))VM_KERN;
diff --git a/arch/riscv/kernel/main.c b/arch/riscv/kernel/main.c
index c8eca26..d4fa7b8 100644
--- a/arch/riscv/kernel/main.c
+++ b/arch/riscv/kernel/main.c
@@ -1,9 +1,37 @@
 #include <apos/debug.h>
 #include <apos/init.h>
 #include <apos/vmem.h>
+#include <apos/mem.h>
+
+#ifdef DEBUG
+
+static void kernel_dbg(void *fdt)
+{
+	struct dbg_info_t dbg = dbg_from_fdt(fdt);
+	dbg_init(dbg.dbg_ptr, dbg.dev);
+}
+
+#else
+#define kernel_dbg(...)
+#endif
+
+static void map_fdt(struct vm_branch_t *branch, vm_t fdt_base, vm_t fdt_top)
+{
+	map_vregion(branch, fdt_base, 0, fdt_top - fdt_base, VM_R | VM_W |VM_V);
+}
 
 void __main main(struct init_data_t d)
 {
+	/* feck */
+	dbg_init(0x10000000, NS16550A);
+	dbg("test\n");
+	init_pmap((void *)d.pmap_base);
+	init_mem(d.max_order, d.bits, d.page_shift);
+	init_vmem(ROOT_PTE, d.tmp_pte);
+	/* TODO: move pmap into vmem somewhere */
+	/* TODO: mark kernel area as used */
+	map_fdt(ROOT_PTE, d.fdt_base, d.fdt_top);
+	kernel_dbg((void *)d.fdt_base);
 	/* TODO: approximate order of business:
 	 * setup debugging in vmem (requires mapping fdt)
 	 * free unnecessary init
diff --git a/common/mem.c b/common/mem.c
index 48bd776..cff11b7 100644
--- a/common/mem.c
+++ b/common/mem.c
@@ -9,19 +9,19 @@ size_t __mm_sizes[10];
 size_t __mm_page_shift;
 size_t __mm_max_order;
 
-void init_mem(size_t max_order, size_t widths[10], size_t page_shift)
+void init_mem(size_t max_order, size_t bits[10], size_t page_shift)
 {
 	__mm_max_order = max_order;
 	__mm_page_shift = page_shift;
 
 	__mm_shifts[0] = 0;
-	__mm_widths[0] = 1 << widths[0];
+	__mm_widths[0] = 1 << bits[0];
 	__mm_sizes[0]  = 1 << __mm_page_shift;
 
-	for(size_t i = 0; i <= __mm_max_order; ++i){
-		__mm_widths[i] = 1 << widths[i];
-		__mm_shifts[i] = __mm_shifts[i - 1] + __mm_widths[i - 1];
-		__mm_sizes[i] = 1 << __mm_shifts[i] << __mm_page_shift;
+	for(size_t i = 1; i <= __mm_max_order; ++i){
+		__mm_widths[i] = 1 << bits[i];
+		__mm_shifts[i] = __mm_shifts[i - 1] + bits[i - 1];
+		__mm_sizes[i] = 1UL << __mm_shifts[i] << __mm_page_shift;
 	}
 }
 
diff --git a/common/vmem.c b/common/vmem.c
index c30f99b..c874e3d 100644
--- a/common/vmem.c
+++ b/common/vmem.c
@@ -33,22 +33,22 @@ struct mm_block_region_t {
 	struct mm_block_t *first;
 };
 
-static struct mm_block_region_t *root_region = 0;
-/* a block represents free regions */
+static struct mm_block_region_t *root_region = (struct mm_block_region_t *)ROOT_REGION;
 static struct mm_block_t *root_block = 0;
 
 #define NODE_REGION(x) ((struct mm_block_region_t *)(((vm_t)x) & (__mm_page_shift - 1)))
+#define NEXT_BLOCK() ((struct mm_block_region_t *)\
+((region_counter * __o_size(MM_O0)) + ROOT_PTE))
 
+/* TODO: mark all used regions, also blocks */
 static struct mm_block_t *get_free_block(struct vm_branch_t *branch)
 {
 	struct mm_block_region_t *region = root_region;
 	size_t region_counter = 1;
-	for(;region; region = region->next){
+	for(; region; region = region->next){
 		if(region->next == 0){
 			pm_t next_pa = alloc_page(MM_O0, 0);
-			struct mm_block_region_t *next_va =
-				(struct mm_block_region_t *)
-				(region_counter * __o_size(MM_O0));
+			struct mm_block_region_t *next_va = NEXT_BLOCK();
 
 			map_vmem(branch, next_pa, (vm_t)next_va,
 					VM_W | VM_R | VM_V, MM_O0);
@@ -71,15 +71,23 @@ static struct mm_block_t *get_free_block(struct vm_branch_t *branch)
 				return &block[i];
 			}
 		}
+
+		region_counter++;
 	}
 
 	return 0;
 }
 
-void init_vmem(struct vm_branch_t *branch)
+void init_vmem(struct vm_branch_t *branch, vm_t tmp_pte)
 {
+#if defined(KERNEL)
+	arch_init_vmem(branch, tmp_pte);
+#else
+	(void)tmp_pte;
+#endif
+
 	pm_t first_block = alloc_page(MM_O0, 0);
-	map_vmem(branch, first_block, 0, VM_W | VM_R | VM_V, MM_O0);
+	map_vmem(branch, first_block, (vm_t)root_region, VM_W | VM_R | VM_V, MM_O0);
 	memset(root_region, 0, __o_size(MM_O0));
 
 	root_region->max_blocks = (__o_size(MM_O0) - sizeof(struct mm_block_region_t))
diff --git a/include/apos/init.h b/include/apos/init.h
index 05e24ba..2dd8a27 100644
--- a/include/apos/init.h
+++ b/include/apos/init.h
@@ -2,6 +2,7 @@
 #define APOS_INIT_H
 
 #include <apos/pmem.h>
+#include <apos/vmem.h>
 
 struct init_data_t {
 	pm_t init_base;
@@ -10,6 +11,9 @@ struct init_data_t {
 	pm_t initrd_base;
 	pm_t initrd_top;
 
+	pm_t pmap_base;
+	pm_t pmap_top;
+
 	pm_t fdt_base;
 	pm_t fdt_top;
 
@@ -17,12 +21,11 @@ struct init_data_t {
 	pm_t stack_top;
 
 	struct vm_branch_t *kernel_vm_base;
+	vm_t tmp_pte;
 
-	/*
-	   TODO: is this necessary?
-	   pm_t pmap_base;
-	   pm_t pmap_top;
-	*/
+	size_t max_order;
+	size_t bits[10];
+	size_t page_shift;
 };
 
 #endif /* APOS_INIT_H */
diff --git a/include/apos/sizes.h b/include/apos/sizes.h
index 58a68bc..392e122 100644
--- a/include/apos/sizes.h
+++ b/include/apos/sizes.h
@@ -1,6 +1,8 @@
 #ifndef APOS_SIZES_H
 #define APOS_SIZES_H
 
+#if defined(__ASSEMBLER__)
+
 #define SZ_1				0x000000000001
 #define SZ_2				0x000000000002
 #define SZ_4				0x000000000004
@@ -56,4 +58,63 @@
 #define SZ_256T				0x400000000000
 #define SZ_512T				0x800000000000
 
+#else
+
+#define SZ_1				0x000000000001UL
+#define SZ_2				0x000000000002UL
+#define SZ_4				0x000000000004UL
+#define SZ_8				0x000000000008UL
+#define SZ_16				0x000000000010UL
+#define SZ_32				0x000000000020UL
+#define SZ_64				0x000000000040UL
+#define SZ_128				0x000000000080UL
+#define SZ_256				0x000000000100UL
+#define SZ_512				0x000000000200UL
+
+#define SZ_1K				0x000000000400UL
+#define SZ_2K				0x000000000800UL
+#define SZ_4K				0x000000001000UL
+#define SZ_8K				0x000000002000UL
+#define SZ_16K				0x000000004000UL
+#define SZ_32K				0x000000008000UL
+#define SZ_64K				0x000000010000UL
+#define SZ_128K				0x000000020000UL
+#define SZ_256K				0x000000040000UL
+#define SZ_512K				0x000000080000UL
+
+#define SZ_1M				0x000000100000UL
+#define SZ_2M				0x000000200000UL
+#define SZ_4M				0x000000400000UL
+#define SZ_8M				0x000000800000UL
+#define SZ_16M				0x000001000000UL
+#define SZ_32M				0x000002000000UL
+#define SZ_64M				0x000004000000UL
+#define SZ_128M				0x000008000000UL
+#define SZ_256M				0x000010000000UL
+#define SZ_512M				0x000020000000UL
+
+#define SZ_1G				0x000040000000UL
+#define SZ_2G				0x000080000000UL
+#define SZ_4G				0x000100000000UL
+#define SZ_8G				0x000200000000UL
+#define SZ_16G				0x000400000000UL
+#define SZ_32G				0x000800000000UL
+#define SZ_64G				0x000400000000UL
+#define SZ_128G				0x000800000000UL
+#define SZ_256G				0x001000000000UL
+#define SZ_512G				0x002000000000UL
+
+#define SZ_1T				0x004000000000UL
+#define SZ_2T				0x008000000000UL
+#define SZ_4T				0x010000000000UL
+#define SZ_8T				0x020000000000UL
+#define SZ_16T				0x040000000000UL
+#define SZ_32T				0x080000000000UL
+#define SZ_64T				0x100000000000UL
+#define SZ_128T				0x200000000000UL
+#define SZ_256T				0x400000000000UL
+#define SZ_512T				0x800000000000UL
+
+#endif
+
 #endif /* APOS_SIZES_H */
diff --git a/include/apos/vmem.h b/include/apos/vmem.h
index d2dcedc..e72f783 100644
--- a/include/apos/vmem.h
+++ b/include/apos/vmem.h
@@ -28,4 +28,12 @@ vm_t map_vregion(struct vm_branch_t *branch, pm_t base, vm_t start, size_t size,
 		uint8_t flags);
 void unmap_vregion(struct vm_branch_t *branch, vm_t start);
 
+void init_vmem(struct vm_branch_t *branch, vm_t tmp_pte);
+
+#if defined(KERNEL)
+void arch_init_vmem(struct vm_branch_t *branch, vm_t tmp_pte);
+#else
+struct vm_branch_t *arch_get_tmp_pte(struct vm_branch_t *branch);
+#endif
+
 #endif /* APOS_VMEM_H */