parsing mem operand

stages/as0/as0.asm

@@ -613,6 +613,28 @@ is_digit:
     setb al
     ret
 
+;; converts char $dil to a digit with radix $rsi, returning it in $edx. $al is set to 1 if the char is a valid digit, and 0 otherwise.
+to_digit:
+    lea eax, [rsi - 2]
+    cmp eax, 35
+    jae .invalid
+    movzx rdi, dil
+    lea edx, [rdi - 65] ; 'A' = 65
+    and edx, -33        ; convert to uppercase
+    add edx, 10         ; 'A' should map to 10
+    lea eax, [rdi - 48] ; '0' = 48
+    cmp esi, 11
+    cmovb edx, eax              ; if radix <= 10, then take the difference from '0'
+    cmp edi, 58                  
+    cmovb edx, eax              ; or if char < '9', then take the difference from '0'
+    xor eax, eax
+    cmp edx, esi
+    setb al                     ; al = edx < radix
+    ret
+.invalid:
+    xor eax, eax
+    ret
+
 exit:
     mov rax, 60
     syscall
@@ -628,10 +650,9 @@ exit:
 struc operand
 .type resd 1
 .size resd 1
-.reg resb 1
+.reg resq 1
-.index resb 1
+.index resq 1
-.scale resb 1
+.scale resd 1
-.pad resb 1
 .disp resd 1
 endstruc
 
@@ -649,10 +670,13 @@ try_parse_reg:
     push rbx                    ; buf
     push r12                    ; reg_num
     push r13                    ; size
+
+    push rdi
+
     xor r12d, r12d
     xor r13d, r13d
 
-    lea rbx, [rel buf]
+    mov rbx, rdi
     mov dil, byte [rbx]
     call is_alpha
     jnb .done
@@ -723,8 +747,8 @@ try_parse_reg:
 
 .done:
     mov rax, rbx
-    lea rbx, [rel buf]
+    pop rdi
-    sub rax, rbx                ; number of bytes consumed
+    sub rax, rdi                ; number of bytes consumed
 
     movzx rdx, r12d             ; reg_num
     shl rdx, 32
@@ -755,6 +779,253 @@ try_parse_reg:
     sub r13d, eax
     jmp .r_or_e
 
