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

.include "./arch/riscv/constants.s"

.section .rodata

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

.section .text.boot

.type _start, @function
.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.
	# 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.
	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 a1, s2
	mv a2, sp
	mv a3, s4
	jal imem_initialize
	mv a1, s3
	jal imem_mirror

	# Save the address where the initial page tables end in a variable.
	add a0, a0, s3
	la t0, pmem_address
	sw a0, (t0)

	call set_up_trap_routine
	add sp, sp, s4 # Move the stack to the virtual memory.

	# Zero BSS and kernel stack.
	lui a0, %hi(__bss_start)
	addi a0, a0, %lo(__bss_start)
	mv a1, sp
	call imem_zero

	mv a0, s1

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

# 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:
	li t0, 0xffc00000
	and a0, a0, t0
	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:
.Limem_zero_loop:
	bgeu a0, a1, .Limem_zero_end
	sw zero, (a0)
	addi a0, a0, 4

	j .Limem_zero_loop

.Limem_zero_end:
	ret

# Copies the identity mappings to the higher half.
#
# Parameters:
# a0 - Size of additional page tables.
# a1 - Start of the kernel virtual space.
#
# Returns a pointer pointing to the end of the last page table
# with higher half mappings in a0.
.type imem_mirror, @function
imem_mirror:
	csrr t0, satp
	slli t0, t0, 12

	li t3, PAGE_SIZE
	add t1, t0, t3 # Original page table.
	add t2, t1, a0 # Copied page table.

	li t5, 0xffffffcf
	mv t3, t2

.Limem_mirror_copy:
	lw t4, (t1)
	and t4, t4, t5 # Clear "dirty" and "accessed" bits.
	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.
	srli t4, a1, 20 # VPN[1] * PTESIZE.
	add t4, t0, t4 # a + VPN[1] * PTESIZE.
	srli t5, t2, 2
	ori t5, t5, 0b1
	sw t5, (t4)

	mv a0, t3

	ret

# Activates paging and creates root page table.
# Allocates enough pages for the kernel and its stack.
#
# Parameters:
# a0 - Size of additional page tables.
# a1 - Physical start of the text section.
# a2 - Page aligned address following the kernel (kernel end).
# a3 - Offset from the physical to virtual memory.
#
# Returns the size of the created page tables (a0) in a0.
.type imem_initialize, @function
imem_initialize:
	# Prologue.
	addi sp, sp, -32
	sw ra, 28(sp)
	sw s0, 24(sp)
	addi s0, sp, 32

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

	# Set s1 to the beginning of free memory (skipping root page table).
	li s1, PAGE_SIZE
	add s1, a2, s1

	# Set boundaries between segments.
	lui s2, %hi(etext)
	sub s2, s2, a3
	lui s3, %hi(_data)
	sub s3, s3, a3

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

	# Zero memory for page tables.
	mv a0, a2
	add a1, s1, a0
	call imem_zero

	# Executable section.
	mv a0, s1
	lw a1, 4(sp)
	addi a2, s2, %lo(etext)
	li a3, PAGE_READ | PAGE_EXECUTE
	jal imem_map_range

	# Read only section.
	addi a1, s2, %lo(etext)
	addi a2, s3, %lo(_data)
	li a3, PAGE_READ
	jal imem_map_range

	# Data section.
	addi a1, s3, %lo(_data)
	lw a2, 20(sp)
	slli a2, a2, 1 # Multiply by two to account for the copy in virtual memory.
	# s1 is a2 + PAGE, where PAGE is a free page reserved for future use,
	# since the size of the root page table is already included in a0.
	# s1 is used to avoid calculating a2 + PAGE_SIZE twice.
	add a2, s1, a2
	li a3, PAGE_READ | PAGE_WRITE
	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)
	lw a0, 20(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 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.
	addi sp, sp, -32
	sw ra, 28(sp)
	sw s0, 24(sp)
	addi s0, sp, 32

	sw s1, 20(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 s1, s2, .Limem_map_range_end

	mv a1, s1
	mv a2, s4
	jal imem_map_at

	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.
	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 - Free memory page that can be used to create a leaf page table.
# a1 - Physical page address.
# a2 - Flags.
#
# 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 t1, 0xffc # 10 bit * PTESIZE mask.

	srli t2, a1, 20
	and t2, t2, t1 # VPN[1] * PTESIZE.
	add t2, t0, t2 # a + VPN[1] * PTESIZE.

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

	beqz t4, .Limem_map_at_create

	li t4, 0xfffffc00
	and t0, t3, t4
	slli t0, t0, 2

	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, a0
	srli t3, t0, 2
	ori t3, t3, 0b1
	sw t3, (t2)

	li t4, PAGE_SIZE
	add a0, a0, t4

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

	srli t3, a1, 2 # Physical page number mapped to 12 bits of the page table entry.
	ori a2, a2, 0b1
	or t3, t3, a2 # Execute, write, read and valid.
	sw t3, (t2)
	ret

.section .text

.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