1 // Copyright 2024 The Pigweed Authors
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License"); you may not
4 // use this file except in compliance with the License. You may obtain a copy of
5 // the License at
6 //
7 // https://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
11 // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12 // License for the specific language governing permissions and limitations under
13 // the License.
14
15 #include "pw_async2/coro.h"
16
17 #include "pw_allocator/null_allocator.h"
18 #include "pw_allocator/testing.h"
19 #include "pw_async2/dispatcher_base.h"
20 #include "pw_status/status.h"
21
22 namespace {
23
24 using ::pw::OkStatus;
25 using ::pw::Result;
26 using ::pw::Status;
27 using ::pw::allocator::Allocator;
28 using ::pw::allocator::GetNullAllocator;
29 using ::pw::allocator::test::AllocatorForTest;
30 using ::pw::async2::Context;
31 using ::pw::async2::Coro;
32 using ::pw::async2::CoroContext;
33 using ::pw::async2::Dispatcher;
34 using ::pw::async2::Pending;
35 using ::pw::async2::Poll;
36 using ::pw::async2::Ready;
37 using ::pw::async2::Task;
38 using ::pw::async2::Waker;
39
40 class ExpectCoroTask final : public Task {
41 public:
ExpectCoroTask(Coro<pw::Status> && coro)42 ExpectCoroTask(Coro<pw::Status>&& coro) : coro_(std::move(coro)) {}
43
44 private:
DoPend(Context & cx)45 Poll<> DoPend(Context& cx) final {
46 Poll<Status> result = coro_.Pend(cx);
47 if (result.IsPending()) {
48 return Pending();
49 }
50 EXPECT_EQ(*result, OkStatus());
51 return Ready();
52 }
53 Coro<pw::Status> coro_;
54 };
55
ImmediatelyReturnsFive(CoroContext &)56 Coro<Result<int>> ImmediatelyReturnsFive(CoroContext&) { co_return 5; }
57
StoresFiveThenReturns(CoroContext & coro_cx,int & out)58 Coro<Status> StoresFiveThenReturns(CoroContext& coro_cx, int& out) {
59 PW_CO_TRY_ASSIGN(out, co_await ImmediatelyReturnsFive(coro_cx));
60 co_return OkStatus();
61 }
62
63 class ObjectWithCoroMethod {
64 public:
ObjectWithCoroMethod(int x)65 ObjectWithCoroMethod(int x) : x_(x) {}
CoroMethodStoresField(CoroContext &,int & out)66 Coro<Status> CoroMethodStoresField(CoroContext&, int& out) {
67 out = x_;
68 co_return OkStatus();
69 }
70
71 private:
72 int x_;
73 };
74
TEST(CoroTest,BasicFunctionsWithoutYieldingRun)75 TEST(CoroTest, BasicFunctionsWithoutYieldingRun) {
76 AllocatorForTest<256> alloc;
77 CoroContext coro_cx(alloc);
78 int output = 0;
79 ExpectCoroTask task = StoresFiveThenReturns(coro_cx, output);
80 Dispatcher dispatcher;
81 dispatcher.Post(task);
82 EXPECT_TRUE(dispatcher.RunUntilStalled().IsReady());
83 EXPECT_EQ(output, 5);
84 }
85
TEST(CoroTest,AllocationFailureProducesInvalidCoro)86 TEST(CoroTest, AllocationFailureProducesInvalidCoro) {
87 CoroContext coro_cx(GetNullAllocator());
88 EXPECT_FALSE(ImmediatelyReturnsFive(coro_cx).IsValid());
89 int x = 0;
90 EXPECT_FALSE(StoresFiveThenReturns(coro_cx, x).IsValid());
91 }
92
TEST(CoroTest,ObjectWithCoroMethodIsCallable)93 TEST(CoroTest, ObjectWithCoroMethodIsCallable) {
94 AllocatorForTest<256> alloc;
95 CoroContext coro_cx(alloc);
96 ObjectWithCoroMethod obj(4);
97 int out = 22;
98 ExpectCoroTask task = obj.CoroMethodStoresField(coro_cx, out);
99 Dispatcher dispatcher;
100 dispatcher.Post(task);
101 EXPECT_TRUE(dispatcher.RunUntilStalled().IsReady());
102 EXPECT_EQ(out, 4);
103 }
104
105 struct MockPendable {
MockPendable__anon50e29b260111::MockPendable106 MockPendable() : poll_count(0), return_value(Pending()), last_waker() {}
107 MockPendable(const MockPendable&) = delete;
108 MockPendable& operator=(const MockPendable&) = delete;
109 MockPendable(MockPendable&&) = delete;
110 MockPendable& operator=(MockPendable&&) = delete;
111
Pend__anon50e29b260111::MockPendable112 Poll<int> Pend(Context& cx) {
113 ++poll_count;
114 PW_ASYNC_STORE_WAKER(
115 cx, last_waker, "MockPendable is waiting for last_waker");
116 return return_value;
117 }
118
119 int poll_count;
120 Poll<int> return_value;
121 Waker last_waker;
122 };
123
AddTwo(CoroContext &,MockPendable & a,MockPendable & b)124 Coro<Result<int>> AddTwo(CoroContext&, MockPendable& a, MockPendable& b) {
125 co_return co_await a + co_await b;
126 }
127
AddTwoThenStore(CoroContext & alloc,MockPendable & a,MockPendable & b,int & out)128 Coro<Status> AddTwoThenStore(CoroContext& alloc,
129 MockPendable& a,
130 MockPendable& b,
131 int& out) {
132 PW_CO_TRY_ASSIGN(out, co_await AddTwo(alloc, a, b));
133 co_return OkStatus();
134 }
135
TEST(CoroTest,AwaitMultipleAndAwakenRuns)136 TEST(CoroTest, AwaitMultipleAndAwakenRuns) {
137 AllocatorForTest<512> alloc;
138 CoroContext coro_cx(alloc);
139 MockPendable a;
140 MockPendable b;
141 int output = 0;
142 ExpectCoroTask task = AddTwoThenStore(coro_cx, a, b, output);
143 Dispatcher dispatcher;
144 dispatcher.Post(task);
145
146 EXPECT_TRUE(dispatcher.RunUntilStalled().IsPending());
147 EXPECT_EQ(a.poll_count, 1);
148 EXPECT_EQ(b.poll_count, 0);
149
150 EXPECT_TRUE(dispatcher.RunUntilStalled().IsPending());
151 EXPECT_EQ(a.poll_count, 1);
152 EXPECT_EQ(b.poll_count, 0);
153
154 int a_value = 4;
155 a.return_value = a_value;
156 std::move(a.last_waker).Wake();
157 EXPECT_TRUE(dispatcher.RunUntilStalled().IsPending());
158 EXPECT_EQ(a.poll_count, 2);
159 EXPECT_EQ(b.poll_count, 1);
160
161 int b_value = 5;
162 b.return_value = b_value;
163 std::move(b.last_waker).Wake();
164 EXPECT_TRUE(dispatcher.RunUntilStalled().IsReady());
165 EXPECT_EQ(a.poll_count, 2);
166 EXPECT_EQ(b.poll_count, 2);
167 EXPECT_EQ(output, a_value + b_value);
168 }
169
170 } // namespace
171