smallbrain.s

.equ EOF, -1
.equ O_RDONLY, 0
.equ EXIT_FAILURE, 1

// Opcodes
ENUM(OP_ADD)
ENUM(OP_SUB)
ENUM(OP_RIGHT)
ENUM(OP_LEFT)
ENUM(OP_LOOP_START)
ENUM(OP_LOOP_END)
ENUM(OP_READ)
ENUM(OP_WRITE)
ENUM(OP_ZERO)
ENUM(OP_COPY)

.global main

.data
zero_pattern:	.asciz "[-]"
read_mode:	.asciz "r"
die_fmt:	.asciz "%s: %s\n"
usage_fmt:	.asciz "Usage: %s script\n"
func_open:	.asciz "open"
func_fstat:	.asciz "fstat"
func_read:	.asciz "read"
func_malloc:	.asciz "malloc"
memory:		.zero 30000

.bss
bytecode:	.quad 0
program:	.quad 0

.text
// ==================
// Description:
//	The entry point of the program.
//
// Args:
//	%rdi: The number of command line arguments
//	%rsi: An array of command line arguments ((%rsi) is the program name)
//
// Return:
//	%rax: The programs exit code
// ==================
main:
	// Make sure that the right number of arguments were passed
	cmpl	$2, %edi
	jne	usage

	movq	8(%rsi), %rdi // Move the specified script filename into %rdi
	call	read_file     // Read the file
	call	compile       // Compile into bytecode
	call	execute       // Execute the program

	// Return successfully
	xorl	%eax, %eax
	ret

// ==================
// Description:
//	Read the specified file into a program buffer which the user must later free themselves.
//
// Args:
//	%rdi: The scripts filename
// ==================
read_file:
	pushq	%rbp
	movq	%rsp, %rbp

	// Allocate space for local variables
	// 	4 bytes for the file descriptor
	//	144 bytes for the struct stat (check `extras/sizeof_struct_stat.c` for numbers)
	//	12 bytes of padding to align %rsp on a 16 byte boundary
	subq	$160, %rsp

	// Open the file in read-only mode with the open(2) syscall
	movl	$O_RDONLY, %esi
	call	open

	cmpl	$-1, %eax      // Check if open(1) returned -1 (it failed)
	je	open_die       // If so, exit the program
	movl	%eax, -4(%rbp) // If not store the file descriptor

	movl	%eax, %edi       // Set the first argument to the file descriptor
	leaq	-160(%rbp), %rsi // Set the second argument to the address of the struct stat
	call	fstat            // Call fstat(2) to populate the struct stat

	// Error check fstat(2) just like open(2)
	cmpl	$-1, %eax
	je	fstat_die

	// Allocate a buffer for the programs contents
	movq	-112(%rbp), %rdi // Put st_size (the filesize) in %rdi
	incq	%rdi             // Make space for the NUL byte
	pushq	%rdi             // Store the value of %rdi for the next 2 malloc calls
	call	malloc           // Allocate the memory
	testq	%rax, %rax       // Check if %rax is NULL
	je	malloc_die       // If it is then malloc failed and we exit
	movq	%rax, (program)  // Store the address of the allocated memory in (program)
	addq	-112(%rbp), %rax // Point to the last element of the program array
	movb	$0, (%rax)       // NULL terminate the program

	// Allocate a buffer for the bytecode
	popq	%rdi             // Retrieve %rdi from the stack
	shlq	$4, %rdi         // Multiply the size of the buffer by 16 to make space for opcodes
	call	malloc           // Allocate the memory
	movq	%rax, (bytecode) // Store the address of the allocated memory in (bytecode)
	testq	%rax, %rax       // Check if %rax is NULL
	je	malloc_die       // If it is then malloc failed and we exit

	// Get a FILE* from the file descriptor with fdopen(3)
	movl	-4(%rbp), %edi
	movq	$read_mode, %rsi
	call	fdopen

	pushq	%rax            // Push the FILE* to the stack for the later call to fclose(3)
	movq	(program), %r15 // Move the program buffer into %r15

