LOAD_ADDR equ 0x00010000 ; address this program is loaded at

TOKEN_TABLE_ADDR equ 0x00060000 ; address the token table is loaded at
TOKEN_TABLE_SIZE equ 0x1000     ; max length of table

TEST_ARENA_ADDR equ 0x00060000 ; address to run tests at
TEST_ARENA_SIZE equ 0x1000     ; maximum size tests can use

OUTPUT_ADDR equ 0x00070000 ; address of outputed binary
OUTPUT_SIZE equ 0x1000     ; max length of outputed binary

STACK_ADDR equ 0x00060000 ; address to put the 64-bit stack at

[bits 64]
[org LOAD_ADDR]

start:
  mov rsp, STACK_ADDR ; we might need more stack space, let's just be safe

  mov rsi, msg_welcome
  call print

  call run_tests

  call clear_token_table

  jmp halt

; ------------------------------------------------------------------------------
; tokenising
; ------------------------------------------------------------------------------

; ------------------------------------------------------------------------------
; copy_token 
;
; description:
; copies a token from one spot in memory to another
;
; parameters:
; rdi -> start of buffer to be read
; rsi -> start of buffer to be written
;
; returned:
; rax -> last byte read
; rdx -> last byte written
; ------------------------------------------------------------------------------

copy_token:
  .loop:

    mov dl, [rdi] ; move bit to compare to current byte in read buffer

    push rdi ; push incrementors to call elemb
    push rsi ;

    mov rdi, 8                  ; length of terminator list
    mov rsi, token_terminator_8 ; start of terminator list
                                ; dl set before pushing rdi
    call elemb

    pop rsi ;
    pop rdi ; pop incrementors after call

    cmp rax, 1 ; check if the next character is a token terminator
    je .break  ; > if so, break the function

    ; rdi and rsi set from previous loop iteration
    call copy_byte ; if not, copy the current byte in read buffer

    inc rdi ; read pointer
    inc rsi ; write pointer

    jmp .loop

  .break:
    mov rax, rdi ; -> last byte read
    mov rdx, rsi ; -> last byte written
    ret
    
; ------------------------------------------------------------------------------
; copy_byte
;
; description:
; copies a byte from one spot in memory to another
;
; parameters:
; rdi -> word to be read
; rsi -> word to be written
;
; returned:
; al = byte that was read; the rest of rax is zeroed
; ------------------------------------------------------------------------------

copy_byte:
  xor rax, rax ; zero out so it returns fine
  mov al, [rdi]
  mov [rsi], al
  ret

; ------------------------------------------------------------------------------
; utilities
; ------------------------------------------------------------------------------

; ------------------------------------------------------------------------------
; print
;
; description:
; prints a null-terminated string
;
; parameters:
; rsi -> start of null-terminated string
; ------------------------------------------------------------------------------

print:
  push rdx
  mov rdx, 0x3F8
  .loop:
    mov al, [rsi]
    test al, al
    jz .done
    out dx, al
    inc rsi
    jmp .loop
  .done:
    pop rdx
    ret

; ------------------------------------------------------------------------------
; halt
;
; description:
; halts the program, silly :)
; ------------------------------------------------------------------------------

halt:
  mov rsi, msg_halt
  call print
  hlt
  jmp halt

; ------------------------------------------------------------------------------
; elemb
;
; description:
; checks if given byte is element of the specified list
;
; parameters:
; rdi = size of list
; rsi -> start of list
; dl =  given byte
;
; returned:
; rax = 0: is not an element
;       1: is an element
; ------------------------------------------------------------------------------

elemb:
  .loop
    cmp rdi, 0    ; check if remaining length 0
    je .not_found ; if so, break; dl not an element of list

    mov al, [rsi]
    cmp al, dl    ; check if current byte in list is the desired byte
    je .found     ; if so, break; dl an element of list

    inc rsi ; move to next byte
    dec rdi ; and reduce remaining length

    jmp .loop

  .not_found
    xor rax, rax ; return 0; dl not an element of list
    ret

  .found
    xor rax, rax
    mov rax, 1 ; return 1; dl an element of list
    ret

  .f db "found", 0x0D, 0x0A, 0x00
  .nf db "not found", 0x0D, 0x0A, 0x00

; ------------------------------------------------------------------------------
; clear_token_table
;
; description:
; clears the token table as specified by TOKEN_TABLE_SIZE and TOKEN_TABLE_ADDR
; ------------------------------------------------------------------------------

