Fixing tail calls in the jit.

src/alloy/backend/x64/lowering/lowering_sequences.cc

@@ -245,17 +245,30 @@ void IssueCall(X64Emitter& e, FunctionInfo* symbol_info, uint32_t flags) {
 
   // Actually jump/call to rax.
   if (flags & CALL_TAIL) {
+    // Pass the callers return address over.
+    e.mov(e.rdx, e.qword[e.rsp + StackLayout::GUEST_RET_ADDR]);
+
     e.add(e.rsp, (uint32_t)e.stack_size());
     e.jmp(e.rax);
   } else {
+    // Return address is from the previous SET_RETURN_ADDRESS.
+    e.mov(e.rdx, e.qword[e.rsp + StackLayout::GUEST_CALL_RET_ADDR]);
+
     e.call(e.rax);
   }
 }
 void IssueCallIndirect(X64Emitter& e, Value* target, uint32_t flags) {
-  // Resolve address to the function to call and store in rax.
-  // TODO(benvanik): caching/etc. For now this makes debugging easier.
   Reg64 r;
   e.BeginOp(target, r, 0);
+
+  // Check if return.
+  if (flags & CALL_POSSIBLE_RETURN) {
+    e.cmp(r.cvt32(), e.dword[e.rsp + StackLayout::GUEST_RET_ADDR]);
+    e.je("epilog", CodeGenerator::T_NEAR);
+  }
+
+  // Resolve address to the function to call and store in rax.
+  // TODO(benvanik): caching/etc. For now this makes debugging easier.
   if (r != e.rdx) {
     e.mov(e.rdx, r);
   }
@@ -264,9 +277,15 @@ void IssueCallIndirect(X64Emitter& e, Value* target, uint32_t flags) {
 
   // Actually jump/call to rax.
   if (flags & CALL_TAIL) {
+    // Pass the callers return address over.
+    e.mov(e.rdx, e.qword[e.rsp + StackLayout::GUEST_RET_ADDR]);
+
     e.add(e.rsp, (uint32_t)e.stack_size());
     e.jmp(e.rax);
   } else {
+    // Return address is from the previous SET_RETURN_ADDRESS.
+    e.mov(e.rdx, e.qword[e.rsp + StackLayout::GUEST_CALL_RET_ADDR]);
+
     e.call(e.rax);
   }
 }
@@ -434,6 +453,14 @@ table->AddSequence(OPCODE_RETURN_TRUE, [](X64Emitter& e, Instr*& i) {
   return true;
 });
 
+table->AddSequence(OPCODE_SET_RETURN_ADDRESS, [](X64Emitter& e, Instr*& i) {
+  XEASSERT(i->src1.value->IsConstant());
+  e.mov(e.qword[e.rsp + StackLayout::GUEST_CALL_RET_ADDR],
+        i->src1.value->AsUint64());
+  i = e.Advance(i);
+  return true;
+});
+
 // --------------------------------------------------------------------------
 // Branches
 // --------------------------------------------------------------------------

src/alloy/backend/x64/lowering/tracers.cc

@@ -29,10 +29,12 @@ namespace lowering {
 #define DFLUSH()
 #define DPRINT
 
+#define TARGET_THREAD 1
+
 #define IFLUSH() fflush(stdout)
-#define IPRINT if (thread_state->thread_id() == 1) printf
+#define IPRINT if (thread_state->thread_id() == TARGET_THREAD) printf
 #define DFLUSH() fflush(stdout)
-#define DPRINT DFLUSH(); if (thread_state->thread_id() == 1) printf
+#define DPRINT DFLUSH(); if (thread_state->thread_id() == TARGET_THREAD) printf
 
 
 void TraceString(void* raw_context, const char* str) {

src/alloy/backend/x64/x64_emitter.cc

@@ -158,6 +158,10 @@ int X64Emitter::Emit(HIRBuilder* builder, size_t& out_stack_size) {
   if (emit_prolog) {
     sub(rsp, (uint32_t)stack_size);
     mov(qword[rsp + StackLayout::GUEST_RCX_HOME], rcx);
+    mov(qword[rsp + StackLayout::GUEST_RET_ADDR], rdx);
+    mov(qword[rsp + StackLayout::GUEST_CALL_RET_ADDR], 0);
+    // ReloadRDX:
+    mov(rdx, qword[rcx + 8]); // membase
   }
 
   auto lowering_table = backend_->lowering_table();

src/alloy/backend/x64/x64_function.cc

@@ -48,6 +48,6 @@ int X64Function::CallImpl(ThreadState* thread_state, uint64_t return_address) {
   thunk(
       machine_code_,
       thread_state->raw_context(),
-      thread_state->memory()->membase());
+      (void*)return_address);
   return 0;
 }

src/alloy/backend/x64/x64_thunk_emitter.cc

@@ -34,6 +34,7 @@ HostToGuestThunk X64ThunkEmitter::EmitHostToGuestThunk() {
 
   const size_t stack_size = StackLayout::THUNK_STACK_SIZE;
   // rsp + 0 = return address
+  mov(qword[rsp + 8 * 3], r8);
   mov(qword[rsp + 8 * 2], rdx);
   mov(qword[rsp + 8 * 1], rcx);
   sub(rsp, stack_size);
@@ -88,6 +89,7 @@ HostToGuestThunk X64ThunkEmitter::EmitHostToGuestThunk() {
   add(rsp, stack_size);
   mov(rcx, qword[rsp + 8 * 1]);
   mov(rdx, qword[rsp + 8 * 2]);
+  mov(r8, qword[rsp + 8 * 3]);
   ret();
 
   void* fn = Emplace(stack_size);

src/alloy/backend/x64/x64_thunk_emitter.h

@@ -104,6 +104,10 @@ namespace x64 {
  *  +------------------+
  *  | rcx / context    | rsp + 64
  *  +------------------+
+ *  | guest ret addr   | rsp + 72
+ *  +------------------+
+ *  | call ret addr    | rsp + 80
+ *  +------------------+
  *    ... locals ...
  *  +------------------+
  *  | (return address) |
@@ -115,8 +119,10 @@ class StackLayout {
 public:
   const static size_t THUNK_STACK_SIZE = 120;
 
-  const static size_t GUEST_STACK_SIZE = 72;
+  const static size_t GUEST_STACK_SIZE = 88;
   const static size_t GUEST_RCX_HOME = 64;
+  const static size_t GUEST_RET_ADDR = 72;
+  const static size_t GUEST_CALL_RET_ADDR = 80;
 };