merge block code

make

@@ -83,7 +83,7 @@ def link():
     #     **dict(globals(), **locals())
     # ))
     udis_libs = ' '.join(glob.glob('deps/udis86/libudis86/.libs/*.o'))
-    # we are not using the compiler interface, just the assembler, would be nice if we could strip all the functions
+    # we are not using the compiler interface, just the assembler
     asmjit_libs = ' '.join(filter(lambda x: 'compiler' not in x, glob.glob('build/asmjit/CMakeFiles/asmjit.dir/src/asmjit/*/*.o')))
     Run('{LD} {LD_FLAGS} -shared -fPIC -Wl,-Bsymbolic -Wl,-soname,libredmagic.so.1.0.0 -o build/libredmagic.so.1.0.0 {objs} {udis_libs} {asmjit_libs} {LIBS}'.format(
         **dict(globals(), **locals())
@@ -161,7 +161,6 @@ def unit():
     Shell('./' + UNIT_TARGET)
 
 def deps():
-    # udis86 version 1.7.2
     if not os.path.isdir('build'):
         Shell('mkdir -p build')
     if not os.path.isfile('deps/udis86/libudis86/.libs/libudis86.so') or not os.path.isfile('deps/udis86/libudis86/itab.h'):
@@ -174,7 +173,7 @@ def deps():
         cm_args = '-DASMJIT_DISABLE_COMPILER=1 -DASMJIT_CFLAGS=\'==REPLACE_ME==\' -DCMAKE_CXX_COMPILER=g++ -DCMAKE_C_COMPILER=gcc'
         if RELEASE:
             Shell('cd build/asmjit && cmake ../../deps/asmjit {} -DASMJIT_RELEASE=1'.format(cm_args), shell=True)
-            asm_flags += '\-O2'
+            asm_flags += '\-O2\ \-flto'
         else:
             Shell('cd build/asmjit && cmake ../../deps/asmjit {} -DASMJIT_DEBUG=1'.format(cm_args), shell=True)
             asm_flags += '\-ggdb'

src/align_udis_asmjit.cc

@@ -210,7 +210,7 @@ const asmjit::Operand AlignedInstructions::get_asm_op(unsigned int i) {
       int scale = 0;
 
       switch(info->index_scale) {
-      case  0: scale = 0; break;
+      case  1: scale = 0; break;
       case  2: scale = 1; break;
       case  4: scale = 2; break;
       case  8: scale = 3; break;

src/asm_interface.S

@@ -66,12 +66,15 @@ red_asm_restart_trace:
   movq %r10, %rdi
   movq %r9, %rsi
   movq %rsp, %rdx
+  movq %r8, %rcx
 
-  // restore the origional values of r14&15
+  // restore the origional values of r10,9,8
   movq 216(%rsp), %r10
   movq 224(%rsp), %r9
+  movq 232(%rsp), %r8
   movq %r10, 56(%rsp)
   movq %r9, 64(%rsp)
+  movq %r8, 72(%rsp)
 
   call red_resume_trace@plt
   m_pop_all_regs

src/asm_macros.S

@@ -76,7 +76,7 @@
   //add $16, %rax // just include this add in the offset of fxsaveq
   and $~15, %rax
 
-  fxsaveq 0x100(%rax)
+  fxsaveq 0x180(%rax)
 
   .endm
 
@@ -86,7 +86,7 @@
   // extra registers
   movq %rsp, %rax
   and $~15, %rax
-  fxrstorq 0x100(%rax)
+  fxrstorq 0x180(%rax)
 
   // eflags
   movq 144(%rsp), %r14

src/jit_internal.h

@@ -62,6 +62,9 @@ namespace redmagic {
     void *temp_disable(void *resume_pc);
     void *temp_enable(void *resume_pc);
 
+    void *begin_merge_block();
+    void *end_merge_block();
+
     void disable_branch(void *id);
 
     void* is_traced_call();
@@ -86,6 +89,9 @@ namespace redmagic {
       Tracer *tracer = nullptr;
       void *starting_point = nullptr;
       bool disabled = false;
+      int64_t traced_instruction_count = 0;
+      int sub_branches = 0;
+      uint64_t *trace_loop_counter = nullptr;
     };
 
     std::unordered_map<void*, branch_info> branches;
@@ -119,6 +125,7 @@ namespace redmagic {
   //extern thread_local bool is_traced;
   extern Manager *manager;
   extern thread_local bool protected_malloc;
+  extern std::atomic<Tracer*> free_tracer_list;
 
 
   class CodeBuffer final {

src/main.cc

@@ -15,6 +15,7 @@ enum {
   COND_JUMP,
   COND_SET,
   COND_ADD,
+  MERGE_TEST,
 };
 
 int main(int argc, char* argv[]) {
@@ -36,6 +37,7 @@ int main(int argc, char* argv[]) {
     COND_SET, 100,
     PRINT_NUM, 1,
     COND_ADD, -1,
+    MERGE_TEST, 0,
     COND_JUMP, 4,  // print num
     EXIT,
   };
@@ -82,6 +84,16 @@ int main(int argc, char* argv[]) {
       cond_var += program[pc + 1];
       pc += 2;
       break;
+    case MERGE_TEST:
+      redmagic_begin_merge_block();
+      if(cond_var % 2 == 0) {
+        cout << "a\n";
+      } else {
+        cout << "b\n";
+      }
+      pc += 2;
+      redmagic_end_merge_block();
+      break;
     case EXIT:
       cout << "hitting normal exit case\n" << a;
       exit(0);

src/manager.cc

@@ -98,6 +98,16 @@ extern "C" void* red_user_temp_enable(void *_, void *ret_addr) {
   //return NULL;
 }
 
+extern "C" void* red_user_begin_merge_block(void *_, void *ret_addr) {
+  UnprotectMalloc upm;
+  return manager->begin_merge_block();
+}
+
+extern "C" void* red_user_end_merge_block(void *_, void *ret_addr) {
+  UnprotectMalloc upm;
+  return manager->end_merge_block();
+}
+
 extern "C" void *__real_malloc(size_t);
 
 extern "C" void redmagic_start() {
@@ -166,7 +176,10 @@ static const char *avoid_inlining_methods[] = {
   "redmagic_temp_enable",
   "redmagic_is_traced",
   "redmagic_disable_branch",
-  "redmagic_do_not_trace_function"
+  "redmagic_do_not_trace_function",
+  "redmagic_begin_merge_block",
+  "redmagic_end_merge_block",
+  NULL,
 };
 
 // namespace redmagic {
@@ -249,6 +262,7 @@ void* Manager::begin_trace(void *id, void *ret_addr) {
     ret = l->Start(trace_pc);
     new_head->is_traced = true;
     info->starting_point = l->get_start_location();
+    info->trace_loop_counter = l->get_loop_counter();
   }
   //return NULL;
   return ret;
@@ -461,6 +475,22 @@ void* Manager::temp_enable(void *resume_pc) {
   // return NULL;
 }
 
+void* Manager::begin_merge_block() {
+  auto head = get_tracer_head();
+  if(head->tracer) {
+    return head->tracer->BeginMergeBlock();
+  }
+  return NULL;
+}
+
+void* Manager::end_merge_block() {
+  auto head = get_tracer_head();
+  if(head->tracer) {
+    return head->tracer->EndMergeBlock();
+  }
+  return NULL;
+}
+
 void* Manager::is_traced_call() {
   auto head = get_tracer_head();
   if(head->is_traced) {

src/redmagic.h

@@ -44,6 +44,11 @@ void redmagic_disable_branch(void *);
 // return 0 if not traced, non zero otherwise
 unsigned long redmagic_is_traced(void);
 
+// A merge block allows you to have some section of code that performs branches but at the end of the block merge back into the same instruction stream
+// useful for things like reference counting where hitting zero will not have a significant impact on the following control flow
+void redmagic_begin_merge_block(void);
+void redmagic_end_merge_block(void);
+
 #ifdef __cplusplus
 }
 

src/simple_compiler.cc

@@ -150,7 +150,7 @@ void SimpleCompiler::SetRegister(int reg, register_t val) {
   mov(get_register(reg), imm_u(val));
 }
 
-CodeBuffer SimpleCompiler::TestRegister(mem_loc_t resume_pc, int reg, register_t val) {
+CodeBuffer SimpleCompiler::TestRegister(mem_loc_t resume_pc, int reg, register_t val, mem_loc_t *merge_addr) {
   auto r = get_register(reg);
   asmjit::X86GpReg scr;
   if(val > 0x7fffffff)
@@ -165,6 +165,8 @@ CodeBuffer SimpleCompiler::TestRegister(mem_loc_t resume_pc, int reg, register_t
   auto resume_cb = resume_block.finalize_bottom();
   set_label_address(failure, resume_cb.getRawBuffer());
 
+  do_merge_addr(resume_cb, merge_addr);
+
   pushf();
   if(val > 0x7fffffff) {
     mov(scr, imm_u(val));
@@ -181,7 +183,7 @@ CodeBuffer SimpleCompiler::TestRegister(mem_loc_t resume_pc, int reg, register_t
   return resume_cb;
 }
 
-CodeBuffer SimpleCompiler::TestMemoryLocation(mem_loc_t resume_pc, mem_loc_t where, register_t val) {
+CodeBuffer SimpleCompiler::TestMemoryLocation(mem_loc_t resume_pc, mem_loc_t where, register_t val, mem_loc_t *merge_addr) {
   //Label success = newLabel();
 
   Label failure = newLabel();
@@ -198,6 +200,8 @@ CodeBuffer SimpleCompiler::TestMemoryLocation(mem_loc_t resume_pc, mem_loc_t whe
   auto resume_cb = resume_block.finalize_bottom();
   set_label_address(failure, resume_cb.getRawBuffer());
 
+  do_merge_addr(resume_cb, merge_addr);
+
   pushf();
   mov(scr, imm_ptr(where));
   if(val > 0x7fffffff) {
@@ -214,7 +218,7 @@ CodeBuffer SimpleCompiler::TestMemoryLocation(mem_loc_t resume_pc, mem_loc_t whe
   return resume_cb;
 }
 
-CodeBuffer SimpleCompiler::TestOperand(mem_loc_t resume_pc, const asmjit::Operand& opr, register_t val) {
+CodeBuffer SimpleCompiler::TestOperand(mem_loc_t resume_pc, const asmjit::Operand& opr, register_t val, mem_loc_t *merge_addr) {
   Label failure = newLabel();
   asmjit::X86GpReg scr, scr2;
   scr = get_scratch_register();
@@ -231,6 +235,7 @@ CodeBuffer SimpleCompiler::TestOperand(mem_loc_t resume_pc, const asmjit::Operan
 
   auto resume_cb = resume_block.finalize_bottom();
   set_label_address(failure, resume_cb.getRawBuffer());
+  do_merge_addr(resume_cb, merge_addr);
 
   pushf();
   // mov(src, opr);
@@ -262,8 +267,11 @@ uint64_t* SimpleCompiler::MakeCounter() {
   // mov(x86::ptr(label), scr);
 
   // TODO: make this add to this location
-  add(x86::ptr(label), 1);
-  assert(0);
+  //auto op = x86::ptr_abs((Ptr)cptr);
+  auto op = x86::ptr(label);
+  op.setSize(8);
+  add(op, 1);
+  //assert(0);
 
   return cptr;
 }
@@ -303,8 +311,8 @@ void SimpleCompiler::Push64bitValue(uint64_t value) {
 
 }
 
-CodeBuffer SimpleCompiler::ConditionalJump(mem_loc_t resume_pc, enum asmjit::X86InstId mnem) {
-  CodeBuffer rb = MakeResumeTraceBlock(resume_pc);
+CodeBuffer SimpleCompiler::ConditionalJump(mem_loc_t resume_pc, enum asmjit::X86InstId mnem, mem_loc_t *merge_addr) {
+  CodeBuffer rb = MakeResumeTraceBlock(resume_pc, merge_addr);
   auto label = newLabel();
   set_label_address(label, rb.getRawBuffer());
   emit(mnem, label);
@@ -314,11 +322,20 @@ CodeBuffer SimpleCompiler::ConditionalJump(mem_loc_t resume_pc, enum asmjit::X86
 
 extern "C" void red_asm_restart_trace();
 
-CodeBuffer SimpleCompiler::MakeResumeTraceBlock(mem_loc_t resume_pc) {
+CodeBuffer SimpleCompiler::MakeResumeTraceBlock(mem_loc_t resume_pc, mem_loc_t *merge_addr) {
   SimpleCompiler resume_block(buffer);
   //resume_block.clobbered_registers |= clobbered_registers;
   resume_block.ResumeBlockJump(resume_pc);
-  return resume_block.finalize_bottom();
+  auto ret = resume_block.finalize_bottom();
+  do_merge_addr(ret, merge_addr);
+  return ret;
+}
+
+void SimpleCompiler::do_merge_addr(CodeBuffer &buff, mem_loc_t *merge_addr) {
+  mem_loc_t* ma = buff.find_stump<mem_loc_t>(0xfbfbfbfbfbfbfbfb);
+  // make a linked list out of address where we are going to write this merge address
+  *ma = *merge_addr;
+  *merge_addr = (mem_loc_t)ma;
 }
 
 void SimpleCompiler::ResumeBlockJump(mem_loc_t resume_pc) {
@@ -330,12 +347,14 @@ void SimpleCompiler::ResumeBlockJump(mem_loc_t resume_pc) {
   bind(label_top);
   mov(x86::ptr(x86::rsp, -TRACE_STACK_OFFSET + 216), x86::r10);
   mov(x86::ptr(x86::rsp, -TRACE_STACK_OFFSET + 224), x86::r9);
+  mov(x86::ptr(x86::rsp, -TRACE_STACK_OFFSET + 232), x86::r8);
   mov(x86::r10, imm_u(resume_pc));
   // TODO: have this load the address of the instruction that jumped here instead of just this block
   // this will allow for it to easily write in a direct jump, as being designed now, we will have to redirect the jump through this indirection
   // so first conditional jump followed by direct jump
   // also, this will not work with concurrent threads
   lea(x86::r9, x86::ptr(label));
+  mov(x86::r8, imm_u(0xfbfbfbfbfbfbfbfb));
   jmp(imm_ptr(&red_asm_restart_trace));
 
   // for identifying which instruction it jumped from

src/simple_compiler.h

@@ -45,17 +45,17 @@ namespace redmagic {
     void PushMemoryLocationValue(mem_loc_t where);
     void Push64bitValue(uint64_t value);
 
-    CodeBuffer TestRegister(mem_loc_t resume_pc, int reg, register_t val);
-    CodeBuffer TestMemoryLocation(mem_loc_t resume_pc, mem_loc_t where, register_t val);
+    CodeBuffer TestRegister(mem_loc_t resume_pc, int reg, register_t val, mem_loc_t *merge_addr);
+    CodeBuffer TestMemoryLocation(mem_loc_t resume_pc, mem_loc_t where, register_t val, mem_loc_t *merge_addr);
 
-    CodeBuffer TestOperand(mem_loc_t resume_pc, const asmjit::Operand& opr, register_t val);
+    CodeBuffer TestOperand(mem_loc_t resume_pc, const asmjit::Operand& opr, register_t val, mem_loc_t *merge_addr);
 
-    CodeBuffer MakeResumeTraceBlock(mem_loc_t resume_pc);
+    CodeBuffer MakeResumeTraceBlock(mem_loc_t resume_pc, mem_loc_t *merge_addr);
 
     uint64_t* MakeCounter();
 
     // return the CodeBuffer of the destination jump block instead of self
-    CodeBuffer ConditionalJump(mem_loc_t resume_pc, enum asmjit::X86InstId mnem);
+    CodeBuffer ConditionalJump(mem_loc_t resume_pc, enum asmjit::X86InstId mnem, mem_loc_t *merge_addr);
 
   private:
     void ResumeBlockJump(mem_loc_t resume_pc);
@@ -79,6 +79,8 @@ namespace redmagic {
 
     void set_label_address(const asmjit::Label &label, mem_loc_t addr);
 
+    void do_merge_addr(CodeBuffer &buff, mem_loc_t *merge_addr);
+
   public:
     virtual size_t _relocCode(void *_dst, asmjit::Ptr baseAddress) const noexcept override;