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

.globl pmem_address
pmem_address: .word

.section .text

# a0 - First pointer.
# a1 - Second pointer.
# a2 - The length to compare.
#
# Returns 0 in a0 if memory regions are equal.
.type memcmp, @function
.globl memcmp
memcmp:
	mv t0, a0
	li a0, 0

.Lmemcmp_loop:
	beqz a2, .Lmemcmp_end

	lbu t1, (t0)
	lbu t2, (a1)
	sub a0, t1, t2

	bnez a0, .Lmemcmp_end

	addi t0, t0, 1
	addi a1, a1, 1
	addi a2, a2, -1

	j .Lmemcmp_loop

.Lmemcmp_end:
	ret

# Prints a string.
#
# Parameters:
# a0 - String pointer.
# a1 - Length of the string.
.type write_s, @function
.globl write_s
write_s:
	# 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)

	mv s1, a1
	mv s2, a0

.Lwrite_s_if:
	beqz s1, .Lwrite_s_end

	lw a0, (s2)
	call write_c

	addi s1, s1, -1
	addi s2, s2, 1

	j .Lwrite_s_if

.Lwrite_s_end:
	lw s1, 20(sp)
	lw s2, 16(sp)

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

# Prints a character from a0.
#
# Arguments:
# a0 - Character.
.type write_c, @function
.globl write_c
write_c:
	li a1, 0
	li a2, 0
	li a3, 0
	li a4, 0
	li a5, 0
	li a6, 0
	li a7, 1 # sbi_console_putchar.

	ecall
	ret

# Searches for the occurences of a character in the given memory block.
#
# Parameters:
# a0 - Memory block.
# a1 - Needle.
# a2 - Memory size.
#
# 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.

	lbu t0, (a0) # Load the character from the memory block.
	beq t0, a1, .Lmemchr_end # Exit if the character was found.

	# Otherwise, continue with the next character.
	addi a0, a0, 1
	addi a2, a2, -1

	j .Lmemchr_loop

.Lmemchr_nil:
	li a0, 0

.Lmemchr_end:
	ret

# Locates a substring.
#
# Parameters:
# a0 - Haystack.
# a1 - Haystack size.
# a2 - Needle.
# a3 - Needle size.
#
# 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
	sw ra, 20(sp)
	sw s0, 16(sp)
	addi s0, sp, 24

	# Save preserved registers. They are used to keep arguments.
	sw s1, 12(sp)
	sw s2, 8(sp)
	sw s3, 4(sp)
	sw s4, 0(sp)

	mv s1, a0
	mv s2, a1
	mv s3, a2
	mv s4, a3

.Lmemmem_loop:
	blt s2, s3, .Lmemmem_nil # Exit if the needle length is greater than memory.

	mv a0, s1
	mv a1, s3
	mv a2, s4
	call memcmp

	mv t0, a0 # memcmp result.
	mv a0, s1 # Memory pointer for the case the substring was found.
	beqz t0, .Lmemmem_end

	addi s1, s1, 1
	add s2, s2, -1

	j .Lmemmem_loop

.Lmemmem_nil:
	li a0, 0

.Lmemmem_end:

	# Restore the preserved registers.
	lw s1, 12(sp)
	lw s2, 8(sp)
	lw s3, 4(sp)
	lw s4, 0(sp)

	# Epilogue.
	lw ra, 20(sp)
	lw s0, 16(sp)
	add sp, sp, 24
	ret

# Copies memory.
#
# Parameters:
# a0 - Destination.
# a1 - Source.
# a2 - Size.
#
# Preserves a0.
.type memcpy, @function
.globl memcpy
memcpy:
	mv t0, a0

.Lmemcpy_loop:
	beqz a2, .Lmemcpy_end

	lbu t1, (a1)
	sb t1, (t0)

	addi t0, t0, 1
	addi a1, a1, 1
	addi a2, a2, -1

	j .Lmemcpy_loop

.Lmemcpy_end:
	ret

# Sets a memory region to the given byte value.
#
# Parameters:
# a0 - Memory chunk to fill.
# a1 - Constant byte.
# a2 - Memory size.
.type memset, @function
.globl memset
memset:
	mv t0, a0

.Lmemset_loop:
	beqz a2, .Lmemset_end

	sb a1, (t0)

	addi a2, a2, -1
	addi t0, t0, 1

	j .Lmemset_loop

.Lmemset_end:
	ret

# Maps a page of physical memory to the identical virtual memory.
#
# Parameters:
# a0 - Corresponding virtual address.
# a1 - Physical page address.
# a2 - Page table level. If a2 is 0 a superpage is allocated, otherwise a normal page.
# a3 - 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 a2, .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:
	la a7, pmem_address
	# 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.
	lw t0, (a7)
	srli t3, t0, 2
	ori t3, t3, 0b1
	sw t3, (t2)

	li t6, PAGE_SIZE
	add t0, t0, t6
	sw t0, (a7)

.Lpmem_map_at_next:
	# Calculate the next page table address for the next level.
	addi t1, t1, -10
	addi a2, a2, -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 a3, a3, 0b1
	or t3, t3, a3 # Execute, write, read and valid.
	sw t3, (t2)

	mv a0, a2
	ret