add disabling branches, used for yield since their control flow is sc…

…rewy

src/jit_internal.h

@@ -62,6 +62,8 @@ namespace redmagic {
     void *temp_disable(void *resume_pc);
     void *temp_enable(void *resume_pc);
 
+    void disable_branch(void *id);
+
     void* is_traced_call();
 
     void do_not_trace_method(void *addr);
@@ -83,6 +85,7 @@ namespace redmagic {
       int count = 0;
       Tracer *tracer = nullptr;
       void *starting_point = nullptr;
+      bool disabled = false;
     };
 
     std::unordered_map<void*, branch_info> branches;

src/manager.cc

@@ -92,6 +92,10 @@ extern "C" void redmagic_do_not_trace_function(void *function_pointer) {
   manager->do_not_trace_method(function_pointer);
 }
 
+extern "C" void redmagic_disable_branch(void *id) {
+  manager->disable_branch(id);
+}
+
 static const char *avoid_inlining_methods[] = {
   // inlining the allocator doesn't really help since it will have a lot of branching
   // in trying to find where there is open memory
@@ -163,6 +167,11 @@ uint32_t Manager::get_thread_id() {
 }
 
 void* Manager::begin_trace(void *id, void *ret_addr) {
+
+  branch_info *info = &branches[id];
+  if(info->disabled)
+    return NULL; // do not trace this loop
+
   auto old_head = get_tracer_head();
 
   void *ret;
@@ -181,7 +190,6 @@ void* Manager::begin_trace(void *id, void *ret_addr) {
     auto new_head = push_tracer_stack();
 
     //assert(tracer == nullptr);
-    branch_info *info = &branches[id];
     assert(info->tracer == nullptr);
 
     auto buff = make_shared<CodeBuffer>(4 * 1024 * 1024);
@@ -210,6 +218,7 @@ void* Manager::end_trace(void *id) {
   Tracer *l;
   auto head = get_tracer_head();
   branch_info *info = &branches[id];
+  assert(!info->disabled);
   assert(head->trace_id == id);
   assert(head->tracer == info->tracer);
   l = head->tracer;
@@ -273,6 +282,11 @@ void* Manager::backwards_branch(void *id, void *ret_addr) {
   } else {
     branch_info *info = &branches[id];
 
+    if(info->disabled) {
+      // this loop is disabled
+      return NULL;
+    }
+
     if(info->tracer != nullptr) {
       // there is already a tracer with some other thread running on this item
       // which means that we are going to wait for that thread to finish its trace
@@ -379,6 +393,13 @@ void* Manager::is_traced_call() {
   return NULL;
 }
 
+void Manager::disable_branch(void *id) {
+  branches[id].disabled = true;
+  for(int i = 0; i < threadl_tracer_stack.size(); i++) {
+    assert(threadl_tracer_stack[i].trace_id != id);
+  }
+}
+
 uint32_t Manager::tracer_stack_size() {
   return threadl_tracer_stack.size();
 }

src/redmagic.h

@@ -35,10 +35,12 @@ void redmagic_ensure_not_traced(void);
 void redmagic_temp_disable(void);
 void redmagic_temp_enable(void);
 
-
 // for setup to tell the jit to not trace some function, eg the call into a gc or some custom malloc
 void redmagic_do_not_trace_function(void* function_pointer);
 
+// disable traceing on a specific branch, eg if there is some disallowed instruction such as yield contained in the branch
+void redmagic_disable_branch(void *);
+
 // return 0 if not traced, non zero otherwise
 unsigned long redmagic_is_traced(void);
 

src/tracer.cc

@@ -1192,6 +1192,8 @@ void Tracer::replace_rip_instruction() {
   enum ud_mnemonic_code mnem = ud_insn_mnemonic(&disassm);
 
   AlignedInstructions ali(&disassm);
+  uint64_t used_registers = ali.registers_used();
+
 
   switch(mnem) {
   case UD_Ilea: {
@@ -1215,12 +1217,16 @@ void Tracer::replace_rip_instruction() {
     assert(0);
   }
 
-  case UD_Imov /*... UD_Imovzx*/: {
+
+    /*... UD_Imovzx:*/
+  case UD_Imov:
+  case UD_Imovsxd: {
     const ud_operand_t *opr1 = ud_insn_opr(&disassm, 0); // dest address
     const ud_operand_t *opr2 = ud_insn_opr(&disassm, 1); // source address
     assert(opr1 != NULL && opr2 != NULL);
-    if(opr2->base == UD_R_RIP && opr2->index == UD_NONE && opr1->type == UD_OP_REG) {
-      // // then we are just reading some offset from this address
+    if(opr2->base == UD_R_RIP && opr2->index == UD_NONE && opr1->type == UD_OP_REG && ali.get_op(0)->register_i.size == 64) {
+      // then we are just reading some offset from this address, which is constant
+      // also we have a register which we are about to clobber, so we can use it to store the address instead of allocating a new register
       opr_value val = get_opr_value(opr2);
       assert(val.is_ptr);
       // assert(opr1->type == UD_OP_REG);
@@ -1230,7 +1236,7 @@ void Tracer::replace_rip_instruction() {
       auto r = compiler.get_register(dest);
       compiler.mov(r, asmjit::imm_u(val.address));
       // use the aligned instruction so that the correct size is read from memory
-      compiler.emit(asmjit::kX86InstIdMov, ali.get_asm_op(0), asmjit::x86::ptr(r));
+      compiler.emit(ali.get_asm_mnem(), ali.get_asm_op(0), asmjit::x86::ptr(r));
 
       return;
     }
@@ -1259,12 +1265,19 @@ void Tracer::replace_rip_instruction() {
   case UD_Iadd:
   case UD_Isub:
   case UD_Icmp:
+  case UD_Itest:
   case UD_Ixor:
+  case UD_Idec:
+  case UD_Iinc:
+
+
+
+  case UD_Icmpxchg: // todo: don't want this here....
+
   _auto_rewrite_register:
     {
       // automatically rewrite the registers that are being used
       // and then use the compiler to generate the approperate bytes
-      uint64_t used_registers = ali.registers_used();
       assert((used_registers & (1 << RIP)) != 0);
       assert((used_registers & (1 << RSP)) == 0);
       used_registers &= ~(1 << RIP);