.global _entry, __bss, __bss_end, __stack_top

.section .rodata

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

.section .text

bzero:
	la t0, __bss
	la t1, __bss_end

.Lbzero_loop:
	bgt t0, t1, .Lbzero_end
	sw zero, (t0)
	addi t0, t0, 4

.Lbzero_end:
	ret

# Maps a page of physical memory to the identical virtual memory.
#
# Parameters:
# a0 - Physical page address.
# a1 - Free memory page that can be used to create a leaf page table.
# a2 - Page table level. If a2 is 0 a superpage is allocated, otherwise a normal page.
#
# 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 t1, 0xffc # 10 bit * PTESIZE mask.
	srli t4, a0, 2 # Physical page number.

	# Multiply VPN[1] by 4 and add to the start address of the page table.
	# This gives the physical address of the page table entry.
	srli t2, a0, 20
	and t2, t2, t1 # VPN[1] * PTESIZE.
	add t2, t0, t2 # a + VPN[1] * PTESIZE.

	beqz a2, .Limem_map_at_super

	srli t3, a1, 2
	ori t3, t3, 0b1
	sw t3, (t2)

	srli t2, a0, 10
	and t2, t2, t1 # VPN[0] * PTESIZE.
	add t2, a1, t2 # a + VPN[0] * PTESIZE.

	j .Limem_map_at_leaf

.Limem_map_at_super:
	# Calculate PPN[1] (10 most significant bits of the physical address
	# and map it to 12 bits of the table page entry.
	li t1, 0xfff00000
	and t4, t4, t1

.Limem_map_at_leaf:
	ori t4, t4, 0b1111 # Execute, write, read and valid.
	sw t4, (t2)
	ret

# Activates paging and creates root page table.
# Allocates enough pages for the kernel and its stack.
#
# Parameters:
#
# a0 - Start (base) address of the kernel.
# 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

	li t0, 64 * 1024
	add t0, a1, t0
	sw t0, 20(sp)
	sw a1, 16(sp)

	srli t0, a1, 12
	csrw satp, t0

.Limem_initialize_loop:
	lw t0, 20(sp)
	lw a1, 16(sp)

	bge a0, t0, .Limem_initialize_activate

	li t0, 4096
	add a1, a1, t0
	li a2, 1
	call imem_map_at

	li t0, 4096
	add a0, a0, t0

	j .Limem_initialize_loop

.Limem_initialize_activate:
	csrr t0, satp
	li t1, 0x80000000
	or t0, t0, t1

	sfence.vma
	csrw satp, t0
	sfence.vma

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

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

	la t0, kernel_entry
	csrw stvec, t0

	# Map flat device tree to the virtual memory.
	lw a0, 16(sp)
	lw a1, 20(sp)
	li a2, 0
	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 kernel_entry, @function
kernel_entry:
	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

.section .text.boot
_entry:
	# Set the stack pointer.
	la t0, end
	li t1, 64 * 1024
	add sp, t0, t1

	mv s1, a1

	la a0, .text.boot
	mv a1, sp
	call imem_initialize

	lw t0, (end)
	la t1, next_paddr
	sw t0, (t1)

	call bzero
	mv a0, sp
	mv a1, s1
	j kernel_main