; bfc : a brainfuck compiler & interpreter
; Copyright (C) 2018 Weber Yann
; 
; This program is free software; you can redistribute it and/or modify
; it under the terms of the GNU General Public License as published by
; the Free Software Foundation; either version 3 of the License, or
; any later version.
; 
; This program is distributed in the hope that it will be useful,
; but WITHOUT ANY WARRANTY; without even the implied warranty of
; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
; GNU General Public License for more details.
; 
; You should have received a copy of the GNU General Public License
; along with this program.  If not, see <http://www.gnu.org/licenses/>.
;
; A brainfuck compiler & interpreter :
; Build : nasm -felf64 bfc.asm && ld -s -melf_x86_64 bfc.o -o bfc
;
; ./bfc [-h] [-e [-o a.out]] FILE.bf
; Options :
; 	-h print usage and exit
;	-e tell bfc to produce a elf file
;	-o with -e indicate the file to create
;	FILE.bf the brainfuck source file to compile
[bits 64]
%use smartalign
ALIGNMODE k8

%define lbl_incresize 0x2e
%define lbl_decresize 0x50

%define MAP_INC_SIZE 0x1000
%define MAP_INC_MASK 0x0FFF
%define BFMEM_INIT_SZ MAP_INC_SIZE

section .data

	bf_start:
		jmp .start
		.mremap: ; rbx is resize size
			cmp rbx, MAP_INC_SIZE
			jg .remap_cont
			xor rbx, rbx
			.remap_cont:
			xor rbx, MAP_INC_MASK
			add rbx, MAP_INC_SIZE
			add r14, rbx
			mov rax, 0x19	; mremap
			mov rdi, r15	; addr
			mov rsi, r14 	; oldlen
			jc .erremap ; overflow, too much mem !!!!
			mov rdx, r14	; newlen
			xor r10, r10	; maymove
			xor r11, r11
			syscall
			cmp rax, 0
			jle .erremap
			mov r15, rax
			ret
			.erremap:
				mov rax, 0x3c
				mov rdi, 0x2
				syscall
			align 8
		.lbl_incresize: equ $ - bf_start
		.incresize: ; rbx is resize size
			call .mremap
			ret
			align 8
		.lbl_decresize: equ $ - bf_start
		.decresize: ; rbx is decrement
			push rsi
			push r14
			call .incresize
			pop rcx ; old size
			mov rdi, rsi
			inc rdi
			.decresize_cpy:
				movsb
				loop .decresize_cpy
			mov byte [r15], 0
			pop rsi
			ret
		.errmap:
			mov rax, 0x3c
			mov rdi, 0x1
			syscall	
		.start:
		;map init
		mov rax, 0x9    ;mmap
		xor rdi, rdi
		mov rsi, BFMEM_INIT_SZ    ;len
		mov r14, rsi
		mov rdx, (0x1|0x2)	;PROT_READ | PROT_WRITE
		mov r10, (0x2 | 0x20)	;flags MAP_PRIVATE | MAP_ANONYMOUS
		mov r8, -1      ;fd
		xor r9, r9
		syscall
		cmp rax, 0
		jle .errmap
		mov r15, rax
		mov rsi, r15
		add rsi, (BFMEM_INIT_SZ / 2) ; BF ptr
		align 8
	bf_start_sz: equ $ - bf_start

	
	; In piece of code call jump is achieved by adding
	; an offset to the JIT map base addr
	; this base address has to be on top of the stack
	; when executing this small piece of code
	;
	; the first instruction has to be a mov of a byte
	; in a register. This operation will be updated to
	; "pass" a parameter
	bf_decptr:
		mov rbx, strict qword 0x1
		push rbx
		cmp rsi, rbx
		jge .end
		mov rax, [rsp+8]
		add rax, bf_start.lbl_decresize
		call rax
		.end:
		pop rbx
		sub rsi, rbx
	bf_decptr_sz: equ $ - bf_decptr

	bf_incptr:
		mov rbx, strict qword 0x1
		push rbx
		mov rax, rsi
		sub rax, r15
		cmp rax, r14
		jl .end
		mov rax, [rsp+8]
		add rax, bf_start.lbl_incresize
		call rax
		.end:
		pop rbx
		add rsi, rbx
	bf_incptr_sz: equ $ - bf_incptr

	bf_incval:
		mov rbx, strict qword 0x1
		xor rax, rax
		mov al,[rsi]
		add rax, rbx
		mov [rsi], al
	bf_incval_sz: equ $ - bf_incval

	bf_decval:
		mov rbx, strict qword 0x1
		xor rax, rax
		mov al, [rsi]
		sub rax, rbx
		mov [rsi], al
	bf_decval_sz: equ $ - bf_decval

	bf_readval:
		mov rdx, strict qword 0x1
		push rsi
		xor rax, rax ; read
		xor rdi, rdi ; stdin
		syscall
		test rax, rax
		jnz .end
		mov byte [rsi], 0
		.end:
		pop rsi
	bf_readval_sz: equ $ - bf_readval

	bf_writeval:
		mov rdx, strict qword 0x1
		push rsi
		xor rax, rax ; write
		inc rax
		mov rdi, rax ; stdout
		syscall
		pop rsi
	bf_writeval_sz: equ $ - bf_writeval

	bf_loopstart:
		mov rbx, strict qword 0x1
		xor rdx, rdx
		mov dl, [rsi]
		cmp dl, 0
		jnz .end
		jmp rbx
		.end:
	bf_loopstart_sz: equ $ - bf_loopstart

	bf_loopend:
		mov rbx, strict qword 0x1
		xor rdx, rdx
		mov dl, [rsi]
		cmp dl, 0
		jz .end
		jmp rbx
		.end:
	bf_loopend_sz: equ $ - bf_loopend

	bf_exit:
		mov rax, 0x3c
		xor rdi, rdi
		syscall
	bf_exit_sz: equ $ - bf_exit

	miss_open: db "Missing opening '[' matching closing ']'"
	miss_open_sz: equ $ - miss_open
	chr_list : db ": ", 0xA
	read_error: db "Error reading file "
	read_error_sz: equ $ - read_error
	usage_err: db "Usage : FILE.BF"
	usage_err_sz: equ $ - usage_err
	open_err: db "Error opening file", 0xa
	open_err_sz: equ $ - open_err