// cmpjeq - Compare Jump Equal Quadword
// ====================================
// Compare the value of 'val' to %rax and if they match jump to 'jump'.
.macro cmpjeq val, jump
	cmpq	\val, %rax
	je	\jump
.endm

// cmpjeb - Compare Jump Equal Byte
// ================================
// Compare the value of 'val' to %al and if they match jump to 'jump'.
.macro cmpjeb val, jump
	cmpb	\val, %al
	je	\jump
.endm

read_file_loop:
	movq	-168(%rbp), %rdi // Move the FILE* into %rdi
	call	fgetc            // Read a character from the file
	cmpl	$EOF, %eax       // Check if we reached EOF
	je	read_file_eof    // If we did then end the loop

	// If we match any non-comment character jump to read_file_not_comment
	cmpjeb	$'+', read_file_not_comment
	cmpjeb	$'-', read_file_not_comment
	cmpjeb	$'>', read_file_not_comment
	cmpjeb	$'<', read_file_not_comment
	cmpjeb	$'[', read_file_not_comment
	cmpjeb	$']', read_file_not_comment
	cmpjeb	$',', read_file_not_comment
	cmpjeb	$'.', read_file_not_comment
	// DEFAULT CASE (its a comment we can ignore)
	jmp	read_file_loop

read_file_not_comment:
	movb	%al, (%r15)    // Read the character into the program buffer
	incq	%r15           // Point to the next empty slot in the buffer
	jmp	read_file_loop // Loop again

read_file_eof:
	// NUL terminate the buffer
	movb	$0, (%r15)

	// Close the file, don't error check this
	movq	-168(%rbp), %rdi
	call	fclose

	leave
	ret

// ==================
// Description:
//	Compile the program into a bytecode which is an optimized version of the raw program. Each
//	opcode is a "struct" where the higher 8 bytes are an opcode and the lower 8 are data for
//	the instruction.
// ==================
compile:
	movq	(program), %r15  // Store the address of the program pointer into %r15
	movq	(bytecode), %r14 // Store the address of the bytecode pointer into %r14
compile_loop:
	// Load the current command into %rax
	movq	(%r15), %rax

	// Jump to a different label depending on which instruction we hit
	cmpjeb	$'+', compile_add
	cmpjeb	$'-', compile_sub
	cmpjeb	$'>', compile_right
	cmpjeb	$'<', compile_left
	cmpjeb	$'[', compile_loop_start
	cmpjeb	$']', compile_loop_end
	cmpjeb	$',', compile_read
	cmpjeb	$'.', compile_write

compile_add:
	movq	$OP_ADD, (%r14) // Specify the ADD opcode
	movq	$1, 8(%r14)     // Write the count of '+'s to the data portion
compile_add_loop:
	incq	%r15             // Move to the next instruction
	cmpb	$'+', (%r15)     // Check if there is another +
	jne	compile_out      // If not, exit this loop
	incq	8(%r14)          // Increment the accumulator
	jmp	compile_add_loop // Loop again

compile_sub:
	movq	$OP_SUB, (%r14) // Specify the SUB opcode
	movq	$1, 8(%r14)     // Write the count of '-'s to the data portion
compile_sub_loop:
	incq	%r15             // Move to the next instruction
	cmpb	$'-', (%r15)     // Check if there is another -
	jne	compile_out      // If not, exit this loop
	incq	8(%r14)          // Increment the accumulator
	jmp	compile_sub_loop // Loop again

compile_right:
	movq	$OP_RIGHT, (%r14) // Specify the RIGHT opcode
	movq	$1, 8(%r14)       // Write the count of '>'s to the data portion
