# Copyright (c) 2024 Freya Murphy

.data
.align 2

heap_start:
	.space 4
heap_len:
	.space 4
heap_free:
	.space 4
heap_ptr:
	.space 4

null:
	.space 4

.text
.align 2
.globl main

# init the heap
heap_init:
	# sbrk(0)
	li	$v0, 9
	li	$a0, 0
	syscall
	sw	$v0, heap_start
	sw	$v1, heap_len

	# heap is empty, ptr at start
	sw	$v0, heap_ptr
	# heap is empty, size = len
	sw	$v1, heap_free

	jr	$ra



# $a0 : amount of bytes to increase
heap_increase:
	# t0 : OLD heap len + 1
	# t1 : heap ptr
	# t2 : heap end
	# t2 : heap free
	# t3 : a0

	# save current length
	lw	$t0, heap_len

	# save a0
	move	$t3, $a0

	# sbrk(amt)
	li	$v0, 9
	syscall
	# sbrk(0)
	li	$v0, 9
	li	$a0, 0
	syscall
	sw	$v0, heap_start
	sw	$v1, heap_len

	lw	$t1, heap_ptr    # heap ptr
	add	$t2, $v0, $v1   # heap end
	sub	$t3, $t2, $t1   # heap free
	sw	$t3, heap_free

	# set return code to 1 if fail to allocate
	# i.e. $v1 = $t0 or ($v1 < $t0 + 1)
	# i.e. heap size hasen't changed
	addi	$t0, $t0, 1
	slt	$v0, $v1, $t0
	jr	$ra



# $a0 : amount of bytes to allocate
malloc:
	# t0 : heap free + 1
	# t1 : if enough mem is free
	# t2 : alloc size [(a0 align 4) + 4096 + 4]
	# t3 : temp for heap_ptr
	# t7 : alloc amt (a0 align 4) + 4
	# t8 : temp for modulo

	# save $a0
	move	$t7, $a0
	addi	$t7, $t7, 4 # add 4 bytes to save INTERAL ptr size

	# align $t7 by 4
	# allows us to use lw and sw in realloc
	li	$t8, 4
	div	$t8, $t7, $t8
	mfhi	$t8
	beq	$t8, $zero, malloc_aligned
	addi	$t7, $t7, 4
	sub	$t7, $t7, $t8

malloc_aligned:
	# check if we have enough memory free
	lw	$t0, heap_free
	addi	$t0, $t0, 1
	slt	$t1, $t7, $t0
	bne	$t1, $zero, malloc_hasmem

	# set needed mem to alloc
	# page size + alloc request

	# save $ra and $t7
	addi	$sp, $sp, -8
	sw	$ra, 0($sp)
	sw	$t7, 4($sp)

	li	$t2, 4096
	add	$t2, $t2, $t7
	move	$a0, $t2
	jal	heap_increase

	# pop $ra and $t7
	lw	$ra, 0($sp)
	lw	$t7, 4($sp)
	addi	$sp, $sp, 8

	# check heap_increase return code
	beq	$v0, $zero, malloc_hasmem

malloc_error:
	# failed to sbrk, return null
	la	$v0, null
	jr	$ra

malloc_hasmem:
	# set return value, and save ptr size in it
	# add 4 to return ptr, since first 4 bytes are INTERNAL
	addi	$t3, $t7, -4
	lw	$v0, heap_ptr
	sw	$t3, ($v0)
	addi	$v0, $v0, 4

	lw	$t3, heap_ptr
	add	$t3, $t3, $t7 # increase ptr by alloc size
	sw	$t3, heap_ptr

	jr	$ra



# $a0 : address of ptr
free:
	# haha hehe hoho
	jr	$ra



# $a0 : new size
# $a1 : old ptr
realloc:
	# t0 : old ptr size
	# t2 : new ptr addr

	# check if $a0 is zero, if so then just free
	bne	$a0, $zero, realloc_inner
	move	$a0, $a1
	la	$v0, null
	j	free

	# save to stack
	addi	$sp, $sp, -12
	sw	$ra, 0($sp)
	sw	$a0, 4($sp)
	sw	$a1, 8($sp)

	jal	malloc

	# pop to stack
	lw	$ra, 0($sp)
	lw	$a0, 4($sp)
	lw	$a1, 8($sp)
	addi	$sp, $sp, 12

realloc_inner:
	addi	$a1, $a1, -4
	lw	$t0, ($a1)
	addi	$a1, $a1, 4

realloc_loop:
	# loop until $t0 or $t1 is zero
	beq	$t0, $zero, realloc_end
	beq	$a0, $zero, realloc_end

	addi	$a1, $a1, 4
	addi	$t0, $t0, -4
	addi	$a0, $a0, -4
	#lw	$t3,

	j	realloc_loop

realloc_end:


realloc_free:
	jr	$ra


main:
	# push return address
	addi	$sp, $sp, -4
	sw	$ra, ($sp)

	jal	heap_init

	li	$a0, 24
	jal	malloc

	move	$a0, $v0
	li	$v0, 1
	syscall

	li	$v0, 11
	li	$a0, 10
	syscall

	li	$a0, 24
	jal	malloc

	move	$a0, $v0
	li	$v0, 1
	syscall

	li	$v0, 11
	li	$a0, 10
	syscall

	# pop return address
	lw	$ra, ($sp)
	addi	$sp, $sp, 4

exit:
	# exit with code 0
	li	$v0, 0
	jr	$ra