section .bss
	read_buff: resb 128

section .text
global _start
_start:
	mov rcx, [rsp] ; argc
	cmp rcx, 2
	jne .badarg

	; JIT code map init
	; rsi map size
	mov rsi, 0x10
	call initmap

	mov rax, 0x2 ; open
	mov rdi, [rsp+16] ; argv[1]
	xor rsi, rsi ; no flags
	xor rdx, rdx ; O_RDONLY
	syscall
	cmp rax, 0
	jl .err_open

	call compile_bf

	; set code map perm
	mov rax, 0xA ; mprotect
	mov rdi, r15
	mov rsi, r14
	mov rdx, 0x4 | 0x1 ; PROT_EXEC | PROT_READ
	syscall

	push r15
	jmp r15

	.err_open:
		mov rax, 1 ; write
		mov rdi, 2 ; stderr
		mov rsi, open_err
		mov rdx, open_err_sz
		syscall
	.badarg:
		mov rax, 1 ;write
		mov rdi, 2 ; stderr
		mov rsi, usage_err
		mov rdx, 8 ; "Usage : "
		syscall

		mov rsi, [rsp+8] ; argv[0]
		xor rdx, rdx
		xor rcx, rcx
		.argv0len:
			inc rsi
			inc rdx
			mov cl, [rsi]
			test cl, cl
			jnz .argv0len

		mov rax, 1
		mov rdi, 2
		mov rsi, [rsp+8] ; argv[0]
		syscall

		mov rax, 1 ;write
		mov rdi, 2 ; stderr
		mov rsi, usage_err + 7
		mov rdx, usage_err_sz - 7 ; usage opts
		syscall

		mov rax, 1
		mov rdi, 2
		mov rsi, chr_list + 2 ; \n
		mov rdx, 1
		syscall

		mov rax, 0x3c ; exit
		mov rdi, 1
		syscall

; Init a writable memory map
; len in rsi
; ret : r14 map size
; 	r15 map addr
initmap:
	mov r14, rsi
	mov rax, 0x9    ;mmap
	xor rdi, rdi	;addr
	; rsi is len
	mov rdx, (0x1|0x2)	;PROT_READ | PROT_WRITE
	mov r10, (0x2 | 0x20)	;flags MAP_PRIVATE | MAP_ANONYMOUS
	mov r8, -1      ;fd
	xor r9, r9
	syscall
	cmp rax, 0
	jle .err
	mov r15, rax
	ret
	.err:
		mov rax, 0x3c
		mov rdi, 0x2
		syscall