compile_right_loop:
	incq	%r15               // Move to the next instruction
	cmpb	$'>', (%r15)       // Check if there is another >
	jne	compile_out        // If not, exit this loop
	incq	8(%r14)            // Increment the accumulator
	jmp	compile_right_loop // Loop again

compile_left:
	movq	$OP_LEFT, (%r14) // Specify the LEFT opcode
	movq	$1, 8(%r14)      // Write the count of '<'s to the data portion
compile_left_loop:
	incq	%r15              // Move to the next instruction
	cmpb	$'<', (%r15)      // Check if there is another <
	jne	compile_out       // If not, exit this loop
	incq	8(%r14)           // Increment the accumulator
	jmp	compile_left_loop // Loop again

compile_loop_start:
	// When we reach a '[' the first thing we want to do is check to see if it matches the
	// pattern '[-]'. This pattern is one that sets a memory cell to 0, so we can optimize that.
	movq	%r15, %rdi          // Compare the current position in the program string
	movq	$zero_pattern, %rsi // Compare it against the zero pattern '[-]'
	movl	$4, %ecx            // We want to compare 3 bytes (the instruction requires +1)
	repe	cmpsb               // Keep looping CMPSB while bytes match
	jrcxz	compile_zero        // Jump to compile_zero if the strings matched

	movq	%r15, %rdi                // Move the current instruction pointer into %rdi
	call	copy_loop_checker         // Call the copy loop checker
	testq	%rax, %rax                // Check to see if we hit a copy loop
	jz	compile_loop_start_normal // If we didn't, this is a regular loop
	movq	%rax, %r15                // Otherwise, set %r15 to the new location
	incq	%r15                      // Then point to the next instruction
	jmp	compile_out

compile_loop_start_normal:
	// Push the address of the loop start to the stack for the next ']'
	pushq	%r14

	movq	$OP_LOOP_START, (%r14) // Specify the LOOP_START opcode
	movq	$0, 8(%r14)            // Zero the data section
	incq	%r15                   // Increment the instruction pointer
	jmp	compile_out

compile_loop_end:
	popq	8(%r14)              // Pop the address of the previous loop start to the data section
	movq	$OP_LOOP_END, (%r14) // Push the address of the loop end to the stack
	incq	%r15                 // Increment the instruction pointer
	jmp	compile_out

compile_read:
	movq	$OP_READ, (%r14) // Specify the READ opcode
	movq	$0, 8(%r14)      // Zero the data section
	incq	%r15             // Increment the instruction pointer
	jmp	compile_out

compile_write:
	movq	$OP_WRITE, (%r14) // Specify the WRITE opcode
	movq	$0, 8(%r14)       // Zero the data section
	incq	%r15              // Increment the instruction pointer
	jmp	compile_out

compile_zero:
	movq	$OP_ZERO, (%r14) // Specify the ZERO opcode
	movq	$0, 8(%r14)      // Zero the data section
	addq	$3, %r15         // '[-]' is a 3 byte instruction
	// FALLTHROUGH

compile_out:
	addq	$16, %r14    // Move to the next opcode
	movb	(%r15), %al  // Move the current instruction into %al
	testb	%al, %al     // Check if we have reached the NUL byte
	jne	compile_loop // If we haven't, loop
	movq	$0, (%r14)   // Otherwise, NUL terminate the bytecode

	// Now that we have traversed the entire program, we do a 2nd pass backwards so that we can
	// set the jump addresses for the '[' commands now that the ']' commands have the addresses
	// set.
compile_backwards:
	// We are at the NUL terminator, so move backwards
	subq	$16, %r14

	cmpq	$OP_LOOP_END, (%r14)   // Check if we hit a ']'
	jne	compile_backwards_next // If we didn't move to the next check
	pushq	%r14                   // Otherwise push the address of the opcode to the stack
	jmp	compile_backwards_out

compile_backwards_next:
	cmpq	$OP_LOOP_START, (%r14) // Check if we hit a '['
	jne	compile_backwards_out  // If we didn't just keep looping
	popq	8(%r14)                // If we did then pop the corresponding ']' address

