From 9fef7136b7100488fff4e2fa9b1e16d7c6cd0197 Mon Sep 17 00:00:00 2001
From: Niek Bouman <niek.git@gmail.com>
Date: Sat, 16 Nov 2024 22:57:00 +0100
Subject: [PATCH] Proof of principle of generator 'send' functionality

- not production ready
- 'int' type currently hardcoded
- input_value is currently 'public'; not sure if that's desired
- copy_awaiter does not yet have support for 'send'

- credits for this approach go to 'Nicol Bolas'
https://stackoverflow.com/questions/64083104/making-python-generator-via-c20-coroutines
---
 include/seastar/coroutine/generator.hh | 24 +++++++++++++++++++-----
 tests/unit/generator_test.cc           | 21 +++++++++++++++++++++
 2 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/include/seastar/coroutine/generator.hh b/include/seastar/coroutine/generator.hh
index 8a6302d13d..5d5e227c26 100644
--- a/include/seastar/coroutine/generator.hh
+++ b/include/seastar/coroutine/generator.hh
@@ -78,7 +78,8 @@ protected:
     // pointer to the denoted object in the promise as long as the result of
     // dereferencing this pointer is convertible to the Ref type.
     std::add_pointer_t<Yielded> _value = nullptr;
-
+public:
+    int input_value;
 protected:
     std::exception_ptr _exception;
     std::coroutine_handle<> _consumer;
@@ -107,11 +108,14 @@ protected:
     class yield_awaiter final {
         generator_promise_base* _promise;
         std::coroutine_handle<> _consumer;
+        int& _ret;
+
     public:
         yield_awaiter(generator_promise_base* promise,
-                      std::coroutine_handle<> consumer) noexcept
+                      std::coroutine_handle<> consumer, int& ret) noexcept
             : _promise{promise}
             , _consumer{consumer}
+            , _ret{ret}
         {}
         bool await_ready() const noexcept {
             return false;
@@ -127,7 +131,7 @@ protected:
             }
             return _consumer;
         }
-        void await_resume() noexcept {}
+        int& await_resume() noexcept { return _ret; }
     };
 
     class copy_awaiter {
@@ -177,7 +181,7 @@ public:
 
     yield_awaiter final_suspend() noexcept {
         _value = nullptr;
-        return yield_awaiter{this, this->_consumer};
+        return yield_awaiter{this, this->_consumer, this->input_value};
     }
 
     void unhandled_exception() noexcept {
@@ -186,7 +190,7 @@ public:
 
     yield_awaiter yield_value(Yielded value) noexcept {
         this->_value = std::addressof(value);
-        return yield_awaiter{this, this->_consumer};
+        return yield_awaiter{this, this->_consumer, this->input_value};
     }
 
     copy_awaiter yield_value(const std::remove_reference_t<Yielded>& value)
@@ -418,6 +422,16 @@ public:
     [[nodiscard]] std::default_sentinel_t end() const noexcept {
         return {};
     }
+
+    void send(value_type const& input)
+    {
+        _coro.promise().input_value = input;
+    }
+
+    void send(value_type&& input)
+    {
+        _coro.promise().input_value = std::move(input);
+    }
 };
 
 template <typename Ref, typename Value>
diff --git a/tests/unit/generator_test.cc b/tests/unit/generator_test.cc
index e0fc60b951..01ca05019c 100644
--- a/tests/unit/generator_test.cc
+++ b/tests/unit/generator_test.cc
@@ -382,3 +382,24 @@ SEASTAR_TEST_CASE(test_batch_generator_convertible) {
         BOOST_REQUIRE_EQUAL(*n, expected_n++);
     }
 }
+
+coroutine::experimental::generator<int>
+bidirectional_generator() {
+    int num = 0;
+    for (;;) {
+        num = co_yield num+1;
+    }
+}
+
+SEASTAR_TEST_CASE(test_generator_input) {
+
+    auto numbers = bidirectional_generator();
+    auto n = co_await numbers.begin();
+
+    std::cout << *n << ' '; // will print 1
+
+    numbers.send(42);
+    co_await ++n;
+
+    std::cout << *n << ' '; // will print 43
+}