# Custom first-fit malloc routines by Lameguy64
# Part of the PSn00bSDK Project

.set noreorder

.set ND_PREV,	0	# Address to previous block (NULL if starting block)
.set ND_NEXT,	4	# Address to next block (NULL if end block)
.set ND_SIZE,	8	# Size of block
.set ND_HSIZ,	12

.section .text


# Stupid small function just to get bss end
# due to GCC insisting externs to be gp relative
.global GetBSSend
.type GetBSSend, @function
GetBSSend:
	la		$v0, _end
	jr		$ra
	nop
	

# Initializes the heap for malloc
#   a0 - Starting address of heap
#   a1 - Size of memory heap
#
.global InitHeap
.type InitHeap, @function
InitHeap:
	la		$v0, _malloc_addr
	sw		$a0, 0($v0)
	la		$v0, _malloc_size
	sw		$a1, 0($v0)
	
	sw		$0 , ND_PREV($a0)		# Set heap header
	sw		$0 , ND_NEXT($a0)
	jr		$ra
	sw		$0 , ND_SIZE($a0)
	

# Allocates a block of memory in the heap
#   a0 - Size of memory block to allocate.
#
.global malloc
.type malloc, @function
malloc:
	addiu	$a0, 3					# Round size to a multiple of 4
	srl		$a0, 2
	
	la		$a2, _malloc_addr
	lw		$a2, 0($a2)
	sll		$a0, 2
	
.find_next:
	
	move	$a1, $a2
	
	lw		$a2, ND_NEXT($a1)		# Get block header
	lw		$v1, ND_SIZE($a1)
	
	subu	$v0, $a2, $a1			# Compute space between current and next
	
	beqz	$v1, .empty_block		# Occupy empty block (if size = 0)
	nop
	
	beqz	$a2, .new_block			# Allocate a new block (if no next)
	nop
	
	addiu	$v0, -(ND_HSIZ*2)		# Compute remaining space of block
	subu	$v0, $v1
	
	blt		$v0, $a0, .find_next	# Search for the next block if space is not big enough
	nop
	
	# Perform a block split using remaining space of current block

	addiu	$v0, $a1, ND_HSIZ		# Compute address for new header
	addu	$v0, $v1
	
	sw		$a1, ND_PREV($v0)		# Set the new block header
	sw		$a2, ND_NEXT($v0)
	sw		$a0, ND_SIZE($v0)
	
	sw		$v0, ND_NEXT($a1)		# Update previous and next blocks
	sw		$v0, ND_PREV($a2)
	
	jr		$ra
	addiu	$v0, ND_HSIZ
	
.empty_block:	# Occupy an empty block

	beqz	$a2, .no_next			# Skip size calculation if there's no next
	nop
	
	addiu	$v0, -ND_HSIZ
	blt		$v0, $a0, .find_next
	nop
	
	b		.skip_space_check
	nop
	
.no_next:

	la		$v1, _malloc_addr		# Check if there's enough space for a block
	lw		$v1, 0($v1)
	la		$v0, _malloc_size
	lw		$v0, 0($v0)
	
	subu	$v1, $a1, $v1
	addu	$v1, $a0
	addiu	$v1, ND_HSIZ
	
	bgt		$v1, $v0, .no_space
	nop
	
.skip_space_check:

	sw		$a0, ND_SIZE($a1)
	jr		$ra						# Return address
	addiu	$v0, $a1, ND_HSIZ
	
.new_block:		# Create a new block

	addu	$a2, $a1, $v1			# Compute address for new block
	addiu	$a2, ND_HSIZ
	
	la		$v1, _malloc_addr
	lw		$v1, 0($v1)
	la		$v0, _malloc_size
	lw		$v0, 0($v0)
	
	subu	$v1, $a2, $v1
	addu	$v1, $a0
	addiu	$v1, ND_HSIZ
	
	bgt		$v1, $v0, .no_space		# Reject if it exceeds specified size
	nop
	
	sw		$a1, ND_PREV($a2)
	sw		$0 , ND_NEXT($a2)
	sw		$a0, ND_SIZE($a2)
	
	sw		$a2, ND_NEXT($a1)
	
	jr		$ra						# Return address
	addiu	$v0, $a2, ND_HSIZ
	
.no_space:		# Return a null if no space can be found
	jr		$ra
	move	$v0, $0

	
# Allocates a block of memory in block units and zero fills the
# allocated block.
#   a0 - Block size.
#	a1 - Number of blocks to allocate
#
.global calloc
.type calloc, @function
calloc:
	mult	$a0, $a1
	addiu	$sp, -4
	sw		$ra, 0($sp)
	
	jal		malloc
	mflo	$a0
	
	move	$a0, $v0
	mflo	$a1
.clear_loop:
	sw		$0 , 0($a0)
	addi	$a1, 4
	bgtz	$a1, .clear_loop
	addiu	$a0, 4
	
	lw		$ra, 0($sp)
	addiu	$sp, 4
	jr		$ra
	nop

	
# Deallocates an allocated block
#	a0 - An address returned by malloc to deallocate
#
.global free
.type free, @function
free:

	addiu	$a0, -ND_HSIZ
	lw		$a1, ND_PREV($a0)
	lw		$a2, ND_NEXT($a0)
	
	beqz	$a1, .is_start			# Check if block is a starting block
	nop
	
	beqz	$a2, .is_end
	nop

	# Unlink
	
	sw		$a2, ND_NEXT($a1)
	jr		$ra
	sw		$a1, ND_PREV($a2)
	
.is_end:		# Unlinks the ending block
	jr		$ra
	sw		$0 , ND_NEXT($a1)
.is_start:		# Simply set size to 0 if starting block
	jr		$ra
	sw		$0 , ND_SIZE($a0)


# Internal variables
.comm _malloc_addr, 4, 4
.comm _malloc_size, 4, 4