uJIT: support 64 bit operands for TEST. Use it to check for zero

It's one byte shorter than `cmp reg64, 0`. To illustrate:

```
48 83 f9 00        cmp    rcx, 0x0
48 85 c9           test   rcx, rcx
```

3 files changed, 6 insertions, 5 deletions

diff --git a/ujit_asm.c b/ujit_asm.c
index 3eea61cf69..939f238b0e 100644
--- a/ujit_asm.c
+++ b/ujit_asm.c
@@ -1609,10 +1609,8 @@ void test(codeblock_t* cb, x86opnd_t rm_opnd, x86opnd_t test_opnd)
     }
     else
     {
-        // For now, 32-bit operands only
         assert (test_opnd.num_bits == rm_opnd.num_bits);
-        assert (test_opnd.num_bits == 32);
-        cb_write_rm(cb, false, false, test_opnd, rm_opnd, 0xFF, 1, 0x85);
+        cb_write_rm(cb, false, rm_opnd.num_bits == 64, test_opnd, rm_opnd, 0xFF, 1, 0x85);
     }
 }
 
diff --git a/ujit_asm_tests.c b/ujit_asm_tests.c
index 59b98c443b..14be0017ac 100644
--- a/ujit_asm_tests.c
+++ b/ujit_asm_tests.c
@@ -334,6 +334,9 @@ void run_tests()
     cb_set_pos(cb, 0); test(cb, mem_opnd(8, RSI, 16), imm_opnd(1)); check_bytes(cb, "F6461001");
     cb_set_pos(cb, 0); test(cb, mem_opnd(8, RSI, -16), imm_opnd(1)); check_bytes(cb, "F646F001");
     cb_set_pos(cb, 0); test(cb, mem_opnd(32, RSI, 64), EAX); check_bytes(cb, "854640");
+    cb_set_pos(cb, 0); test(cb, mem_opnd(64, RDI, 42), RAX); check_bytes(cb, "4885472A");
+    cb_set_pos(cb, 0); test(cb, RAX, RAX); check_bytes(cb, "4885C0");
+    cb_set_pos(cb, 0); test(cb, RAX, RSI); check_bytes(cb, "4885F0");
     cb_set_pos(cb, 0); test(cb, mem_opnd(64, RSI, 64), imm_opnd(~0x08)); check_bytes(cb, "48F74640F7FFFFFF");
 
     // xor
diff --git a/ujit_codegen.c b/ujit_codegen.c
index e1bf8d85b7..db8b48a2e8 100644
--- a/ujit_codegen.c
+++ b/ujit_codegen.c
@@ -1132,7 +1132,7 @@ jit_protected_guard(jitstate_t *jit, codeblock_t *cb, const rb_callable_method_e
     // VALUE rb_obj_is_kind_of(VALUE obj, VALUE klass);
     call_ptr(cb, REG0, (void *)&rb_obj_is_kind_of);
     ujit_load_regs(cb);
-    cmp(cb, RAX, imm_opnd(0));
+    test(cb, RAX, RAX);
     jz_ptr(cb, COUNTED_EXIT(side_exit, oswb_se_protected_check_failed));
 }
 
@@ -1668,7 +1668,7 @@ gen_leave(jitstate_t* jit, ctx_t* ctx)
 
     // If the return address is NULL, fall back to the interpreter
     int FALLBACK_LABEL = cb_new_label(cb, "FALLBACK");
-    cmp(cb, REG1, imm_opnd(0));
+    test(cb, REG1, REG1);
     jz_label(cb, FALLBACK_LABEL);
 
     // Jump to the JIT return address