; resize the memory map
; addr in r15 (0 if init)
; len in r14
; inc size in bytes in r10
mremap:
	mov rax, 0x19    ;mremap
	mov rdi, r15	;addr
	mov rsi, r14    ;oldlen
	mov rdx, rsi
	add rdx, MAP_INC_SIZE ; newlen
	xor r10, r10 ; MAYMOVE
	xor r11, r11
	syscall
	cmp rax, 0
	jle .err
	mov r15, rax
	add r14, MAP_INC_SIZE
	ret
	.err:
		mov rax, 0x3c
		mov rdi, 0x3
		syscall

; JIT brainfuck compiler. Read bf code from fd and write
; machine code in a dedicated anon map with PROT_READ | PROT_EXEC
; Using heap to read file
; args :
;	rax source fd
;	r14 map size
;	r15 map addr
; ret :
;	rax map addr
compile_bf:
	; Allocating growing heap to store various datas
	; heap start will be stored in r13
%define fd [r13]
%define map_ptr [r13+0x8]
%define base_rsp [r13+0x10]
%define chr_repeat [r13+0x18]
%define line_count [r13+0x20]
%define chr_count [r13+0x28]
%define chr_buff_off 0x30
%define chr_buff [r13+chr_buff_off]
%define prev_chr_off 0x31
%define prev_chr [r13+prev_chr_off]
%define heap_size 0x32
	push rax ; source fd

	mov rax, 0xc ; brk
	xor rdi, rdi
	syscall

	push rax ; heap start

	mov rdi, rax
	add rdi, heap_size ; new heap addr
	mov rax, 0xc ; brk
	syscall

	pop rdi ; heap start
	mov r13, rdi
	pop rax ; source fd
	mov fd, rax

	; init heap
	mov byte prev_chr, 0
	mov qword chr_count, 0
	mov qword line_count, 0
	mov qword chr_repeat, 1
	mov base_rsp, rsp ; save rsp in heap

	; copy code map header
	mov rdi, r15
	mov rsi, bf_start
	mov rdx, bf_start_sz
	call code_cpy
	mov map_ptr, rax ; new map ptr

	; read first char in prev_chr
	xor rax, rax ;read
	mov rdi, fd ; fd
	mov rsi, r13
	add rsi, prev_chr_off ; chr_prev
	mov rdx, 1 ; read 1 byte
	syscall
	cmp rax, 0
	jle .read_error

	.readloop:
		xor rax, rax ;read
		mov rdi, fd ; fd
		mov rsi, r13
		add rsi, chr_buff_off ; buff byte
		mov rdx, 1 ; read 1 byte
		syscall
		cmp rax, 0
		je .endread
		jl .read_error ; error

		mov rax, chr_count ; chr counter
		inc rax
		mov chr_count, rax

		mov al, chr_buff

		; arg for loop is not a repeat counter
		cmp al, 0x5b ; '['
		je .cmpchar
		cmp al, 0x5d ; '['
		je .cmpchar

		cmp al, prev_chr
		je .incnum ; same instruction, incrementing counter

		.cmpchar:
		mov rdi, map_ptr ; prepare to copy in code map
		; compare previous char and store current in prev
		; note : chr_repeat has to be reset by .nxtinstr
		; after jump
		xchg prev_chr, al
		cmp al, 0x3c ; '<'
		je .lptr
		cmp al, 0x3e ; '>'
		je .rptr
		cmp al, 0x2b ; '+'
		je .incval
		cmp al, 0x2d ; '-'
		je .decval
		cmp al, 0x2e ; '.'
		je .wrval
		cmp al, 0x2c ; ','
		je .rdval
		cmp al, 0x5b ; '['
		je .loopstart
		cmp al, 0x5d ; ']'
		je .loopend

		cmp al, 0x0a ; '\n'
		je .line
		
		; chr is not an instruction, printing them
		; on stderr
		mov chr_buff, al
		mov rcx, chr_repeat
		.errchr:
			push rcx
			mov rax, 1 ; write
			mov rdi, 2 ; stderr
			mov rsi, r13
			add rsi, chr_buff_off ; heap buff
			mov rdx, rax ; sz 1
			syscall
			pop rcx
			loop .errchr
		jmp .nxtinstr

		.line: ; increment line counter in heap
			mov rax, line_count
			add rax, chr_repeat
			mov line_count, rax
			mov rcx, chr_repeat
			jmp .errchr ; print the newline

		

		; following ref copy assume rdi to be map_ptr
		.incval:
			mov rsi, bf_incval
			mov rdx, bf_incval_sz
			push rdx
			jmp .callcpy
		.decval:
			mov rsi, bf_decval
			mov rdx, bf_decval_sz
			push rdx
			jmp .callcpy
		.lptr:
			mov rsi, bf_decptr
			mov rdx, bf_decptr_sz
			push rdx
			jmp .callcpy
		.rptr:
			mov rsi, bf_incptr
			mov rdx, bf_incptr_sz
			push rdx
			jmp .callcpy
		.wrval:
		 	mov rsi, bf_writeval
			mov rdx, bf_writeval_sz
			push rdx
			jmp .callcpy
		.rdval:
			mov rsi, bf_readval
			mov rdx, bf_readval_sz
			push rdx
			jmp .callcpy
		.loopstart:
			push qword map_ptr ; ret addr
			mov rsi, bf_loopstart
			mov rdx, bf_loopstart_sz
			push rdx
			jmp .callcpy
		.loopend:
			cmp rsp, base_rsp
			je .loop_err_miss_open
			mov rsi, bf_loopend
			mov rdx, bf_loopend_sz
			push rdx
			call code_cpy
			mov map_ptr, rax
			pop rdx
			sub rax, rdx
			add rax, 2 ;arg addr in code map
			pop rbx ; loop_start code addr
			mov [rax], rbx ; loop end jump to start
			mov rax, map_ptr
			mov [rbx+2], rax ; start jump to end
			jmp .nxtinstr

		.callcpy:
			call code_cpy
			mov map_ptr, rax
			; set the 1st instr rgs in the mapping
			; and reinit chr_repeat
			pop rdx
			sub rax, rdx
			add rax, 2 ; arg addr in code map
			mov rbx, chr_repeat
			mov [rax], rbx
			.nxtinstr:
				; reinit chr_repeat
				mov qword chr_repeat, 1
				jmp .readloop

		.incnum:
			; same instruction found incrementing
			; chr_repeat
			mov rbx, chr_repeat
			inc rbx
			jc .incoverflow
			mov chr_repeat, rbx
			jmp .readloop

			.incoverflow:
				dec rbx
				mov chr_repeat, rbx
				jmp .cmpchar
			
		jmp .readloop
	.loop_err_miss_open:
		; miss_open_err
		mov rax, 1
		mov rdi, 2 ; stderr
		mov rsi, miss_open
		mov rdx, miss_open_sz
		syscall
		jmp .exit_error
	
	.read_error:
		xor rax, rax
		inc rax ; write
		mov rsi, 2 ; stderr
		mov rdi, read_error
		mov rdx, read_error_sz
		syscall
		jmp .exit_error

	.exit_error:
		mov rax, 1
		mov rdi, 2
		mov rsi, chr_list + 2
		mov rdx, 1
		syscall
		mov rax, 0x3c
		mov rdi, 0x11
		syscall

	.endread: ; EOF reached
	; fake \0 read to process prev_chr
	mov byte chr_buff, 0
	mov bl, prev_chr
	test bl, bl
	jnz .cmpchar
	; prevchar is 0 copying exit in code map
	.end_compile:
	mov rdi, map_ptr
	mov rsi, bf_exit
	mov rdx, bf_exit_sz
	call code_cpy

	; restoring stack
	mov rsp, base_rsp

	; restore heap
	mov rax, 0xc ; brk
	mov rdi, r13
	syscall

	ret

%undef fd
%undef map_ptr
%undef base_rsp
%undef chr_buff
%undef prev_chr
%undef heap_size

; Copy bf code from data to map
; Use :
; 	r15  map start
;	r14  map size
;	rdi  map ptr
;	rsi  code ptr
;	rdx  code size in bytes
; ret : 
;	rax : new map ptr
code_cpy:
	push rdx
	mov rax, rdi
	sub rax, r15 ; used len in map
	push rax
	mov rcx, rdx
	cmp rax, r14 ; rax is future len (after copy)
	jle .copy
	; resize
	push rsi ; save code_ptr
	call mremap
	pop rsi ; new code_ptr
	.copy:
	mov rdi, r15
	pop rax
	add rdi, rax ; new map ptr
	pop rcx ; size in words to write
	cld ; clear DF
	.copyloop:
		movsb
		loop .copyloop
	mov rax, rdi
	.ret:
	ret