I seem to relearn blr != return once every month or so.

src/alloy/backend/ivm/ivm_function.cc

@@ -105,7 +105,7 @@ void IVMFunction::OnBreakpointHit(ThreadState* thread_state, IntCode* i) {
 
 #undef TRACE_SOURCE_OFFSET
 
-int IVMFunction::CallImpl(ThreadState* thread_state) {
+int IVMFunction::CallImpl(ThreadState* thread_state, uint64_t return_address) {
   // Setup register file on stack.
   auto stack = (IVMStack*)thread_state->backend_data();
   auto register_file = (Register*)stack->Alloc(register_count_);
@@ -122,6 +122,8 @@ int IVMFunction::CallImpl(ThreadState* thread_state) {
   ics.did_saturate = 0;
   ics.access_callbacks = thread_state->runtime()->access_callbacks();
   ics.thread_state = thread_state;
+  ics.return_address = return_address;
+  ics.call_return_address = 0;
 
   volatile int* suspend_flag_address = thread_state->suspend_flag_address();
 

src/alloy/backend/ivm/ivm_function.h

@@ -31,7 +31,8 @@ public:
 protected:
   virtual int AddBreakpointImpl(runtime::Breakpoint* breakpoint);
   virtual int RemoveBreakpointImpl(runtime::Breakpoint* breakpoint);
-  virtual int CallImpl(runtime::ThreadState* thread_state);
+  virtual int CallImpl(runtime::ThreadState* thread_state,
+                       uint64_t return_address);
 
 private:
   IntCode* GetIntCodeAtSourceOffset(uint64_t offset);

src/alloy/backend/ivm/ivm_intcode.cc

@@ -580,7 +580,9 @@ uint32_t IntCode_CALL_XX(IntCodeState& ics, const IntCode* i, uint32_t reg) {
   ics.thread_state->runtime()->ResolveFunction(symbol_info->address(), &fn);
   XEASSERTNOTNULL(fn);
   // TODO(benvanik): proper tail call support, somehow.
-  fn->Call(ics.thread_state);
+  uint64_t return_address =
+      (i->flags & CALL_TAIL) ? ics.return_address : ics.call_return_address;
+  fn->Call(ics.thread_state, return_address);
   if (i->flags & CALL_TAIL) {
     return IA_RETURN;
   }
@@ -645,12 +647,21 @@ int Translate_CALL_TRUE(TranslationContext& ctx, Instr* i) {
 uint32_t IntCode_CALL_INDIRECT_XX(IntCodeState& ics, const IntCode* i, uint32_t reg) {
   uint64_t target = ics.rf[reg].u32;
 
+  // Check if return address - if so, return.
+  if (i->flags & CALL_POSSIBLE_RETURN) {
+    if (target == ics.return_address) {
+      return IA_RETURN;
+    }
+  }
+
   // Real call.
   Function* fn = NULL;
   ics.thread_state->runtime()->ResolveFunction(target, &fn);
   XEASSERTNOTNULL(fn);
   // TODO(benvanik): proper tail call support, somehow.
-  fn->Call(ics.thread_state);
+  uint64_t return_address =
+      (i->flags & CALL_TAIL) ? ics.return_address : ics.call_return_address;
+  fn->Call(ics.thread_state, return_address);
   if (i->flags & CALL_TAIL) {
     return IA_RETURN;
   }
@@ -775,6 +786,14 @@ int Translate_RETURN_TRUE(TranslationContext& ctx, Instr* i) {
   return DispatchToC(ctx, i, fns[i->src1.value->type]);
 }
 
+uint32_t IntCode_SET_RETURN_ADDRESS(IntCodeState& ics, const IntCode* i) {
+  ics.call_return_address = ics.rf[i->src1_reg].u32;
+  return IA_NEXT;
+}
+int Translate_SET_RETURN_ADDRESS(TranslationContext& ctx, Instr* i) {
+  return DispatchToC(ctx, i, IntCode_SET_RETURN_ADDRESS);
+}
+
 uint32_t IntCode_BRANCH_XX(IntCodeState& ics, const IntCode* i, uint32_t reg) {
   return ics.rf[reg].u32;
 }
@@ -4101,6 +4120,7 @@ static const TranslateFn dispatch_table[] = {
   Translate_CALL_EXTERN,
   Translate_RETURN,
   Translate_RETURN_TRUE,
+  Translate_SET_RETURN_ADDRESS,
 
   Translate_BRANCH,
   Translate_BRANCH_TRUE,

src/alloy/backend/ivm/ivm_intcode.h

@@ -48,6 +48,8 @@ typedef struct {
   int8_t        did_saturate;
   runtime::RegisterAccessCallbacks* access_callbacks;
   runtime::ThreadState* thread_state;
+  uint64_t      return_address;
+  uint64_t      call_return_address;
 } IntCodeState;
 
 

src/alloy/backend/x64/x64_function.cc

@@ -42,7 +42,7 @@ int X64Function::RemoveBreakpointImpl(Breakpoint* breakpoint) {
   return 0;
 }
 
-int X64Function::CallImpl(ThreadState* thread_state) {
+int X64Function::CallImpl(ThreadState* thread_state, uint64_t return_address) {
   auto backend = (X64Backend*)thread_state->runtime()->backend();
   auto thunk = backend->host_to_guest_thunk();
   thunk(

src/alloy/backend/x64/x64_function.h

@@ -33,7 +33,8 @@ public:
 protected:
   virtual int AddBreakpointImpl(runtime::Breakpoint* breakpoint);
   virtual int RemoveBreakpointImpl(runtime::Breakpoint* breakpoint);
-  virtual int CallImpl(runtime::ThreadState* thread_state);
+  virtual int CallImpl(runtime::ThreadState* thread_state,
+                       uint64_t return_address);
 
 private:
   void*   machine_code_;