clear_token_table:
  xor rax, rax                  ; value to write
  mov rcx, TOKEN_TABLE_SIZE / 4 ; number of double words
  mov rdi, TOKEN_TABLE_ADDR     ; address to start
  rep stosd
  ret

; ------------------------------------------------------------------------------
; clear_test_arena
;
; description:
; clears the test arena as specified by TEST_ARENA_SIZE and TEST_ARENA_ADDR
; ------------------------------------------------------------------------------

clear_test_arena:
  xor rax, rax                  ; value to write
  mov rcx, TOKEN_TABLE_SIZE / 4 ; number of double words
  mov rdi, TOKEN_TABLE_ADDR     ; address to start
  rep stosd
  ret

; ------------------------------------------------------------------------------
; tests
; ------------------------------------------------------------------------------

; ------------------------------------------------------------------------------
; run_tests
;
; description:
; runs all tests
; ------------------------------------------------------------------------------

run_tests:
  mov rsi, .msg
  call print

  call clear_test_arena
  call test_copy_byte

  call clear_test_arena
  call test_copy_token

  call clear_test_arena
  call test_elemb

  ret
  .msg db "running test suite...", 0x0D, 0x0A, 0x00

; ------------------------------------------------------------------------------
; test_copy_byte
;
; description:
; tests copy_byte described functionality
; ------------------------------------------------------------------------------

test_copy_byte:
  mov rsi, .msg
  call print

  mov rdi, test_byte       ; byte to be copied
  mov rsi, TEST_ARENA_ADDR ; location of test
  call copy_byte

  mov cx, [rsi]
  and ax, 0xFF ; only compare bottom byte
  and cx, 0xFF
  cmp ax, cx ; compare returned byte to copied byte
  jne .fail
  cmp al, [test_byte] ; compare returned byte to expected byte
  jne .fail

  .pass:
    mov rsi, msg_pass
    call print
    ret
  .fail:
    mov rsi, msg_fail
    call print
    ret
  .msg db "test_copy_byte...", 0x00

; ------------------------------------------------------------------------------
; test_copy_token
;
; description:
; tests copy_token described functionality
; ------------------------------------------------------------------------------

test_copy_token:
  mov rsi, .msg
  call print

  ; test case: space terminated

  mov rdi, test_token_space ; read buffer
  mov rsi, TEST_ARENA_ADDR  ; write buffer
  call copy_token

  ; check reported final indicies with the expected final indicies
  cmp rax, test_token_space + 8 ; last byte read
  jne .fail
  cmp rdx, TEST_ARENA_ADDR + 8 ; last byte written
  jne .fail

  mov rsi, TEST_ARENA_ADDR
  mov rcx, [rsi]
  cmp rcx, [test_token_space] ; check if copied token matches expected token
  jne .fail                   ; if not, fail

  ; test case: null terminated

  mov rdi, test_token_null ; read buffer
  mov rsi, TEST_ARENA_ADDR ; write buffer
  call copy_token

  ; check reported final indicies with the expected final indicies
  cmp rax, test_token_null + 8 ; last byte read
  jne .fail
  cmp rdx, TEST_ARENA_ADDR + 8 ; last byte written
  jne .fail

  mov rsi, TEST_ARENA_ADDR
  mov rcx, [rsi]
  cmp rcx, [test_token_null] ; check if copied token matches expected token
  jne .fail                  ; if not, fail

  .pass:
    mov rsi, msg_pass
    call print
    ret
  .fail:
    mov rsi, msg_fail
    call print
    ret
  .msg db "test_copy_token...", 0x00

; ------------------------------------------------------------------------------
; test_elemb
;
; description:
; tests elemb described functionality
; ------------------------------------------------------------------------------

test_elemb:
  mov rsi, .msg
  call print

  ; [0]
  mov rdi, 5
  mov rsi, test_elemb_5
  mov dl, [test_elemb_5]
  call elemb
  cmp al, 1
  jne .fail

  ; [n - 1]
  mov rdi, 5
  mov rsi, test_elemb_5
  mov dl, [test_elemb_5 + 4]
  call elemb
  cmp al, 1
  jne .fail

  ; [1]
  mov rdi, 5
  mov rsi, test_elemb_5
  mov dl, [test_elemb_5 + 1]
  call elemb
  cmp al, 1
  jne .fail

  ; not present
  mov rdi, 5
  mov rsi, test_elemb_5
  mov dl, 0xDA
  call elemb
  cmp al, 0
  jne .fail

  ; 0 length list
  mov rdi, 0
  mov rsi, test_elemb_0
  mov dl, 0x34
  call elemb
  cmp al, 0
  jne .fail

  .pass:
    mov rsi, msg_pass
    call print
    ret
  .fail:
    mov rsi, msg_fail
    call print
    ret
  .msg db "test_elemb...", 0x00

