1 // Copyright 2021 gRPC authors.
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     http://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,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 #ifndef GRPC_SRC_CORE_LIB_PROMISE_LOOP_H
16 #define GRPC_SRC_CORE_LIB_PROMISE_LOOP_H
17 
18 #include <grpc/support/port_platform.h>
19 
20 #include <type_traits>
21 
22 #include "absl/status/status.h"
23 #include "absl/status/statusor.h"
24 #include "absl/types/variant.h"
25 
26 #include "src/core/lib/gprpp/construct_destruct.h"
27 #include "src/core/lib/promise/detail/promise_factory.h"
28 #include "src/core/lib/promise/poll.h"
29 
30 namespace grpc_core {
31 
32 // Special type - signals to loop to take another iteration, instead of
33 // finishing
34 struct Continue {};
35 
36 // Result of polling a loop promise - either Continue looping, or return a value
37 // T
38 template <typename T>
39 using LoopCtl = absl::variant<Continue, T>;
40 
41 namespace promise_detail {
42 
43 template <typename T>
44 struct LoopTraits;
45 
46 template <typename T>
47 struct LoopTraits<LoopCtl<T>> {
48   using Result = T;
49   static LoopCtl<T> ToLoopCtl(LoopCtl<T> value) { return value; }
50 };
51 
52 template <typename T>
53 struct LoopTraits<absl::StatusOr<LoopCtl<T>>> {
54   using Result = absl::StatusOr<T>;
55   static LoopCtl<Result> ToLoopCtl(absl::StatusOr<LoopCtl<T>> value) {
56     if (!value.ok()) return value.status();
57     const auto& inner = *value;
58     if (absl::holds_alternative<Continue>(inner)) return Continue{};
59     return absl::get<T>(inner);
60   }
61 };
62 
63 template <>
64 struct LoopTraits<absl::StatusOr<LoopCtl<absl::Status>>> {
65   using Result = absl::Status;
66   static LoopCtl<Result> ToLoopCtl(
67       absl::StatusOr<LoopCtl<absl::Status>> value) {
68     if (!value.ok()) return value.status();
69     const auto& inner = *value;
70     if (absl::holds_alternative<Continue>(inner)) return Continue{};
71     return absl::get<absl::Status>(inner);
72   }
73 };
74 
75 template <typename F>
76 class Loop {
77  private:
78   using Factory = promise_detail::RepeatedPromiseFactory<void, F>;
79   using PromiseType = decltype(std::declval<Factory>().Make());
80   using PromiseResult = typename PromiseType::Result;
81 
82  public:
83   using Result = typename LoopTraits<PromiseResult>::Result;
84 
85   explicit Loop(F f) : factory_(std::move(f)) {}
86   ~Loop() {
87     if (started_) Destruct(&promise_);
88   }
89 
90   Loop(Loop&& loop) noexcept : factory_(std::move(loop.factory_)) {}
91 
92   Loop(const Loop& loop) = delete;
93   Loop& operator=(const Loop& loop) = delete;
94 
95   Poll<Result> operator()() {
96     if (!started_) {
97       started_ = true;
98       Construct(&promise_, factory_.Make());
99     }
100     while (true) {
101       // Poll the inner promise.
102       auto promise_result = promise_();
103       // If it returns a value:
104       if (auto* p = promise_result.value_if_ready()) {
105         //  - then if it's Continue, destroy the promise and recreate a new one
106         //  from our factory.
107         auto lc = LoopTraits<PromiseResult>::ToLoopCtl(*p);
108         if (absl::holds_alternative<Continue>(lc)) {
109           Destruct(&promise_);
110           Construct(&promise_, factory_.Make());
111           continue;
112         }
113         //  - otherwise there's our result... return it out.
114         return absl::get<Result>(lc);
115       } else {
116         // Otherwise the inner promise was pending, so we are pending.
117         return Pending();
118       }
119     }
120   }
121 
122  private:
123   GPR_NO_UNIQUE_ADDRESS Factory factory_;
124   GPR_NO_UNIQUE_ADDRESS union {
125     GPR_NO_UNIQUE_ADDRESS PromiseType promise_;
126   };
127   bool started_ = false;
128 };
129 
130 }  // namespace promise_detail
131 
132 // Looping combinator.
133 // Expects F returns LoopCtl<T> - if it's Continue, then run the loop again -
134 // otherwise yield the returned value as the result of the loop.
135 template <typename F>
136 promise_detail::Loop<F> Loop(F f) {
137   return promise_detail::Loop<F>(std::move(f));
138 }
139 
140 }  // namespace grpc_core
141 
142 #endif  // GRPC_SRC_CORE_LIB_PROMISE_LOOP_H
143