compile_backwards_out:
	cmpq	%r14, (bytecode)  // Check if we've seen every opcode
	jne	compile_backwards // If not keep looping
	movq	(program), %rdi   // Otherwise, move the program buffer to %rdi
	call	free              // Free it
	ret                       // And return

// ==================
// Description:
//	Try to figure out if we are at a copy loop and optimize it. A copy loop follows the pattern
//	of a loop ([]) beginning with a '-' followed by N occurances of '>' followed by a '+' and N
//	occurances of '<'. This sequence copies the current cell to the cell at offset N and clears
//	the current cell afterwards.
//
// Args:
//	%rdi: A pointer to the first '[' of the potential copy loop
//
// Return:
//	0 if not a copy loop, otherwise the new position of the instruction pointer.
// ==================
copy_loop_checker:
	// Skip '['
	incq	%rdi

	// All copy loops must begin with a '-'
	cmpb	$'-', (%rdi)
	jne	copy_loop_fail
	incq	%rdi

	// Zero %rax so we can use it to count the copy offset
	xorl	%eax, %eax
copy_loop_count_offset:
	cmpb	$'>', (%rdi)           // Check for '>'
	jne	copy_loop_next         // If we don't match anymore then move to the next step
	incq	%rax                   // Increment the offset counter
	incq	%rdi                   // Increment the instruction pointer
	jmp	copy_loop_count_offset // Loop again

copy_loop_next:
	cmpb	$'+', (%rdi)   // Check if we see the mandatory '+'
	jne	copy_loop_fail // If we don't then fail
	incq	%rdi           // Otherwise move to the next instruction

	// The following code is the exact same as what we just did to count the offset but we are
	// now using %rcx and decrementing for each '<'. This is so we can make sure that the copy
	// loop is a working one.
	cmpb	$'<', (%rdi)
	jne	copy_loop_fail
	movq	%rax, %rcx
copy_loop_verify_offset:
	cmpb	$'<', (%rdi)
	jne	copy_loop_next_2
	decq	%rcx
	incq	%rdi

copy_loop_next_2:
	cmpb	$']', (%rdi)   // Ensure this is the end of the loop
	jne	copy_loop_fail // If its not then fail
	testq	%rcx, %rcx     // Otherwise make sure that our offsets line up
	jnz	copy_loop_fail // If they don't then fail

	movq	$OP_COPY, (%r14)   // Create an OP_COPY opcode
	movq	%rax, 8(%r14)      // Set the offset to copy to
	movq	$OP_ZERO, 16(%r14) // Create an OP_ZERO opcode
	movq	$0, 24(%r14)       // Set an empty data section
	addq	$16, %r14          // Increment the opcode pointer

	// Return the address of the instruction pointer
	movq	%rdi, %rax
	ret
copy_loop_fail:
	xorl	%eax, %eax
	ret

// ==================
// Description:
//	Execute the brainfuck bytecode.
// ==================
execute:
	movq	(bytecode), %r15 // Store the address of the program pointer into %r15
	movq	$memory, %r14    // Store the address of the first memory cell into %r14 # TODO make sure this handles overflows normally
execute_loop:
	// Load the current command into %rax
	movq	(%r15), %rax

	// Jump to a different label depending on which instruction we hit
	cmpjeq	$OP_ADD, execute_add
	cmpjeq	$OP_SUB, execute_sub
	cmpjeq	$OP_RIGHT, execute_right
	cmpjeq	$OP_LEFT, execute_left
	cmpjeq	$OP_LOOP_START, execute_loop_start
	cmpjeq	$OP_LOOP_END, execute_loop_end
	cmpjeq	$OP_READ, execute_read
	cmpjeq	$OP_WRITE, execute_write
	cmpjeq	$OP_ZERO, execute_zero
	// OP_COPY
	jmp	execute_copy