+try_parse_num:
+    push rbx                    ; buf
+    push r12                    ; radix
+    push r13                    ; acc
+
+    push rdi
+
+    mov r12, 10
+    xor r13d, r13d
+    xor eax, eax
+    mov rbx, rdi
+
+    ; eat leading - and flip sign
+    cmp byte [rbx], '-'
+    sete al
+    sub r13, rax
+    lea rbx, [rbx + rax]
+
+    ; eat any leading +
+    cmp byte [rbx], '+'
+    sete al
+    lea rbx, [rbx + rax]
+
+    mov dil, byte [rbx]
+    call is_digit
+    jnb .not_num
+
+    cmp dil, '0'
+    jne .loop
+    inc rbx
+    mov dil, byte [rbx]
+    and dil, 0xdf               ; convert to uppercase
+    cmp dil, 'X'
+    sete al
+    lea ecx, [eax + eax * 2]
+    shl ecx, 1
+    add r12d, ecx
+    cmp dil, 'O'
+    sete al
+    shl al, 1
+    sub r12d, eax
+    cmp dil, 'B'
+    sete al
+    shl al, 3
+    sub r12d, eax
+
+    cmp r12d, 10
+    setne al
+    lea rbx, [rbx + rax]
+
+.loop:
+    mov dil, byte [rbx]
+    test dil, dil
+    jz .done
+
+    mov esi, r12d
+    call to_digit
+    mov ecx, edx
+    test al, al
+    jz .done
+    mov rax, r13
+    imul r12
+    add rax, rcx
+    mov r13, rax
+    inc rbx
+    jmp .loop
+
+.done:
+    mov rax, rbx
+    pop rdi
+    sub rax, rdi                ; number of bytes consumed
+    mov rdx, r13                ; acc
+
+    pop r13
+    pop r12
+    pop rbx
+    ret
+
+.not_num:
+    xor rax, rax
+    xor rdx, rdx
+    pop rdi
+    pop r13
+    pop r12
+    pop rbx
+    ret
+
+skip_whitespace_rbx:
+    mov al, byte [rbx]
+    cmp al, ' '
+    ja .done
+    inc rbx
+    jmp skip_whitespace_rbx
+.done:
+    ret
+
+    global try_parse_mem
+try_parse_mem:
+    push rbx                    ; buf
+    push r12                    ; size
+    push rdi
+
+    mov rbx, rdi                ; buf ptr
+    mov r12, rsi                ; output operand ptr
+
+    mov dword [r12 + operand.type], KIND_MEM
+    mov dword [r12 + operand.size], 0
+    mov qword [r12 + operand.reg], -1
+    mov qword [r12 + operand.index], -1
+    mov dword [r12 + operand.scale], 0
+    mov dword [r12 + operand.disp], 0
+
+    mov edi, dword [rbx]
+
+    cmp dil, '['
+    je .bracket
+
+    lea rbx, [rbx + 4]
+
+    cmp edi, 'byte'
+    je .bracket
+
+    mov dword [r12 + operand.size], 1 ; 16-bit
+    cmp edi, 'word'
+    je .bracket
+
+    lea rbx, [rbx - 4]
+    mov dl, byte [rbx]
+    inc rbx
+    mov edi, dword [rbx]
+
+    xor eax, eax
+    mov dword [r12 + operand.size], 2 ; 16-bit
+    cmp dl, 'q'
+    sete al
+    cmp dl, 'd'
+    sete dl
+    or al, dl
+    jz .not_mem
+    add dword [r12 + operand.size], eax ; 32-bit?
+
+    cmp edi, 'word'
+    jne .not_mem
+    lea rbx, [rbx + 4]
+
+.bracket:
+    call skip_whitespace_rbx
+    mov dl, byte [rbx]
+    cmp dl, '['
+    jne .not_mem
+    inc rbx
+
+    call skip_whitespace_rbx
+    mov rdi, rbx
+    call try_parse_reg
+    add rbx, rax
+
+    call skip_whitespace_rbx
+    cmp byte [rbx], '*'
+    jne .base
+    mov qword [r12 + operand.index], rdx
+
+.scale:
+    inc rbx
+    call skip_whitespace_rbx
+    mov dil, byte [rbx]
+    mov esi, 10
+    call to_digit
+    mov dword [r12 + operand.scale], edx
+    add rbx, rax
+
+    call skip_whitespace_rbx
+    mov dil, byte [rbx]
+    cmp dil, ']'
+    je .done
+
+
+    ; TODO: allow for label as displacement
+.disp:
+    sub dil, `,`
+    neg dil                     ; sign
+    mov al, dil
+    inc al
+    test al, -3
+    jne .not_mem
+    movsx edi, dil
+    mov dword [r12 + operand.disp], edi
+
+    inc rbx
+    call skip_whitespace_rbx
+    mov rdi, rbx
+    call try_parse_num
+    test eax, eax
+    jz .not_mem
+    add rbx, rax
+    mov eax, dword [r12 + operand.disp]
+    imul edx
+    mov dword [r12 + operand.disp], eax
+
+    call skip_whitespace_rbx
+    cmp byte [rbx], ']'
+    je .done
+    jmp .not_mem
+
+.base:
+    mov qword [r12 + operand.reg], rdx
+    cmp byte [rbx], ']'
+    je .done
+    cmp byte [rbx], '+'
+    jne .not_mem
+    inc rbx
+    call skip_whitespace_rbx
+
+    mov rdi, rbx
+    call try_parse_reg
+    test eax, eax
+    jz .disp
+    add rbx, rax
+    mov qword [r12 + operand.index], rdx
+
+    call skip_whitespace_rbx
+    cmp byte [rbx], '*'
+    je .scale
+    cmp byte [rbx], '+'
+    je .disp
+    cmp byte [rbx], ']'
+    jne .not_mem
+    
+.done:
+    inc rbx
+    mov rax, rbx
+    pop rdi
+    sub rax, rdi                ; number of bytes consumed
+    mov rdx, r12                ; output operand ptr
+
+    pop r12
+    pop rbx
+    ret
+
+.not_mem:
+    xor rax, rax
+    pop rdi
+    pop r12
+    pop rbx
+    ret
+
+
 global _entry_print_table
 _entry_print_table:
     call print_op_table

