diff --git a/arch/riscv/common.s b/arch/riscv/common.s
index 8edd3ef..c121f5c 100644
--- a/arch/riscv/common.s
+++ b/arch/riscv/common.s
@@ -2,6 +2,8 @@
 # v. 2.0. If a copy of the MPL was not distributed with this file, You can
 # obtain one at https://mozilla.org/MPL/2.0/.
 
+.include "./arch/riscv/constants.s"
+
 .section .text
 
 # a0 - First pointer.
@@ -237,3 +239,70 @@ memset:
 
 .Lmemset_end:
 	ret
+
+# Maps a page of physical memory to the identical virtual memory.
+#
+# Parameters:
+# a0 - Corresponding virtual address.
+# a1 - Physical page address.
+# a2 - Free memory page that can be used to create a leaf page table.
+# a3 - Page table level. If a3 is 0 a superpage is allocated, otherwise a normal page.
+# a4 - Flags.
+#
+# If the function creates new page tables in a2, it advances a2 by the page size and
+# returns it in a0.
+.type pmem_map_at, @function
+.globl pmem_map_at
+pmem_map_at:
+	csrr t0, satp
+	slli t0, t0, 12 # Root table address.
+
+	li t4, 0xfffffc00
+	li t5, 0xffc # 10 bit * PTESIZE mask.
+	li t1, 20
+
+.Lpmem_map_at_loop:
+	# Multiply VPN[i] by 4 and add to the start address of the page table.
+	# This gives the physical address of the page table entry.
+	srl t2, a0, t1
+	and t2, t2, t5 # VPN[i] * PTESIZE.
+	add t2, t0, t2 # a + VPN[i] * PTESIZE.
+
+	beqz a3, .Lpmem_map_at_leaf
+
+	lw t3, (t2)
+	andi t6, t3, 0b1
+
+	beqz t6, .Lpmem_map_at_create
+
+	and t0, t3, t4
+	slli t0, t0, 2
+
+	j .Lpmem_map_at_next
+
+.Lpmem_map_at_create:
+	# Take a chunk of free memory and use it as the next page table.
+	# The beginning addres off the chunk is the base address (a) for the next level.
+	mv t0, a2
+	srli t3, t0, 2
+	ori t3, t3, 0b1
+	sw t3, (t2)
+
+	li t6, PAGE_SIZE
+	add a2, a2, t6
+
+.Lpmem_map_at_next:
+	# Calculate the next page table address for the next level.
+	addi t1, t1, -10
+	addi a3, a3, -1
+
+	j .Lpmem_map_at_loop
+
+.Lpmem_map_at_leaf:
+	srli t3, a1, 2 # Physical page number mapped to 12 bits of the page table entry.
+	ori a4, a4, 0b1
+	or t3, t3, a4 # Execute, write, read and valid.
+	sw t3, (t2)
+
+	mv a0, a2
+	ret
diff --git a/arch/riscv/constants.s b/arch/riscv/constants.s
new file mode 100644
index 0000000..44ecb51
--- /dev/null
+++ b/arch/riscv/constants.s
@@ -0,0 +1 @@
+.equ PAGE_SIZE, 4096
diff --git a/arch/riscv/kernel.s b/arch/riscv/kernel.s
index efbbc7c..d1a78e2 100644
--- a/arch/riscv/kernel.s
+++ b/arch/riscv/kernel.s
@@ -2,6 +2,8 @@
 # v. 2.0. If a copy of the MPL was not distributed with this file, You can
 # obtain one at https://mozilla.org/MPL/2.0/.
 
+.include "./arch/riscv/constants.s"
+
 .section .rodata
 
 panic_message: .asciz "\nPanic"
@@ -104,12 +106,10 @@ set_up_stack:
 # a1 - Memory end pointer.
 .type imem_zero, @function
 imem_zero:
-	mv t0, a0
-
 .Limem_zero_loop:
-	bgeu t0, a1, .Limem_zero_end
-	sw zero, (t0)
-	addi t0, t0, 4
+	bgeu a0, a1, .Limem_zero_end
+	sw zero, (a0)
+	addi a0, a0, 4
 
 	j .Limem_zero_loop
 
@@ -125,7 +125,7 @@ imem_mirror:
 	csrr t0, satp
 	slli t0, t0, 12
 
-	li t3, 0x1000
+	li t3, PAGE_SIZE
 	add t1, t0, t3 # Original page table.
 	add t2, t1, t3 # Copied page table.
 
@@ -167,9 +167,10 @@ imem_initialize:
 	sw s1, 16(sp)
 	sw s2, 12(sp)
 	sw s3, 8(sp)
+	sw a0, 4(sp)
 
 	# Set s1 to the beginning of free memory (skipping root page table).
-	li s1, 4096
+	li s1, PAGE_SIZE
 	add s1, a1, s1
 
 	slli a3, a3, 1 # Multiply by two to account for the copy in virtual memory.
@@ -190,28 +191,28 @@ imem_initialize:
 	srli t0, a1, 12
 	csrw satp, t0
 
+	# Zero memory for page tables.
+	mv a0, a1
+	add a1, s1, a3
+	call imem_zero
+
 	# Executable section.
-	# a0 is coming from the parameter.
-	addi a1, s2, %lo(etext)
-	mv a2, s1
-	li a3, 1
-	li a4, 0b1010
+	mv a0, s1
+	lw a1, 4(sp)
+	addi a2, s2, %lo(etext)
+	li a3, 0b1010 # Read and execute.
 	jal imem_map_range
 
 	# Read only section.