; ------------------------------------------------------------------------------
; data
; ------------------------------------------------------------------------------

tokens:
  .length2
    db "r8"
    dw 0x0008
    db "r9"
    dw 0x0009
    db "ax"
    dw 0x0020
    db "bx"
    dw 0x0021
    db "cx"
    dw 0x0022
    db "dx"
    dw 0x0023
    db "si"
    dw 0x0024
    db "di"
    dw 0x0025
    db "sp"
    dw 0x0026
    db "bp"
    dw 0x0027
    db "al"
    dw 0x0030
    db "bl"
    dw 0x0031
    db "cl"
    dw 0x0032
    db "dl"
    dw 0x0033
    db "ah"
    dw 0x0040
    db "bh"
    dw 0x0041
    db "ch"
    dw 0x0042
    db "dh"
    dw 0x0043
    db "cs"
    dw 0x0044
    db "ds"
    dw 0x0045
    db "es"
    dw 0x0046
    db "fs"
    dw 0x0047
    db "gs"
    dw 0x0048
    db "ss"
    dw 0x0049
  .length3
    db "rax"
    dw 0x0000
    db "rbx"
    dw 0x0001
    db "rcx"
    dw 0x0002
    db "rdx"
    dw 0x0003
    db "rsi"
    dw 0x0004
    db "rdi"
    dw 0x0005
    db "rsp"
    dw 0x0006
    db "rbp"
    dw 0x0007
    db "r10"
    dw 0x000A
    db "r11"
    dw 0x000B
    db "r12"
    dw 0x000C
    db "r13"
    dw 0x000D
    db "r14"
    dw 0x000E
    db "r15"
    dw 0x000F
    db "eax"
    dw 0x0010
    db "ebx"
    dw 0x0011
    db "ecx"
    dw 0x0012
    db "edx"
    dw 0x0013
    db "esi"
    dw 0x0014
    db "edi"
    dw 0x0015
    db "esp"
    dw 0x0016
    db "ebp"
    dw 0x0017
    db "r8d"
    dw 0x0018
    db "r9d"
    dw 0x0019
    db "r8w"
    dw 0x0028
    db "r9w"
    dw 0x0029
    db "sil"
    dw 0x0034
    db "dil"
    dw 0x0035
    db "spl"
    dw 0x0036
    db "bpl"
    dw 0x0037
    db "r8b"
    dw 0x0038
    db "r9b"
    dw 0x0039
    db "cr0"
    dw 0x004A
    db "cr2"
    dw 0x004B
    db "cr3"
    dw 0x004C
    db "cr4"
    dw 0x004D
    db "cr8"
    dw 0x004E
  .length4
    db "r10d"
    dw 0x001A
    db "r11d"
    dw 0x001B
    db "r12d"
    dw 0x001C
    db "r13d"
    dw 0x001D
    db "r14d"
    dw 0x001E
    db "r15d"
    dw 0x001F
    db "r10w"
    dw 0x002A
    db "r11w"
    dw 0x002B
    db "r12w"
    dw 0x002C
    db "r13w"
    dw 0x002D
    db "r14w"
    dw 0x002E
    db "r15w"
    dw 0x002F
    db "r10b"
    dw 0x003A
    db "r11b"
    dw 0x003B
    db "r12b"
    dw 0x003C
    db "r13b"
    dw 0x003D
    db "r14b"
    dw 0x003E
    db "r15b"
    dw 0x003F
  .length5
  .end

msg_welcome db "Welcome to Twasm", 0x0D, 0x0A, 0x00
msg_halt db "halted.", 0x0D, 0x0A, 0x00
msg_pass db "passed.", 0x0D, 0x0A, 0x00
msg_fail db "failed.", 0x0D, 0x0A, 0x00

test_byte db "Q" ; unterminated, just a byte chillin
test_token_null db "TestTokn", 0x00 ; followed by null terminator. Quad word
test_token_space db "TestTokn " ; followed by space. Quad word
test_elemb_0
test_elemb_5 db 0x54, 0x00, 0x21, 0x20, 0x34

token_terminator_8 db 0x00, " ", 0x0A, 0x0D, 0x00, 0x00, 0x00, 0x00

debug_string db "debug_string", 0x0D, 0x0A, 0x00