Avoid recompiling overlapping instruction sequences in ujit

compile.c

@@ -862,22 +862,30 @@ rb_iseq_translate_threaded_code(rb_iseq_t *iseq)
 {
 #if OPT_DIRECT_THREADED_CODE || OPT_CALL_THREADED_CODE
     const void * const *table = rb_vm_get_insns_address_table();
-    unsigned int i;
     VALUE *encoded = (VALUE *)iseq->body->iseq_encoded;
 
-    for (i = 0; i < iseq->body->iseq_size; /* */ )
+    unsigned int insn_idx;
+    unsigned int next_ujit_idx = 0;
+
+    bool ujit_disabled = false /*get_cmdline_flag()*/;
+
+    for (insn_idx = 0; insn_idx < iseq->body->iseq_size; /* */)
     {
-    	int insn = (int)iseq->body->iseq_encoded[i];
+    	int insn = (int)iseq->body->iseq_encoded[insn_idx];
     	int len = insn_len(insn);
 
-        uint8_t* native_code_ptr = ujit_compile_insn(iseq, i);
+        uint8_t* native_code_ptr = NULL;
+
+        // If ujit is enabled and hasn't already compiled this instruction
+        if (!ujit_disabled && insn_idx >= next_ujit_idx)
+            native_code_ptr = ujit_compile_insn(iseq, insn_idx, &next_ujit_idx);
 
         if (native_code_ptr)
-            encoded[i] = (VALUE)native_code_ptr;
+            encoded[insn_idx] = (VALUE)native_code_ptr;
         else
-            encoded[i] = (VALUE)table[insn];
+            encoded[insn_idx] = (VALUE)table[insn];
 
-    	i += len;
+    	insn_idx += len;
     }
 
     FL_SET((VALUE)iseq, ISEQ_TRANSLATED);

ujit_asm_tests.c

@@ -349,6 +349,27 @@ void run_tests()
     cb_set_pos(cb, 0); sub(cb, EAX, imm_opnd(1)); check_bytes(cb, "83E801");
     cb_set_pos(cb, 0); sub(cb, RAX, imm_opnd(2)); check_bytes(cb, "4883E802");
 
+    // test
+    /*
+    test(
+        delegate void (CodeBlock cb) { cb.instr(TEST, AL, 4); },
+        "A804"
+    );
+    test(
+        delegate void (CodeBlock cb) { cb.instr(TEST, CL, 255); },
+        "F6C1FF"
+    );
+    test(
+        delegate void (CodeBlock cb) { cb.instr(TEST, DL, 7); },
+        "F6C207"
+    );
+    test(
+        delegate void (CodeBlock cb) { cb.instr(TEST, DIL, 9); },
+        "",
+        "40F6C709"
+    );
+    */
+
     // xor
     cb_set_pos(cb, 0); xor(cb, EAX, EAX); check_bytes(cb, "31C0");
 

ujit_compile.c

@@ -92,7 +92,7 @@ System V ABI reference:
 https://wiki.osdev.org/System_V_ABI#x86-64
 */
 uint8_t *
-ujit_compile_insn(rb_iseq_t *iseq, size_t insn_idx)
+ujit_compile_insn(rb_iseq_t *iseq, unsigned int insn_idx, unsigned int* next_ujit_idx)
 {
     // If not previously done, initialize ujit
     if (!cb)
@@ -100,6 +100,9 @@ ujit_compile_insn(rb_iseq_t *iseq, size_t insn_idx)
         ujit_init();
     }
 
+    // NOTE: if we are ever deployed in production, we
+    // should probably just log an error and return NULL here,
+    // so we can fail more gracefully
     if (cb->write_pos + 1024 >= cb->mem_size)
     {
         rb_bug("out of executable memory");
@@ -124,8 +127,6 @@ ujit_compile_insn(rb_iseq_t *iseq, size_t insn_idx)
 
         // Get the current opcode
         int opcode = ctx_get_opcode(&ctx);
-        //const char* name = insn_name(insn);
-        //printf("%s\n", name);
 
         // Lookup the codegen function for this instruction
         st_data_t st_gen_fn;
@@ -150,6 +151,9 @@ ujit_compile_insn(rb_iseq_t *iseq, size_t insn_idx)
         insn_idx += insn_len(opcode);
     }
 
+    // Let the caller know how many instructions ujit compiled
+    *next_ujit_idx = insn_idx;
+
     // If no instructions were compiled
     if (num_instrs == 0)
     {
@@ -157,8 +161,7 @@ ujit_compile_insn(rb_iseq_t *iseq, size_t insn_idx)
     }
 
     // Directly return the next PC, which is a constant
-    void *next_pc = &iseq->body->iseq_encoded[insn_idx];
+    mov(cb, RAX, const_ptr_opnd(ctx.pc));
-    mov(cb, RAX, const_ptr_opnd(next_pc));
 
     // Write the post call bytes
     ujit_instr_exit(cb);

ujit_compile.h

@@ -9,6 +9,6 @@ typedef struct rb_iseq_struct rb_iseq_t;
 #define rb_iseq_t rb_iseq_t
 #endif
 
-uint8_t* ujit_compile_insn(rb_iseq_t* iseq, size_t insn_idx);
+uint8_t* ujit_compile_insn(rb_iseq_t *iseq, unsigned int insn_idx, unsigned int* next_ujit_idx);
 
 #endif