execute_add:
	// Increment the current memory cell
	movq	8(%r15), %rax
	addb	%al, (%r14)
	jmp	execute_out

execute_sub:
	// Decrement the current memory cell
	movq	8(%r15), %rax
	subb	%al, (%r14)
	jmp	execute_out

execute_right:
	// Move the memory pointer right
	addq	8(%r15), %r14
	jmp	execute_out

execute_left:
	// Move the memory pointer left
	subq	8(%r15), %r14
	jmp	execute_out

execute_loop_start:
	// If the current memory cell is 0 move to the next ']'
	cmpb	$0, (%r14)
	cmovzq	8(%r15), %r15
	jmp	execute_out

execute_loop_end:
	// If the current memory cell is not 0 move to the last '['
	cmpb	$0, (%r14)
	cmovnzq	8(%r15), %r15
	jmp	execute_out

execute_read:
	// Set the current cell to the character read from stdin
	call	getchar          // Read a character with getchar(3)
	cmpb	$EOF, %al        // Check if the EOF was read
	je	execute_read_eof // If EOF was read, jump to a special handler for that
	movb	%al, (%r14)      // Otherwise move the read character into the current memory cell
	jmp	execute_out

execute_read_eof:
	// If EOF was read, set the current cell to 0
	movb	$0, (%r14)
	jmp	execute_out

execute_write:
	// Print the character at the current memory cell
	movl	(%r14), %edi // Move the current memory cell into %edi
	call	putchar      // Print it with putchar(3)
	jmp	execute_out

execute_zero:
	// Zero the current cell
	movb	$0, (%r14)
	jmp	execute_out

execute_copy:
	// Copy the current memory cells contents elsewhere
	movq	8(%r15), %rax
	movb	(%r14), %cl
	movb	%cl, (%rax, %r14, 1)
	// FALLTHROUGH

execute_out:
	addq	$16 , %r15   // Increment the instruction pointer
	movq	(%r15), %rax // Move the current instruction into %rax
	testq	%rax, %rax   // Check if we have reached the NUL byte
	jne	execute_loop // If we haven't, loop
	ret                  // Otherwise, return

// ==================
// Description:
//	The following die functions all work the same. They simply take the name of the
//	corresponding function and store it in %rdi. Then the die function is called to print a
//	message to stderr and terminate the program.
// ==================

.macro fdie s
	movq	\s, %rdi
	jmp	die
.endm

open_die:	fdie	$func_open
fstat_die:	fdie	$func_fstat
read_die:	fdie	$func_read
malloc_die:	fdie	$func_malloc


// ==================
// Description:
//	Print out an error message in the format "<func name>: <err msg>" then exit via `_exit`
//
// Args:
//	%rdi: The function name
// ==================
die:
	// Store the function name temporarily in %r15
	movq	%rdi, %r15

	call	__errno_location // Call __errno_location to get a pointer to errno
	movq	(%rax), %rdi     // Move errno into %rdi
	call	strerror         // Get the error string with strerror(3)

	movq	%rax, %rcx     // Set the error string
	xorl	%eax, %eax     // Zero %rax
	movq	%r15, %rdx     // Set the function name
	movq	$die_fmt, %rsi // Set the format string
	movq	stderr, %rdi   // Get stderr 
	call	fprintf        // Print the message to stderr

	jmp	_exit

// ==================
// Description:
//	Print a usage message to standard error and exit the program via `_exit`
// ==================
usage:
	xorl	%eax, %eax       // Set rax to 0
	movq	(%rsi), %rdx     // Get argv[0]
	movq	$usage_fmt, %rsi // Set the format string
	movq	stderr, %rdi     // Get stderr 
	call	fprintf          // Print the error
	// FALLTHROUGH

// ==================
// Description:
//	Exit the program with the return code EXIT_FAILURE
// ==================
_exit:
	movl	$EXIT_FAILURE, %eax
	call	exit