stages/as0/test.rs

@@ -2,7 +2,8 @@ unsafe extern "C" {
     // unsafe fn parse_gpr(src: *mut Source, prefix: u8) -> (u64, u64);
     // unsafe fn parse_reg(src: *mut Source) -> (u64, u64);
 
-    unsafe fn try_parse_reg() -> (usize, Register);
+    unsafe fn try_parse_reg(_: *const u8) -> (usize, Register);
+    unsafe fn try_parse_mem(_: *const u8, _: *mut Operand) -> (usize, *const Operand);
 
     #[link_name = "buf"]
     static mut BUF: [u8; 0x100];
@@ -15,6 +16,17 @@ struct Register {
     num: u32,
 }
 
+#[repr(C)]
+#[derive(Debug, PartialEq, Eq)]
+struct Operand {
+    kind: u32,
+    size: u32,
+    reg: Register,
+    index: Register,
+    scale: u32,
+    disp: i32,
+}
+
 impl Register {
     fn new(num: u32, size: u32) -> Self {
         Self {
@@ -22,11 +34,15 @@ impl Register {
             size: size.ilog2(),
         }
     }
+    fn invalid() -> Self {
+        Self { num: !0, size: !0 }
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
+    use core::mem::MaybeUninit;
 
     fn write_to_buf(b: &[u8]) {
         assert!(b.len() <= 0x100, "Buffer overflow");
@@ -36,6 +52,61 @@ mod tests {
     }
 
     #[test]
+    fn parse_mem() {
+        let mut op = MaybeUninit::<Operand>::uninit();
+        // let (n, _) =
+        //     unsafe { try_parse_mem(b"qword [rax + rbx * 2 + 0x10]\0".as_ptr(), op.as_mut_ptr()) };
+        // assert_eq!(
+        //     (n, unsafe { op.assume_init_ref() }),
+        //     (
+        //         28,
+        //         &Operand {
+        //             kind: 2,
+        //             size: 3,
+        //             reg: Register::new(0, 8),
+        //             index: Register::new(3, 8),
+        //             scale: 2,
+        //             disp: 0x10,
+        //         }
+        //     )
+        // );
+
+        let (n, _) =
+            unsafe { try_parse_mem(b"byte [r12 + r13 * 4 - 0x20]\0".as_ptr(), op.as_mut_ptr()) };
+        assert_eq!(
+            (n, unsafe { op.assume_init_ref() }),
+            (
+                27,
+                &Operand {
+                    kind: 2,
+                    size: 0,
+                    reg: Register::new(12, 8),
+                    index: Register::new(13, 8),
+                    scale: 4,
+                    disp: -0x20,
+                }
+            )
+        );
+
+        let (n, _) = unsafe { try_parse_mem(b"word [r8 * 8 + 5]\0".as_ptr(), op.as_mut_ptr()) };
+        assert_eq!(
+            (n, unsafe { op.assume_init_ref() }),
+            (
+                17,
+                &Operand {
+                    kind: 2,
+                    size: 1,
+                    reg: Register::invalid(),
+                    index: Register::new(8, 8),
+                    scale: 8,
+                    disp: 5,
+                }
+            )
+        );
+    }
+
+    #[test]
+    #[ignore]
     fn parse_reg() {
         let cases = [
             (&b"rax"[..], (0, 8)),
@@ -106,8 +177,7 @@ mod tests {
         ];
 
         for (text, (reg, size)) in cases {
-            write_to_buf(text);
+            let (n, result) = unsafe { try_parse_reg(text.as_ptr()) };
-            let (n, result) = unsafe { try_parse_reg() };
             assert_eq!(
                 result,
                 Register::new(reg, size),