Calculate kernel size in pages dinamically

arch/riscv/common.s

@@ -2,9 +2,6 @@
 # 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/.
 
-.global write_s, write_c
-.global memcmp, memchr, memmem, memcpy, bzero
-
 .section .text
 
 # a0 - First pointer.
@@ -13,6 +10,7 @@
 #
 # Returns 0 in a0 if memory regions are equal.
 .type memcmp, @function
+.globl memcmp
 memcmp:
 	mv t0, a0
 	li a0, 0
@@ -41,6 +39,7 @@ memcmp:
 # a0 - String pointer.
 # a1 - Length of the string.
 .type write_s, @function
+.globl write_s
 write_s:
 	# Prologue.
 	addi sp, sp, -32
@@ -80,6 +79,7 @@ write_s:
 # Arguments:
 # a0 - Character.
 .type write_c, @function
+.globl write_c
 write_c:
 	li a1, 0
 	li a2, 0
@@ -102,6 +102,7 @@ write_c:
 # Sets a0 to the pointer to the found character or to null if the character
 # doesn't occur in the memory block.
 .type memchr, @function
+.globl memchr
 memchr:
 .Lmemchr_loop:
 	beqz a2, .Lmemchr_nil # Exit if the length is 0.
@@ -132,6 +133,7 @@ memchr:
 # Sets a0 to the pointer to the beginning of the substring in memory or to 0
 # if the substring doesn't occur in the block.
 .type memmem, @function
+.globl memmem
 memmem:
 	# Prologue.
 	addi sp, sp, -24
@@ -193,6 +195,7 @@ memmem:
 #
 # Preserves a0.
 .type memcpy, @function
+.globl memcpy
 memcpy:
 	mv t0, a0
 
@@ -200,30 +203,37 @@ memcpy:
 	beqz a2, .Lmemcpy_end
 
 	lbu t1, (a1)
-	sb t1, (a0)
+	sb t1, (t0)
 
-	addi a0, a0, 1
+	addi t0, t0, 1
 	addi a1, a1, 1
 	addi a2, a2, -1
 
 	j .Lmemcpy_loop
 
 .Lmemcpy_end:
-	mv a0, t0
 	ret
 
-# Zeroes a chank of memory.
+# Sets a memory region to the given byte value.
 #
 # Parameters:
-# a0 - Memory pointer.
-# a1 - Memory size.
-bzero:
+# a0 - Memory chunk to fill.
+# a1 - Constant byte.
+# a2 - Memory size.
+.type memset, @function
+.globl memset
+memset:
 	mv t0, a0
 
-.Lbzero_loop:
-	bgt t0, a1, .Lbzero_end
-	sb zero, (t0)
+.Lmemset_loop:
+	beqz a2, .Lmemset_end
+
+	sb a1, (t0)
+
+	addi a2, a2, -1
 	addi t0, t0, 1
 
-.Lbzero_end:
+	j .Lmemset_loop
+
+.Lmemset_end:
 	ret

arch/riscv/kernel.ld

@@ -19,14 +19,14 @@ SECTIONS {
 	} :text
 
 	. = ALIGN(CONSTANT(MAXPAGESIZE));
-	PROVIDE(etext = . - _kernel_offset);
+	PROVIDE(etext = .);
 
 	.rodata : AT(ADDR(.rodata) - _kernel_offset) ALIGN(16) {
 		*(.rodata .rodata.*);
 	}
 
 	. = ALIGN(CONSTANT(MAXPAGESIZE));
