diff --git a/kernel/src/arch/riscv/vspace/kmem.rs b/kernel/src/arch/riscv/vspace/kmem.rs
index de19125..39e1683 100644
--- a/kernel/src/arch/riscv/vspace/kmem.rs
+++ b/kernel/src/arch/riscv/vspace/kmem.rs
@@ -106,7 +106,7 @@ pub unsafe fn setup_kernel_paging() {
     info!("Setting up kernel paging");
     assert!(!is_kernel_pagetable_allocated(), "Kernel pagetable already allocated");
 
-    let root_pt = alloc_page();
+    let root_pt = alloc_page(VirtAddr(0));
     let mut kernel_pt = Table::<Level0>::new(mmap_phys_to_virt(root_pt));
 
     macro_rules! map_section {
diff --git a/kernel/src/arch/riscv/vspace/utils.rs b/kernel/src/arch/riscv/vspace/utils.rs
index 3f5eca4..7e2121e 100644
--- a/kernel/src/arch/riscv/vspace/utils.rs
+++ b/kernel/src/arch/riscv/vspace/utils.rs
@@ -16,7 +16,7 @@ fn alloc_callback() {
 }
 
 #[cfg_attr(debug_assertions, tracer::trace_callback(log = true, callback = alloc_callback))]
-pub fn alloc_page() -> PhysAddr {
+pub fn alloc_page(_mapped_addr: VirtAddr) -> PhysAddr {
     let addr = RAM_ALLOCATOR.lock().alloc(PAGE_LAYOUT).expect("Failed to allocate page");
 
     unsafe {
@@ -29,7 +29,7 @@ pub fn alloc_page() -> PhysAddr {
 }
 
 pub fn map_range(
-    mut alloc: impl FnMut() -> PhysAddr,
+    mut alloc: impl FnMut(VirtAddr) -> PhysAddr,
     pt: &mut Table<Level0>,
     from: VirtAddr,
     to: PhysAddr,
@@ -63,8 +63,8 @@ pub fn map_range(
             unsafe {
                 Table::<$level>::new(mmap_phys_to_virt(
                     $pt.lookup_mut($addr).map(|e| e.paddr()).unwrap_or_else(|| {
-                        let page = alloc();
                         let addr = $addr.align_down(<$level as TableLevel>::PREVIOUS::LEVEL_SIZE);
+                        let page = alloc(addr);
                         debug!("Creating new pagetable:\t{:?} -> {:?}", addr, page);
                         $pt.map(addr, page, MapAttr::PAGE_TABLE).unwrap();
                         page