From 9becdea2b956ba41715cef648e0880745696c06b Mon Sep 17 00:00:00 2001
From: andromeda <andromeda@lenovo>
Date: Wed, 25 Mar 2026 21:37:40 +0100
Subject: [PATCH] minor improvements

---
 twasm/asm/main.asm | 48 ++++++++++++++++++++++++----------------------
 1 file changed, 25 insertions(+), 23 deletions(-)

diff --git a/twasm/asm/main.asm b/twasm/asm/main.asm
index 5ebc9cb..69708bc 100644
--- a/twasm/asm/main.asm
+++ b/twasm/asm/main.asm
@@ -650,7 +650,7 @@ tokenise:
 
   .loop:
     cmp rdi, rsi ; if current byte greater than last byte
-    jge .break    ; then break
+    jge .break   ; then break
 
     mov dl, [rdi] ; dl = current byte
 
@@ -667,18 +667,17 @@ tokenise:
     push rdi
     push rax
     push rdx
-    ; TODO probably should not ignore null bytes
-    mov rsi, whitespace_2 ; rsi -> list of whitespace (ignored) bytes
-    mov rdi,            2 ; rdi = size of list in bytes
+    mov rsi, whitespace_2 ; rsi -> list of whitespace bytes
+    mov rdi,            2 ; rdi = size of the list in bytes
                           ; dl = current byte
     call elemb
     ; al = 0 if not whitespace, 1 if whitespace
-    cmp al, 1  ; check if current byte is whitespace
-    pop rdx
-    pop rax
-    pop rdi
-    pop rsi
-    je .skip_byte_whitespace
+    test eax, 1  ; check if current byte is whitespace
+    pop rdx ; dl = current byte
+    pop rax ; rax = number of tokens processed
+    pop rdi ; rdi -> current byte of program
+    pop rsi ; rsi -> last byte of program
+    jnz .skip_byte_whitespace
 
     test byte [.expecting], E_OPERATOR ; check if an operator is expected
     jnz .operator                      ; if so, handle it
@@ -690,17 +689,19 @@ tokenise:
     call print.debug
     mov rsi, .msg_comment
     call print
-    pop rsi
+    pop rsi ; rsi -> last byte of program
+
     test byte [.expecting], E_COMMENT ; make sure a comment is expected
     jz .unexpected_comment            ; if not, error
   .comment_loop:
-    ; TODO range check rdi
     mov dl, [rdi]     ; dl = current byte
 
     cmp dl, 0x0A      ; if current byte is a newline
     je .comment_break ; then break
 
     inc rdi           ; point to next unread byte
+    cmp rdi, rsi
+    jge .break
     jmp .comment_loop
   .comment_break:
     jmp .loop
@@ -729,7 +730,7 @@ tokenise:
     test byte [.expecting], E_COMMA            ; make sure a comma was expected
     jz .unexpected_comma                       ; if not, error
     inc rdi
-    mov [.expecting], E_WHITESPACE | E_OPERAND ; else, make operand expected
+    mov byte [.expecting], E_WHITESPACE | E_OPERAND ; else, make operand expected
     jmp .loop                                  ;       and loop
 
   .newline_mk_flags:
@@ -757,15 +758,12 @@ tokenise:
     call print
     pop rsi
 
-    push rax
+    mov rcx, rax                 ; rcx = number of tokens processed
     xor eax, eax                 ; eax = number of bytes in operator
     mov [.pending_operator], eax ; zero pending operator
 
   .operator_loop:
     ; TODO give this its own error
-    ; TODO make this pop rax
-    cmp eax, 4              ; check that operator is short enough
-    jg .unexpected_operator ; if not, error
 
     mov dl, [rdi] ; next byte
 
@@ -781,8 +779,12 @@ tokenise:
 
     mov [.pending_operator + eax], dl
 
-    inc rax            ; inc byte counter
+    inc eax            ; inc byte counter
     inc rdi            ; inc byte pointer
+    cmp eax, 4              ; check that operator is short enough
+    jg .unexpected_operator ; if not, error
+    cmp rdi, rsi
+    jge .break
     jmp .operator_loop ; and loop
 
   .operator_break:
@@ -791,11 +793,11 @@ tokenise:
 
     mov edi, [.pending_operator] ; edi = operator to be searched
     call identify_operator
-               ; ax = operator's token ID
-    mov cx, ax ; cx = operator's token ID for safe keeping
-
-    pop rdi ; rdi = byte counter
-    pop rax ; rax = tokens processed
+    ; ax = operator's token ID
+    push rcx
+    mov ecx, eax ; cx = operator's token ID
+    pop rax      ; rax = tokens processed
+    pop rdi      ; rdi = byte counter
 
     mov [TOKEN_TABLE_ADDR + rax * 2], cx
     inc rax ; plus 1 token processed