# Alloc and Dealloc memory as requested
#
# Programs using these routines will ask
# for a certain size of memory. We actually
# use more than that size for metadata,
# but we put it at the beginning, before the
# pointer we hand back. We add a size field and
# and AVAILABLE/UNAVAILABLE marker.
#
# The whole memory slot looks like this:
#
####################################################
# Available marker | Size | Actual memory location #
####################################################
#			    ^-- returned pointer points
#				here
#
# The calling program won't see our metadata.

### GLOBAL VARS ###
.section .data

############################### DEBUG INFO ################################

# Originally, these labels heap_begin and current_break are .long long...
# This implicitly means 4 bytes for each variable. This was causing the instructions
# in allocate_init to overlap the variables when writing to them:

# The following instructions were the problem:
#	movq %rax, current_break
#			allocated at 0x402011
#	movq %rax, heap_begin
#			allocated at 0x40200d
#
# Static allocation grows up in memory... But, when writing %rax to heap_begin,
# due the variable sizes being 4bytes, the instruction was corrupting
# current_break variable, causing the program to crash later in the allocation
# loop, because the heap_begin allocation, consequently caused current_break to be
# zero causing a NULL ptr dereference (or simply a SEGFAULT):

# =>   0x000000000040114e <+5>:	mov    0x8(%rax),%rdx
# Program received signal SIGSEGV, Segmentation fault.

# allocate_init() disassemble:
#   0x0000000000401105 <+0>:	push   %rbp
#   0x0000000000401106 <+1>:	mov    %rsp,%rbp
#   0x0000000000401109 <+4>:	mov    $0xc,%rax
#   0x0000000000401110 <+11>:	mov    $0x0,%rdi
#   0x0000000000401117 <+18>:	syscall
#   0x0000000000401119 <+20>:	inc    %rax
#   0x000000000040111c <+23>:	mov    %rax,0x402011
#   0x0000000000401124 <+31>:	mov    %rax,0x40200d
#=> 0x000000000040112c <+39>:	mov    %rbp,%rsp
#   0x000000000040112f <+42>:	pop    %rbp
#   0x0000000000401130 <+43>:	ret
#
# (gdb) info register rax
# rax            0x403001            4206593
# (gdb) p /x *0x40200d
# $5 = 0x403001
# (gdb) p /x *0x402011
# $6 = 0x0		<-- Here, the address of current_break got zeroed after
#			    heap_being has been changed

############################# END OF DEBUG INFO #############################

# Points to the beginning of the memory we are managing
heap_begin:
	.quad 0

# Points to one locaiton past the memory we are managing
current_break:
	.quad 0

### HEADER STRUCTURE INFORMATION ###

# To make things simpler, we use one word for
# each field in the header

.equ HEADER_SIZE, 16	 # size of space for memory region header
.equ HDR_AVAIL_OFFSET, 0 # Location of the 'available' flag in the header
.equ HDR_SIZE_OFFSET, 8	 # Location of the size field in the header


### CONSTANTS ###
.equ UNAVAILABLE, 0
.equ AVAILABLE, 1

.equ SYS_BRK, 12	# syscall number for brk() syscall in x86_64


.section .text

### FUNCTIONS ###

## allocate_init ##
#
# Call this function to initialize the functions by setting heap_begin and
# current_break. No parameters and no return value

.globl allocate_init
.type allocate_init, @function

allocate_init:
	pushq %rbp		# standard function stuff
	movq  %rsp, %rbp

	# If brk() syscall is called with a 0 in %rdi, it returns the last valid
	# usable address
	movq $SYS_BRK, %rax
	movq $0, %rdi
	syscall

	incq %rax		 # brk(0) returns the current break in %rax, we want the
				 # value after that

	movq %rax, current_break # Store the current break (actually the address
				 # after that)

	movq %rax, heap_begin	 # Our heap starts where the break is now. First
				 # address of uninitialized memory. This will
				 # cause the allocate function to get more memory
				 # from Linux the first time it is run.

	movq %rbp, %rsp		# Exit the function
	popq %rbp
	ret

## END of allocate_init ##


## allocate ##
#
# Grab a section of memory.
# - Checks to see if there are any free blocks
# - If not, asks Linux for more memory for the
#   heap through brk() syscall
#
# - Receives on parameter, the size of the memory block
#   we want to allocate
#
# - Returns the address of the allocated memory in %rax, or 0
#   if there is no memory available on the system
#
# ### Process
#
# %rcx - holds the size of requested memory (first/only parameter)
# %rax - current memory region being examined
# %rbx - Memory address past the end of the heap
# %rdx - size of the current memory region
#
# We scan through each memory region starting with heap_begin. We look at the
# size of each one and if it has been allocated. If it's big enough for the
# requested size, and it's available, we grab that one. If we do not find a
# region large enough, we ask Linux for more memory, which in case, moves the
# current_break up.

