.section .rodata

panic_message: .asciz "\nPanic"
.equ PANIC_MESSAGE_SIZE, . - panic_message

.section .text.boot

.type _start, @function
.global _start
_start:
	mv s1, a1 # Save the device tree pointer.
	auipc s2, 0 # Save the physical start of the text section.

	# Set the stack pointer.
	lui sp, %hi(end)
	addi sp, sp, %lo(end)
	li t0, 64 * 1024
	add sp, sp, t0

	mv a0, s2
	mv a1, sp
	jal imem_initialize
	jal imem_mirror

	la a0, __bss_start
	la a1, end
	call bzero

	lui t0, %hi(trap_routine)
	addi t0, t0, %lo(trap_routine)
	csrw stvec, t0

	mv a0, sp
	mv a1, s1

	lui t0, %hi(kernel_main)
	jr t0, %lo(kernel_main)

# Copies the identity mappings to the higher half.
.type imem_mirror, @function
imem_mirror:
	csrr t0, satp
	slli t0, t0, 12

	li t3, 0x1000
	add t1, t0, t3 # Original page table.
	add t2, t1, t3 # Copied page table.

	mv t3, t2

.Limem_mirror_copy:
	lw t4, (t1)
	sw t4, (t3)

	addi t3, t3, 4
	addi t1, t1, 4
	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.
	add t4, t0, t4 # a + VPN[1] * PTESIZE.
	srli t3, t2, 2
	ori t3, t3, 0b1
	sw t3, (t4)

	ret

# Activates paging and creates root page table.
# Allocates enough pages for the kernel and its stack.
#
# Parameters:
# a0 - Physical start of the text section.
# a1 - Page aligned address following the kernel (kernel end).
.type imem_initialize, @function
imem_initialize:
	# Prologue.
	addi sp, sp, -32
	sw ra, 28(sp)
	sw s0, 24(sp)
	addi s0, sp, 32

	sw a1, 20(sp)
	sw s1, 16(sp)
	sw s2, 12(sp)
	sw s3, 8(sp)

	li s1, 4096
	add s1, a1, s1
	# Set boundaries between segments.
	lui s2, %hi(etext)
	lui s3, %hi(_data)

	# Set the root page since it is used by mapping functions,
	# but don't activate paging before the segments are property mapped.
	srli t0, a1, 12
	csrw satp, t0

	# Executable section.
	# a0 is coming from the parameter.
	addi a1, s2, %lo(etext)
	mv a2, s1
	li a3, 1
	li a4, 0b1010
	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
	jal imem_map_range

	# 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
	jal imem_map_range

	# Enable paging.
	csrr t0, satp
	li t1, 0x80000000
	or t0, t0, t1

	sfence.vma
	csrw satp, t0
	sfence.vma

	lw s3, 8(sp)
	lw s2, 12(sp)
	lw s1, 16(sp)

	# Epilogue.
	lw ra, 28(sp)
	lw s0, 24(sp)
	addi sp, sp, 32
	ret

# 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.
.type imem_map_range, @function
imem_map_range:
	# Prologue.
	addi sp, sp, -32
	sw ra, 28(sp)
	sw s0, 24(sp)
	addi s0, sp, 32

	sw s1, 20(sp)
	sw a2, 16(sp)
	sw a3, 12(sp)
	sw a4, 8(sp)

	mv s1, a1

.Limem_map_range_loop:
	bge a0, s1, .Limem_map_range_end

	mv a1, a0
	lw a2, 16(sp)
	lw a3, 12(sp)
	lw a4, 8(sp)
	jal imem_map_at

	li t0, 4096
	add a0, a0, t0

	j .Limem_map_range_loop

.Limem_map_range_end:
	lw s1, 20(sp)

	# Epilogue.
	lw ra, 28(sp)
	lw s0, 24(sp)
	addi sp, sp, 32
	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.
#
# Returns virtual page address in 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

.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

	lw t3, (t2)
	andi t6, t3, 0b1

	beqz t6, .Limem_map_at_create

	and t0, t3, t4
	slli t0, t0, 2

	j .Limem_map_at_next

.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
	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

	j .Limem_map_at_loop

.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.
	sw t3, (t2)

	ret

.section .text

