From 2f7282577a567c23c6639c5d3a9c4357f5a94b71 Mon Sep 17 00:00:00 2001
From: janis <janis@nirgendwo.xyz>
Date: Thu, 9 Jul 2026 14:17:15 +0200
Subject: [PATCH] fixes

---
 stages/as0/as0.asm | 29 +++++++++++++++--------------
 stages/as0/test.rs |  8 +++++++-
 2 files changed, 22 insertions(+), 15 deletions(-)

diff --git a/stages/as0/as0.asm b/stages/as0/as0.asm
index 719c89b..bccb612 100644
--- a/stages/as0/as0.asm
+++ b/stages/as0/as0.asm
@@ -1884,7 +1884,7 @@ encode_inst:
     call write_non_zero_bytes
 
     pop rax
-    shl eax, 6
+    shr eax, 6
     dec eax
     cmp eax, 2
     jge .imm
@@ -1937,11 +1937,11 @@ write_imm:
     jmp write_bytes
 
 write_disp:
-    cmp edi, 0xff
+    cmp edi, 0xff               ; if disp > 0xff, write 4 bytes, else write 1 byte
     seta sil
-    shl sil, 1
-    inc sil
-    movzx rsi, sil
+    movzx esi, sil
+    lea esi, [rsi + rsi * 2]
+    inc esi
     jmp write_bytes
 
 write_bytes:
@@ -1986,8 +1986,8 @@ compute_modrm_sib:
     je .done_mem
     and eax, 0x37               ; clear mod bits, mod = 00b
 
-    cmp dword [rsi + operand.type], KIND_MEM
-    jne .done_mem
+    test dword [rsi + operand.type], KIND_MEM
+    jz .done_mem
 
     cmp dword [rsi + operand.disp], 0
     jne .sib                    ; mod = 01b or 10b
@@ -2023,7 +2023,7 @@ compute_modrm_sib:
 
     mov ecx, dword [rsi + operand.index + register.num]
     cmp ecx, -1
-    jne .done_mem
+    je .done_mem
 
     and ecx, 7
     shl ecx, 3                  ; index bits
@@ -2154,8 +2154,8 @@ struc LabelTable
 endstruc
 
 section .data
-    labels dq 0
-    label_refs dq 0
+    labels dq 0, 0
+    label_refs dq 0, 0
 
     section .text
 ;; calculates the hash of a byte sequence in $rdi of length $rsi, returning it in $eax
@@ -2174,6 +2174,7 @@ fasthash:
 .done:
     ret
 
+    global init_label_tables
 init_label_tables:
     mov edi, 0x100
     mov esi, 8
@@ -2199,19 +2200,19 @@ push_lbl_ref_by_hash:
     push r12
     push r13
 
-    mov rbx, [rel label_refs]
+    lea rbx, [rel label_refs]
     mov r13d, esi               ; label offset
     mov r12d, edi               ; label hash
     jmp push_lbl_common.hash_in_r12d
 
 ;; push label $rdi with length $rsi and offset $rdx into the label ref table
 push_lbl_ref:
-    mov rcx, [rel label_refs]
+    lea rcx, [rel label_refs]
     jmp push_lbl_common
 
 ;; push label $rdi with length $rsi and offset $edx into the label table
 push_lbl:
-    mov rcx, [rel labels]
+    lea rcx, [rel labels]
 push_lbl_common:
     push rbx
     push r12
@@ -2255,7 +2256,7 @@ find_lbl:
     push r12
     push r13
 
-    mov rbx, [rel labels]
+    lea rbx, [rel labels]
 
     call fasthash
     mov r12d, eax               ; store the hash in r12d
diff --git a/stages/as0/test.rs b/stages/as0/test.rs
index e73e2d0..6c5410a 100644
--- a/stages/as0/test.rs
+++ b/stages/as0/test.rs
@@ -14,6 +14,8 @@ unsafe extern "C" {
     #[link_name = "encode_inst"]
     unsafe fn encode_inst_impl(_: *mut Instruction) -> usize;
 
+    unsafe fn init_label_tables();
+
     unsafe fn fasthash(_: *const u8, _: usize) -> u32;
 
     #[link_name = "buf"]
@@ -143,6 +145,9 @@ impl Register {
     }
 }
 
+static mut TABLES_INIT: std::cell::LazyCell<()> =
+    std::cell::LazyCell::new(|| unsafe { init_label_tables() });
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -157,7 +162,8 @@ mod tests {
 
     #[test]
     fn parse_inst() {
-        let res = Instruction::parse(b"mov byte [rax], 4\0");
+        _ = unsafe { *TABLES_INIT };
+        let res = Instruction::parse(b"mov byte ['rax], 4\0");
         // assert_eq!(
         //     res,
         //     Some((