.globl allocate
.type allocate, @function
.equ ST_MEM_SIZE, 16	# Stack position of the memory size to allocate

allocate:
	pushq %rbp
	movq %rsp, %rbp

	movq ST_MEM_SIZE(%rbp), %rcx	# %rcx now holds the memory size we are
				# looking for (first/only parameter)

	movq heap_begin, %rax	# %rax will hold the current search location
	movq current_break, %rbx

	alloc_loop_begin:		# Scan memory regions

		cmpq %rbx, %rax		# heap needs more memory if these are
					# equal. %rax will hold the end of the
					# next memory region at each iteraction
					# until it hits the current_break.

		je move_break

		# Retrieve the size of this mem slot
		movq HDR_SIZE_OFFSET(%rax), %rdx

		cmpq $UNAVAILABLE, HDR_AVAIL_OFFSET(%rax) # If the space is
							  # unavailable, i.e.
							  # already in use.. Go
		je next_location			  # to the next one.

		cmpq %rdx, %rcx		# If the slot is available, check if
		jle allocate_here	# it's big enough.

	next_location:
		addq $HEADER_SIZE, %rax	# Total size of the memory region is the
					# sum of the size requested (currently
					# at %rdx), plus the 16 bytes for the
		addq %rdx, %rax		# header.

		jmp alloc_loop_begin	# Go look at the next location

	allocate_here:			# If we've made it here, that means that
					# the region header of the region to
					# allocate is in %rax

		# Mark space as unavailable
		movq $UNAVAILABLE, HDR_AVAIL_OFFSET(%rax)
		addq $HEADER_SIZE, %rax		# move %rax past the header, so it
						# points to the usable memory. Such
						# address is returned to the user.

		# Normal function return
		movq %rbp, %rsp
		popq %rbp
		ret

	move_break:	# We have exhausted all addressable memory, we need to
			# get more from Linux
			# %rbx holds the current endpoint of the data, and
			# %rcx its size.

		addq $HEADER_SIZE, %rbx	# We need to increase %rbx to where we _want_
					# memory to end. So we account for the header
		addq %rcx, %rbx		# and the user's requested size

		# Ask Linux for more memory

		# We'll need the values in these registers, so save them...
		pushq %rax
		pushq %rbx
		pushq %rcx

		movq $SYS_BRK, %rax	# Set new break by calling brk()
		movq %rbx, %rdi		# x86_64 uses %rdi for first parameter
		syscall

		# brk() returns 0 for error or the address we asked for (or larger if it
		# needs to be rounded up). We don't really care here where it ends up
		# setting the break as long as it isn't 0.

		cmpq $0, %rax		# Check for errors
		je error

		# Restore our registers
		popq %rcx
		popq %rbx
		popq %rax

		# Set this memory as unavailable since we are about to give it away
		movq $UNAVAILABLE, HDR_AVAIL_OFFSET(%rax)
		movq %rcx, HDR_SIZE_OFFSET(%rax) # Set memory size

		addq $HEADER_SIZE, %rax		# Set %rax to the address being
						# returned to the user (user
						# doesn't know anything about the headers

		movq %rbx, current_break	# Save the current break. In
						# reality it may be larger due
						# rounding, but we don't care
						# about memory footprint here

		movq %rbp, %rsp
		popq %rbp
		ret

	error:
		movq $0, %rax		# Return 0 on error
		movq %rbp, %rsp
		popq %rbp
		ret
##### END OF allocate ######

## deallocate ##
#
# Give back a region of memory to the pool after user is done
# with it
#
# Parameter: Address of the memory to be freed
# No return value
#
# The memory address returned to the user starts at 2 words beyond its header,
# all we need to do here is mark the memory slot as free (available).
# By now we don't care about moving the break back.
.globl deallocate
.type deallocate, @function

# Stack position of the memory region to be free (function's parameter)
.equ ST_MEMORY_SEG, 8	# We are not saving %rbp here so we just need to skip ret
			# address

deallocate:
	# Function is too simple, no need for fancy function stuff

	# Get the address of the memory to free (Normally this is 16(%rbp), but
	# since we didn't push %rbp or move %rsp to %rbp, we can just do 8(%rsp)
	movq ST_MEMORY_SEG(%rsp), %rax

	# Get the pointer to the beginning of the region (i.e. to the header)
	subq $HEADER_SIZE, %rax

	# Mark it as available
	movq $AVAILABLE, HDR_AVAIL_OFFSET(%rax)
	ret
##### END OF deallocate #####