.type kernel_main, @function
kernel_main:
	# Prologue.
	addi sp, sp, -32
	sw ra, 28(sp)
	sw s0, 24(sp)
	addi s0, sp, 32

	sw a0, 20(sp)
	sw a1, 16(sp)

	# Map flat device tree to the virtual memory.
	lw a0, 16(sp)
	li t0, 0xffc00000
	and a0, a0, t0
	mv a1, a0
	lw a2, 20(sp)
	li a3, 0
	li a4, 0b0010
	call imem_map_at

	lw a0, 16(sp)
	call device_tree

	# Do nothing in a loop.
.Lkernel_main:
	j .Lkernel_main

	# Epilogue.
	lw ra, 28(sp)
	lw s0, 24(sp)
	addi sp, sp, 32
	ret

.balign 4
.type trap_routine, @function
trap_routine:
	csrw sscratch, sp
	addi sp, sp, -4 * 31

	sw ra, 0(sp)
	sw gp, 4(sp)
	sw tp, 8(sp)
	sw t0, 12(sp)
	sw t1, 16(sp)
	sw t2, 20(sp)
	sw t3, 24(sp)
	sw t4, 28(sp)
	sw t5, 32(sp)
	sw t6, 36(sp)
	sw a0, 40(sp)
	sw a1, 44(sp)
	sw a2, 48(sp)
	sw a3, 52(sp)
	sw a4, 56(sp)
	sw a5, 60(sp)
	sw a6, 64(sp)
	sw a7, 68(sp)
	sw s0, 72(sp)
	sw s1, 76(sp)
	sw s2, 80(sp)
	sw s3, 84(sp)
	sw s4, 88(sp)
	sw s5, 92(sp)
	sw s6, 96(sp)
	sw s7, 100(sp)
	sw s8, 104(sp)
	sw s9, 108(sp)
	sw s10, 112(sp)
	sw s11, 116(sp)

	csrr a0, sscratch
	sw a0, 120(sp)

	mv a0, sp
	call handle_trap

	lw ra, 0(sp)
	lw gp, 4(sp)
	lw tp, 8(sp)
	lw t0, 12(sp)
	lw t1, 16(sp)
	lw t2, 20(sp)
	lw t3, 24(sp)
	lw t4, 28(sp)
	lw t5, 32(sp)
	lw t6, 36(sp)
	lw a0, 40(sp)
	lw a1, 44(sp)
	lw a2, 48(sp)
	lw a3, 52(sp)
	lw a4, 56(sp)
	lw a5, 60(sp)
	lw a6, 64(sp)
	lw a7, 68(sp)
	lw s0, 72(sp)
	lw s1, 76(sp)
	lw s2, 80(sp)
	lw s3, 84(sp)
	lw s4, 88(sp)
	lw s5, 92(sp)
	lw s6, 96(sp)
	lw s7, 100(sp)
	lw s8, 104(sp)
	lw s9, 108(sp)
	lw s10, 112(sp)
	lw s11, 116(sp)
	lw sp, 120(sp)

	sret

.type handle_trap, @function
handle_trap:
	# Prologue.
	addi sp, sp, -32
	sw ra, 28(sp)
	sw s0, 24(sp)
	addi s0, sp, 32

	csrr t0, scause
	csrr t1, stval
	csrr t2, sepc

	li a1, PANIC_MESSAGE_SIZE
	la a0, panic_message
	call write_s

	li a0, ' '
	call write_c

	csrr a0, scause
	la a1, write_c
	call print_i

	call separator
	call write_c

	csrr a0, stval
	la a1, write_c
	call print_i

	call separator
	call write_c

	csrr a0, sepc
	la a1, write_c
	call print_i

	call panic

	# Epilogue.
	lw ra, 28(sp)
	lw s0, 24(sp)
	addi sp, sp, 32
	ret

.type panic, @function
panic:
	# Prologue.
	addi sp, sp, -32
	sw ra, 28(sp)
	sw s0, 24(sp)
	addi s0, sp, 32

	li a0, '\n'
	call write_c
.Lpanic:
	j .Lpanic

	# Epilogue.
	lw ra, 28(sp)
	lw s0, 24(sp)
	addi sp, sp, 32
	ret