From 70cb3df628e636984afa376e34a9b7f2faa015e8 Mon Sep 17 00:00:00 2001
From: Ian Petersen <ispeters@gmail.com>
Date: Wed, 11 Sep 2024 14:05:27 -0700
Subject: [PATCH] Add more get_return_address customizations

This PR continues what was started in #635 and adds more customizations
of `get_return_address`, starting with `task<>`.
---
 include/unifex/task.hpp | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/include/unifex/task.hpp b/include/unifex/task.hpp
index 1a73373f..3878515c 100644
--- a/include/unifex/task.hpp
+++ b/include/unifex/task.hpp
@@ -320,9 +320,11 @@ struct _promise final {
     , _return_value_or_void<T, nothrow>::type {
     using result_type = T;
 
-    typename _task<T, nothrow>::type get_return_object() noexcept {
+    typename _task<T, nothrow>::type get_return_object(
+        instruction_ptr returnAddress =
+            instruction_ptr::read_return_address()) noexcept {
       return typename _task<T, nothrow>::type{
-          coro::coroutine_handle<type>::from_promise(*this)};
+          coro::coroutine_handle<type>::from_promise(*this), returnAddress};
     }
 
     auto final_suspend() noexcept {
@@ -750,8 +752,11 @@ struct _task<T, nothrow>::type
   }
 
 private:
-  explicit type(coro::coroutine_handle<promise_type> h) noexcept
-    : coro_holder(h) {}
+  explicit type(
+      coro::coroutine_handle<promise_type> h,
+      instruction_ptr returnAddress) noexcept
+    : coro_holder(h)
+    , returnAddress_(returnAddress) {}
 
   template <typename Promise>
   friend auto
@@ -789,6 +794,13 @@ struct _task<T, nothrow>::type
       tag_t<with_scheduler_affinity>, type&& task, Scheduler&&) noexcept {
     return {std::move(task)};
   }
+
+  friend instruction_ptr
+  tag_invoke(tag_t<get_return_address>, const type& task) noexcept {
+    return returnAddress_;
+  }
+
+  instruction_ptr returnAddress_;
 };
 
 /**