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_;

src/alloy/frontend/ppc/ppc_emit_control.cc

@@ -35,6 +35,7 @@ int InstrEmit_branch(
   // be correct for returns.
   if (lk) {
     Value* return_address = f.LoadConstant(cia + 4);
+    f.SetReturnAddress(return_address);
     f.StoreLR(return_address);
   }
 
@@ -104,6 +105,10 @@ int InstrEmit_branch(
     //  // TODO(benvanik): evaluate hint here.
     //  c.je(e.GetReturnLabel(), kCondHintLikely);
     //}
+#if 0
+    // This breaks longjump, as that uses blr with a non-return lr.
+    // It'd be nice to move SET_RETURN_ADDRESS semantics up into context
+    // so that we can just use this.
     if (!lk && nia_is_lr) {
       // Return (most likely).
       // TODO(benvanik): test? ReturnCheck()?
@@ -116,7 +121,14 @@ int InstrEmit_branch(
         f.Return();
       }
     } else {
+#else
+    {
+#endif
       // Jump to pointer.
+      bool likely_return = !lk && nia_is_lr;
+      if (likely_return) {
+        call_flags |= CALL_POSSIBLE_RETURN;
+      }
       if (cond) {
         if (!expect_true) {
           cond = f.IsFalse(cond);

src/alloy/hir/hir_builder.cc

@@ -642,6 +642,12 @@ void HIRBuilder::ReturnTrue(Value* cond) {
   EndBlock();
 }
 
+void HIRBuilder::SetReturnAddress(Value* value) {
+  Instr* i = AppendInstr(OPCODE_SET_RETURN_ADDRESS_info, 0);
+  i->set_src1(value);
+  i->src2.value = i->src3.value = NULL;
+}
+
 void HIRBuilder::Branch(Label* label, uint32_t branch_flags) {
   Instr* i = AppendInstr(OPCODE_BRANCH_info, branch_flags);
   i->src1.label = label;

src/alloy/hir/hir_builder.h

@@ -81,6 +81,7 @@ public:
   void CallExtern(runtime::FunctionInfo* symbol_info);
   void Return();
   void ReturnTrue(Value* cond);
+  void SetReturnAddress(Value* value);
 
   void Branch(Label* label, uint32_t branch_flags = 0);
   void Branch(Block* block, uint32_t branch_flags = 0);

src/alloy/hir/opcodes.h

@@ -18,7 +18,8 @@ namespace hir {
 
 
 enum CallFlags {
-  CALL_TAIL       = (1 << 1),
+  CALL_TAIL             = (1 << 1),
+  CALL_POSSIBLE_RETURN  = (1 << 2),
 };
 enum BranchFlags {
   BRANCH_LIKELY   = (1 << 1),
@@ -97,6 +98,7 @@ enum Opcode {
   OPCODE_CALL_EXTERN,
   OPCODE_RETURN,
   OPCODE_RETURN_TRUE,
+  OPCODE_SET_RETURN_ADDRESS,
 
   OPCODE_BRANCH,
   OPCODE_BRANCH_TRUE,

src/alloy/hir/opcodes.inl

@@ -92,6 +92,12 @@ DEFINE_OPCODE(
     OPCODE_SIG_X_V,
     OPCODE_FLAG_BRANCH);
 
+DEFINE_OPCODE(
+    OPCODE_SET_RETURN_ADDRESS,
+    "set_return_address",
+    OPCODE_SIG_X_V,
+    0);
+
 DEFINE_OPCODE(
     OPCODE_BRANCH,
     "branch",

src/alloy/runtime/function.cc

@@ -73,7 +73,7 @@ Breakpoint* Function::FindBreakpoint(uint64_t address) {
   return result;
 }
 
-int Function::Call(ThreadState* thread_state) {
+int Function::Call(ThreadState* thread_state, uint64_t return_address) {
   ThreadState* original_thread_state = ThreadState::Get();
   if (original_thread_state != thread_state) {
     ThreadState::Bind(thread_state);
@@ -94,7 +94,7 @@ int Function::Call(ThreadState* thread_state) {
       result = 1;
     }
   } else {
-    CallImpl(thread_state);
+    CallImpl(thread_state, return_address);
   }
 
   if (original_thread_state != thread_state) {

src/alloy/runtime/function.h

@@ -36,13 +36,14 @@ public:
   int AddBreakpoint(Breakpoint* breakpoint);
   int RemoveBreakpoint(Breakpoint* breakpoint);
 
-  int Call(ThreadState* thread_state);
+  int Call(ThreadState* thread_state, uint64_t return_address);
 
 protected:
   Breakpoint* FindBreakpoint(uint64_t address);
   virtual int AddBreakpointImpl(Breakpoint* breakpoint) { return 0; }
   virtual int RemoveBreakpointImpl(Breakpoint* breakpoint) { return 0; }
-  virtual int CallImpl(ThreadState* thread_state) = 0;
+  virtual int CallImpl(ThreadState* thread_state,
+                       uint64_t return_address) = 0;
 
 protected:
   uint64_t    address_;

src/xenia/cpu/processor.cc

@@ -165,7 +165,7 @@ int Processor::Execute(XenonThreadState* thread_state, uint64_t address) {
   context->lr = lr;
 
   // Execute the function.
-  fn->Call(thread_state);
+  fn->Call(thread_state, lr);
   return 0;
 }
 

tools/alloy-sandbox/alloy-sandbox.cc

@@ -49,7 +49,7 @@ int alloy_sandbox(int argc, xechar_t** argv) {
   ctx->lr = 0xBEBEBEBE;
   ctx->r[5] = 10;
   ctx->r[25] = 25;
-  fn->Call(thread_state);
+  fn->Call(thread_state, ctx->lr);
   auto result = ctx->r[11];
 
   delete thread_state;