-	addi a0, s2, %lo(etext)
-	addi a1, s3, %lo(_data)
-	mv a2, s1
-	li a3, 1
-	li a4, 0b0010
+	addi a1, s2, %lo(etext)
+	addi a2, s3, %lo(_data)
+	li a3, 0b0010 # Read.
 	jal imem_map_range
 
 	# Data section.
-	addi a0, s3, %lo(_data)
-	lw a1, 20(sp)
-	mv a2, s1
-	li a3, 1
-	li a4, 0b0110
+	addi a1, s3, %lo(_data)
+	lw a2, 20(sp)
+	li a3, 0b0110 # Read and write.
 	jal imem_map_range
 
 	# Enable paging.
@@ -236,11 +237,12 @@ imem_initialize:
 # Maps multiple pages of physical memory to the identical virtual memory.
 #
 # Parameters:
-# a0 - Start address.
-# a1 - End address.
-# a2 - Start of the free memory that can be used to create leaf page tables.
-# a3 - Page table level. If a3 is 0 a superpage is allocated, otherwise a normal page.
-# a4 - Flags.
+# a0 - Start of the free memory that can be used to create leaf page tables.
+# a1 - Start address.
+# a2 - End address.
+# a3 - Flags.
+#
+# If the function creates new page tables it advances and returns a0.
 .type imem_map_range, @function
 imem_map_range:
 	# Prologue.
@@ -250,27 +252,30 @@ imem_map_range:
 	addi s0, sp, 32
 
 	sw s1, 20(sp)
-	sw a2, 16(sp)
-	sw a3, 12(sp)
-	sw a4, 8(sp)
+	sw s2, 16(sp)
+	sw s3, 12(sp)
+	sw s4, 8(sp)
 
 	mv s1, a1
+	mv s2, a2
+	li s3, PAGE_SIZE
+	mv s4, a3
 
 .Limem_map_range_loop:
-	bge a0, s1, .Limem_map_range_end
+	bge s1, s2, .Limem_map_range_end
 
-	mv a1, a0
-	lw a2, 16(sp)
-	lw a3, 12(sp)
-	lw a4, 8(sp)
+	mv a1, s1
+	mv a2, s4
 	jal imem_map_at
 
-	li t0, 4096
-	add a0, a0, t0
+	add s1, s1, s3
 
 	j .Limem_map_range_loop
 
 .Limem_map_range_end:
+	lw s4, 8(sp)
+	lw s3, 12(sp)
+	lw s2, 16(sp)
 	lw s1, 20(sp)
 
 	# Epilogue.
@@ -282,62 +287,53 @@ imem_map_range:
 # Maps a page of physical memory to the identical virtual memory.
 #
 # Parameters:
-# a0 - Corresponding virtual address.
+# a0 - Free memory page that can be used to create a leaf page table.
 # a1 - Physical page address.
-# a2 - Free memory page that can be used to create a leaf page table.
-# a3 - Page table level. If a3 is 0 a superpage is allocated, otherwise a normal page.
-# a4 - Flags.
+# a2 - Flags.
 #
-# Returns virtual page address in a0.
+# If the function creates new page tables it advances and returns a0.
 .type imem_map_at, @function
 imem_map_at:
 	csrr t0, satp
 	slli t0, t0, 12 # Root table address.
 
-	li t4, 0xfffffc00
-	li t5, 0xffc # 10 bit * PTESIZE mask.
-	li t1, 20
+	li t1, 0xffc # 10 bit * PTESIZE mask.
 
-.Limem_map_at_loop:
-	# Multiply VPN[i] by 4 and add to the start address of the page table.
-	# This gives the physical address of the page table entry.
-	srl t2, a0, t1
-	and t2, t2, t5 # VPN[i] * PTESIZE.
-	add t2, t0, t2 # a + VPN[i] * PTESIZE.
-
-	beqz a3, .Limem_map_at_leaf
+	srli t2, a1, 20
+	and t2, t2, t1 # VPN[1] * PTESIZE.
+	add t2, t0, t2 # a + VPN[1] * PTESIZE.
 
 	lw t3, (t2)
-	andi t6, t3, 0b1
+	andi t4, t3, 0b1
 
-	beqz t6, .Limem_map_at_create
+	beqz t4, .Limem_map_at_create
 
+	li t4, 0xfffffc00
 	and t0, t3, t4
 	slli t0, t0, 2
 
-	j .Limem_map_at_next
+	j .Limem_map_at_existing
 
 .Limem_map_at_create:
 	# Take a chunk of free memory and use it as the next page table.
 	# The beginning addres off the chunk is the base address (a) for the next level.
-	mv t0, a2
+	mv t0, a0
 	srli t3, t0, 2
 	ori t3, t3, 0b1
 	sw t3, (t2)
 
-.Limem_map_at_next:
-	# Calculate the next page table address for the next level.
-	addi t1, t1, -10
-	addi a3, a3, -1
+	li t4, PAGE_SIZE
+	add a0, a0, t4
 
-	j .Limem_map_at_loop
+.Limem_map_at_existing:
+	srli t2, a1, 10
+	and t2, t2, t1 # VPN[0] * PTESIZE.
+	add t2, t0, t2 # a + VPN[0] * PTESIZE.
 
-.Limem_map_at_leaf:
 	srli t3, a1, 2 # Physical page number mapped to 12 bits of the page table entry.
-	ori a4, a4, 0b1
-	or t3, t3, a4 # Execute, write, read and valid.
+	ori a2, a2, 0b1
+	or t3, t3, a2 # Execute, write, read and valid.
 	sw t3, (t2)
-
 	ret
 
 .section .text
@@ -361,7 +357,7 @@ kernel_main:
 	lw a2, 20(sp)
 	li a3, 0
 	li a4, 0b0010
-	call imem_map_at
+	call pmem_map_at
 
 	lw a0, 16(sp)
 	call device_tree