-	PROVIDE(_data = . - _kernel_offset);
+	PROVIDE(_data = .);
 
 	.data : AT(ADDR(.data) - _kernel_offset) ALIGN(16) {
 		*(.data .data.*);
@@ -38,7 +38,7 @@ SECTIONS {
 	}
 
 	. = ALIGN(CONSTANT(MAXPAGESIZE));
-	PROVIDE(end = . - _kernel_offset);
+	PROVIDE(end = .);
 
 	/DISCARD/ : {
 		*(.eh_frame)

arch/riscv/kernel.s

@@ -1,3 +1,7 @@
+# This Source Code Form is subject to the terms of the Mozilla Public License,
+# 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/.
+
 .section .rodata
 
 panic_message: .asciz "\nPanic"
@@ -6,29 +10,41 @@ panic_message: .asciz "\nPanic"
 .section .text.boot
 
 .type _start, @function
-.global _start
+.globl _start
 _start:
+	# Save the physical start of the text section.
+	# This must be the first instruction.
+	auipc s2, 0
 	mv s1, a1 # Save the device tree pointer.
-	auipc s2, 0 # Save the physical start of the text section.
+	# Virtual address of the start of the text section.
+	lui s3, %hi(_start)
+	addi s3, s3, %lo(_start)
+	sub s4, s3, s2 # Save the offset to the virtual memory address.
 
 	# Set the stack pointer.
-	lui sp, %hi(end)
-	addi sp, sp, %lo(end)
-	li t0, 64 * 1024
-	add sp, sp, t0
+	mv a0, s4
+	call set_up_stack
 
+	# Create identity mapping for the kernel and copy it to higher half.
 	mv a0, s2
 	mv a1, sp
+	call get_page_tables_size
+	mv a3, a0
+	mv a0, s2
+	mv a1, sp
+	mv a2, s4
 	jal imem_initialize
+	mv a0, s3
 	jal imem_mirror
 
-	la a0, __bss_start
-	la a1, end
-	call bzero
+	call set_up_trap_routine
+	add sp, sp, s4 # Move the stack to the virtual memory.
 
-	lui t0, %hi(trap_routine)
-	addi t0, t0, %lo(trap_routine)
-	csrw stvec, t0
+	# Zero BSS and kernel stack.
+	lui a0, %hi(__bss_start)
+	addi a0, a0, %lo(__bss_start)
+	mv a1, sp
+	call imem_zero
 
 	mv a0, sp
 	mv a1, s1
@@ -36,7 +52,74 @@ _start:
 	lui t0, %hi(kernel_main)
 	jr t0, %lo(kernel_main)
 
+# Given the start and end addresses of the kernel, calculates how
+# much additional storage should be mapped to store the page tables.
+
+# The original formula for getting the number of required pages is:
+# ((sp - start) / 4096 + 1 + 1023) / 1024.
+#
+# - sp is the stack pointer (kernel end).
+# - The kernel size is divided by the page size. It is the kernel size in pages.
+# - One additional page is added for the root page table.
+# - 1023 is added so the next division is rounded up and not down.
+#
+# A simplified formula is used in the actual implementation.
+#
+# Parameters:
+# a0 - Kernel start address.
+# a1 - Kernel end address.
+#
+# Returns the result in a0.
+.type get_page_tables_size, @function
+get_page_tables_size:
+	sub a0, a1, a0
+
+	srli a0, a0, 22
+	addi a0, a0, 1
+	slli a0, a0, 12
+
+	ret
+
+# Set an exception handler.
+.type set_up_trap_routine, @function
+set_up_trap_routine:
+	lui t0, %hi(trap_routine)
+	addi t0, t0, %lo(trap_routine)
+	csrw stvec, t0
+	ret
+
+.type set_up_stack, @function
+set_up_stack:
+	lui sp, %hi(end)
+	addi sp, sp, %lo(end)
+	sub sp, sp, a0
+	li t0, 64 * 1024
+	add sp, sp, t0
+	ret
+
+# Zeroes a chunk of memory.
+#
+# Parameters:
+# a0 - Memory start pointer.
+# 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
+
+	j .Limem_zero_loop
+
+.Limem_zero_end:
+	ret
+
 # Copies the identity mappings to the higher half.
+#
+# Parameters:
+# a0 - Start of the kernel virtual space.
 .type imem_mirror, @function
 imem_mirror:
 	csrr t0, satp
@@ -57,8 +140,7 @@ imem_mirror:
 	blt t1, t2, .Limem_mirror_copy
 
 	# Persist the copied page table in the root page table.
-	lui t4, %hi(_start)
-	srli t4, t4, 20 # VPN[1] * PTESIZE.
+	srli t4, a0, 20 # VPN[1] * PTESIZE.
 	add t4, t0, t4 # a + VPN[1] * PTESIZE.
 	srli t3, t2, 2
 	ori t3, t3, 0b1
@@ -72,6 +154,8 @@ imem_mirror:
 # Parameters:
 # a0 - Physical start of the text section.
 # a1 - Page aligned address following the kernel (kernel end).
+# a2 - Offset from the physical to virtual memory.
+# a3 - Size of additional page tables.
 .type imem_initialize, @function
 imem_initialize:
 	# Prologue.
@@ -80,16 +164,26 @@ imem_initialize:
 	sw s0, 24(sp)
 	addi s0, sp, 32
 
-	sw a1, 20(sp)
 	sw s1, 16(sp)
 	sw s2, 12(sp)
 	sw s3, 8(sp)
 
+	# Set s1 to the beginning of free memory (skipping root page table).
 	li s1, 4096
 	add s1, a1, s1
+
+	slli a3, a3, 1 # Multiply by two to account for the copy in virtual memory.
+	# s1 is a1 + PAGE, where PAGE is a free page reserved for future use,
+	# since the size of the root page table is already included in a3.
+	# s1 is used to avoid calculating a1 + PAGE_SIZE twice.
+	add a3, s1, a3
+	sw a3, 20(sp)
+
 	# Set boundaries between segments.
 	lui s2, %hi(etext)
+	sub s2, s2, a2
 	lui s3, %hi(_data)
+	sub s3, s3, a2
 
 	# Set the root page since it is used by mapping functions,
 	# but don't activate paging before the segments are property mapped.
@@ -115,8 +209,6 @@ imem_initialize:
 	# Data section.
 	addi a0, s3, %lo(_data)
 	lw a1, 20(sp)
-	li t0, 4096 * 3
-	add a1, a1, t0
 	mv a2, s1
 	li a3, 1
 	li a4, 0b0110