xref: /aosp_15_r20/external/grpc-grpc/tools/codegen/core/gen_join.py (revision cc02d7e222339f7a4f6ba5f422e6413f4bd931f2) 
1#!/usr/bin/env python3
2
3# Copyright 2023 gRPC authors.
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#     http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17import sys
18
19from mako.template import Template
20
21join_state = Template(
22    """
23template <class Traits, ${",".join(f"typename P{i}" for i in range(0,n))}>
24struct JoinState<Traits, ${",".join(f"P{i}" for i in range(0,n))}> {
25  template <typename T>
26  using UnwrappedType = decltype(Traits::Unwrapped(std::declval<T>()));
27% for i in range(0,n):
28  using Promise${i} = PromiseLike<P${i}>;
29  using Result${i} = UnwrappedType<typename Promise${i}::Result>;
30  union {
31    GPR_NO_UNIQUE_ADDRESS Promise${i} promise${i};
32    GPR_NO_UNIQUE_ADDRESS Result${i} result${i};
33  };
34% endfor
35  GPR_NO_UNIQUE_ADDRESS BitSet<${n}> ready;
36  JoinState(${",".join(f"P{i}&& p{i}" for i in range(0,n))}) {
37% for i in range(0,n):
38    Construct(&promise${i}, std::forward<P${i}>(p${i}));
39% endfor
40  }
41  JoinState(const JoinState& other) {
42    GPR_ASSERT(other.ready.none());
43% for i in range(0,n):
44    Construct(&promise${i}, other.promise${i});
45% endfor
46  }
47  JoinState& operator=(const JoinState& other) = delete;
48  JoinState& operator=(JoinState&& other) = delete;
49  JoinState(JoinState&& other) noexcept : ready(other.ready) {
50% for i in range(0,n):
51    if (ready.is_set(${i})) {
52      Construct(&result${i}, std::move(other.result${i}));
53    } else {
54      Construct(&promise${i}, std::move(other.promise${i}));
55    }
56% endfor
57  }
58  ~JoinState() {
59% for i in range(0,n):
60    if (ready.is_set(${i})) {
61      Destruct(&result${i});
62    } else {
63      Destruct(&promise${i});
64    }
65% endfor
66  }
67  using Result = typename Traits::template ResultType<std::tuple<
68      ${",".join(f"Result{i}" for i in range(0,n))}>>;
69  Poll<Result> PollOnce() {
70% for i in range(0,n):
71    if (!ready.is_set(${i})) {
72      if (grpc_trace_promise_primitives.enabled()) {
73        gpr_log(GPR_DEBUG, "join[%p]: begin poll joint ${i+1}/${n}", this);
74      }
75      auto poll = promise${i}();
76      if (grpc_trace_promise_primitives.enabled()) {
77        auto* p = poll.value_if_ready();
78        gpr_log(GPR_DEBUG, "join[%p]: joint ${i+1}/${n} %s", this,
79                p != nullptr? (Traits::IsOk(*p)? "ready" : "early-error") : "pending");
80      }
81      if (auto* p = poll.value_if_ready()) {
82        if (Traits::IsOk(*p)) {
83          ready.set(${i});
84          Destruct(&promise${i});
85          Construct(&result${i}, Traits::Unwrapped(std::move(*p)));
86        } else {
87          return Traits::template EarlyReturn<Result>(std::move(*p));
88        }
89      }
90    } else if (grpc_trace_promise_primitives.enabled()) {
91      gpr_log(GPR_DEBUG, "join[%p]: joint ${i+1}/${n} already ready", this);
92    }
93% endfor
94    if (ready.all()) {
95      return Traits::FinalReturn(${",".join(f"std::move(result{i})" for i in range(0,n))});
96    }
97    return Pending{};
98  }
99};
100"""
101)
102
103front_matter = """
104#ifndef GRPC_SRC_CORE_LIB_PROMISE_DETAIL_JOIN_STATE_H
105#define GRPC_SRC_CORE_LIB_PROMISE_DETAIL_JOIN_STATE_H
106
107// This file is generated by tools/codegen/core/gen_seq.py
108
109#include <grpc/support/port_platform.h>
110
111#include "src/core/lib/gprpp/construct_destruct.h"
112#include "src/core/lib/promise/detail/promise_like.h"
113#include "src/core/lib/promise/poll.h"
114#include "src/core/lib/gprpp/bitset.h"
115#include <grpc/support/log.h>
116#include <tuple>
117#include <type_traits>
118#include <utility>
119#include "src/core/lib/promise/trace.h"
120
121namespace grpc_core {
122namespace promise_detail {
123template <class Traits, typename... Ps>
124struct JoinState;
125"""
126
127end_matter = """
128}  // namespace promise_detail
129}  // namespace grpc_core
130
131#endif  // GRPC_SRC_CORE_LIB_PROMISE_DETAIL_JOIN_STATE_H
132"""
133
134
135# utility: print a big comment block into a set of files
136def put_banner(files, banner):
137    for f in files:
138        for line in banner:
139            print("// %s" % line, file=f)
140        print("", file=f)
141
142
143with open(sys.argv[0]) as my_source:
144    copyright = []
145    for line in my_source:
146        if line[0] != "#":
147            break
148    for line in my_source:
149        if line[0] == "#":
150            copyright.append(line)
151            break
152    for line in my_source:
153        if line[0] != "#":
154            break
155        copyright.append(line)
156
157copyright = [line[2:].rstrip() for line in copyright]
158
159with open("src/core/lib/promise/detail/join_state.h", "w") as f:
160    put_banner([f], copyright)
161    print(front_matter, file=f)
162    for n in range(2, 10):
163        print(join_state.render(n=n), file=f)
164    print(end_matter, file=f)
165