From cf4021ca780d460dc351b30260532c158a5c8a5e Mon Sep 17 00:00:00 2001
From: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>
Date: Wed, 20 Jan 2021 12:44:24 -0500
Subject: [PATCH] Take VM lock in branch_stub_hit(), fix ractor deadlock.

---
 ujit.h         |  1 +
 ujit_codegen.c |  1 -
 ujit_core.c    | 11 +++++++++--
 ujit_core.h    |  8 ++++++++
 4 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/ujit.h b/ujit.h
index a4a7437ada..04c8dcae29 100644
--- a/ujit.h
+++ b/ujit.h
@@ -41,6 +41,7 @@ bool rb_ujit_enabled_p(void)
     return rb_ujit_enabled;
 }
 
+// Threshold==1 means compile on first execution
 #define UJIT_CALL_THRESHOLD (2u)
 
 void rb_ujit_method_lookup_change(VALUE cme_or_cc);
diff --git a/ujit_codegen.c b/ujit_codegen.c
index 8db6942de1..4e02e1f43f 100644
--- a/ujit_codegen.c
+++ b/ujit_codegen.c
@@ -912,7 +912,6 @@ gen_opt_send_without_block(jitstate_t* jit, ctx_t* ctx)
     for (int32_t i = 0; i < argc + 1; ++i)
     {
         x86opnd_t stack_opnd = mem_opnd(64, RAX, -(argc + 1 - i) * 8);
-        //print_ptr(cb, stack_opnd);
         x86opnd_t c_arg_reg = C_ARG_REGS[i];
         mov(cb, c_arg_reg, stack_opnd);
     }
diff --git a/ujit_core.c b/ujit_core.c
index 6f51440b86..2ca4550a71 100644
--- a/ujit_core.c
+++ b/ujit_core.c
@@ -3,6 +3,7 @@
 #include "builtin.h"
 #include "insns.inc"
 #include "insns_info.inc"
+#include "vm_sync.h"
 #include "ujit_asm.h"
 #include "ujit_utils.h"
 #include "ujit_iface.h"
@@ -162,7 +163,7 @@ uint8_t* gen_entry_point(const rb_iseq_t *iseq, uint32_t insn_idx)
 {
     // The entry context makes no assumptions about types
     blockid_t blockid = { iseq, insn_idx };
-    ctx_t ctx = { 0 };
+    ctx_t ctx = { { 0 }, 0 };
 
     // Write the interpreter entry prologue
     uint8_t* code_ptr = ujit_entry_prologue();
@@ -183,6 +184,10 @@ uint8_t* gen_entry_point(const rb_iseq_t *iseq, uint32_t insn_idx)
 // Triggers compilation of branches and code patching
 uint8_t* branch_stub_hit(uint32_t branch_idx, uint32_t target_idx)
 {
+    uint8_t* dst_addr;
+
+    RB_VM_LOCK_ENTER();
+
     assert (branch_idx < num_branches);
     assert (target_idx < 2);
     branch_t *branch = &branch_entries[branch_idx];
@@ -217,7 +222,7 @@ uint8_t* branch_stub_hit(uint32_t branch_idx, uint32_t target_idx)
     add_incoming(p_block, branch_idx);
 
     // Update the branch target address
-    uint8_t* dst_addr = cb_get_ptr(cb, p_block->start_pos);
+    dst_addr = cb_get_ptr(cb, p_block->start_pos);
     branch->dst_addrs[target_idx] = dst_addr;
 
     // Rewrite the branch with the new jump target address
@@ -230,6 +235,8 @@ uint8_t* branch_stub_hit(uint32_t branch_idx, uint32_t target_idx)
     branch->end_pos = cb->write_pos;
     cb_set_pos(cb, cur_pos);
 
+    RB_VM_LOCK_LEAVE();
+
     // Return a pointer to the compiled block version
     return dst_addr;
 }
diff --git a/ujit_core.h b/ujit_core.h
index c99bd808a0..dc94cd2873 100644
--- a/ujit_core.h
+++ b/ujit_core.h
@@ -20,9 +20,17 @@
 // Maximum number of versions per block
 #define MAX_VERSIONS 5
 
+// Maximum number of temp value types we keep track of
+#define MAX_TEMP_TYPES 8
+
 // Code generation context
 typedef struct CtxStruct
 {
+    // Temporary variable types we keep track of
+    // Values are `ruby_value_type`
+    // T_NONE==0 is the unknown type
+    uint8_t temp_types[MAX_TEMP_TYPES];
+
     // Number of values pushed on the temporary stack
     uint32_t stack_size;