Skip to content

Commit 921bb5c

Browse files
EnjectionEnjection
authored
Merge afa9d26 into 1717439 
2 parents 1717439 + afa9d26 commit 921bb5c

25 files changed

+917
-59
lines changed

‎ydb/core/config/protos/marker.proto

+16
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
import "google/protobuf/descriptor.proto";
2+
3+
package NKikimrConfig.NMarkers;
4+
option java_package = "ru.yandex.kikimr.proto.markers";
5+
6+
extend google.protobuf.MessageOptions {
7+
optional bool Root = 81001;
8+
optional bool CombinedType = 81002;
9+
optional bool WithMapType = 81003;
10+
}
11+
12+
extend google.protobuf.FieldOptions {
13+
repeated string CopyTo = 82001;
14+
repeated string AsMap = 82002;
15+
}
16+

‎ydb/core/config/protos/ya.make

+9
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
PROTO_LIBRARY()
2+
3+
SRCS(
4+
marker.proto
5+
)
6+
7+
EXCLUDE_TAGS(GO_PROTO)
8+
9+
END()

‎ydb/core/config/tools/protobuf_plugin/main.cpp

+298
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,298 @@
1+
#include <map>
2+
3+
#include <util/system/compiler.h>
4+
#include <util/generic/map.h>
5+
#include <util/generic/set.h>
6+
#include <util/generic/maybe.h>
7+
#include <util/generic/ptr.h>
8+
#include <util/generic/vector.h>
9+
#include <util/string/builder.h>
10+
#include <util/string/cast.h>
11+
#include <util/string/printf.h>
12+
#include <util/string/subst.h>
13+
14+
#include <google/protobuf/compiler/code_generator.h>
15+
#include <google/protobuf/compiler/plugin.h>
16+
#include <google/protobuf/io/printer.h>
17+
#include <google/protobuf/io/zero_copy_stream.h>
18+
#include <ydb/core/config/protos/marker.pb.h>
19+
20+
#include <ydb/public/lib/protobuf/base_message_generator.h>
21+
#include <ydb/public/lib/protobuf/helpers.h>
22+
#include <ydb/public/lib/protobuf/macro.h>
23+
24+
using namespace google::protobuf::compiler;
25+
using namespace google::protobuf;
26+
using namespace NKikimr::NProtobuf;
27+
28+
constexpr TStringBuf PLUGIN_NAME = "config";
29+
30+
class TMessageGenerator
31+
: public TBaseMessageGenerator
32+
{
33+
private:
34+
35+
void GenerateConfigRoot(TVars vars) {
36+
for (auto i = 0; i < Message->field_count(); ++i) {
37+
const FieldDescriptor* field = Message->field(i);
38+
if (field->is_repeated()) {
39+
continue;
40+
}
41+
if (auto* fieldMessage = field->message_type()) {
42+
vars["field"] = field->name();
43+
vars["fqFieldClass"] = FullyQualifiedClassName(fieldMessage);
44+
vars["fieldNumber"] = std::to_string(field->number());
45+
46+
WITH_PLUGIN_MARKUP(Header, PLUGIN_NAME) {
47+
Header->Print(vars, "struct T$field$FieldTag {};\n");
48+
49+
Header->Print(vars, "constexpr inline static std::tuple<\n");
50+
WITH_INDENT(Header) {
51+
Header->Print(vars,
52+
"bool ($fqMessageClass$::*)() const,\n"
53+
"const $fqFieldClass$& ($fqMessageClass$::*)() const,\n"
54+
"$fqFieldClass$* ($fqMessageClass$::*)()\n"
55+
);
56+
}
57+
Header->Print(vars, "> GetFieldAccessorsByFieldTag($fqMessageClass$::T$field$FieldTag) {\n");
58+
WITH_INDENT(Header) {
59+
Header->Print(vars, "return std::tuple{\n");
60+
WITH_INDENT(Header) {
61+
Header->Print(vars,
62+
"&$fqMessageClass$::Has$field$,\n"
63+
"&$fqMessageClass$::Get$field$,\n"
64+
"&$fqMessageClass$::Mutable$field$\n"
65+
);
66+
}
67+
Header->Print(vars, "};\n");
68+
}
69+
Header->Print(vars, "}\n");
70+
71+
Header->Print(vars, "constexpr inline static ::NProtoBuf::uint32 GetFieldIdByFieldTag($fqMessageClass$::T$field$FieldTag) {\n");
72+
WITH_INDENT(Header) {
73+
Header->Print(vars, "return $fieldNumber$;\n");
74+
}
75+
Header->Print(vars, "}\n");
76+
}
77+
}
78+
}
79+
}
80+
81+
void GenerateCombinedType(TVars vars) {
82+
TMap<TString, TSet<const FieldDescriptor*>> outputs;
83+
for (auto i = 0; i < Message->field_count(); ++i) {
84+
const FieldDescriptor* field = Message->field(i);
85+
auto opts = field->options();
86+
for (int i = 0; i < opts.ExtensionSize(NKikimrConfig::NMarkers::CopyTo); ++i) {
87+
outputs[opts.GetExtension(NKikimrConfig::NMarkers::CopyTo, i)].insert(field);
88+
}
89+
}
90+
91+
for (const auto& [output, fields] : outputs) {
92+
vars["output"] = output;
93+
WITH_PLUGIN_MARKUP(Header, PLUGIN_NAME) {
94+
Header->Print(vars, "template <class TOut>\n");
95+
Header->Print(vars, "void CopyTo$output$(TOut& out) const {\n");
96+
Header->Indent();
97+
for (const auto* field : fields) {
98+
vars["field"] = field->name();
99+
100+
if (!field->is_repeated()) {
101+
if (field->message_type()) {
102+
Header->Print(vars, "if (Has$field$()) {\n");
103+
WITH_INDENT(Header) {
104+
Header->Print(vars, "out.Mutable$field$()->CopyFrom(Get$field$());\n");
105+
}
106+
Header->Print(vars, "}\n");
107+
} else if (field->is_optional()) {
108+
Header->Print(vars, "if (Has$field$()) {\n");
109+
WITH_INDENT(Header) {
110+
Header->Print(vars, "out.Set$field$(Get$field$());\n");
111+
}
112+
Header->Print(vars, "}\n");
113+
} else {
114+
Header->Print(vars, "out.Set$field$(Get$field$());\n");
115+
}
116+
} else {
117+
if (field->message_type()) {
118+
Header->Print(vars, "for (size_t i = 0; i < $field$Size(); ++i) {\n");
119+
WITH_INDENT(Header) {
120+
Header->Print(vars, "out.Add$field$()->CopyFrom(Get$field$(i));\n");
121+
}
122+
Header->Print(vars, "}\n");
123+
} else {
124+
Header->Print(vars, "for (const auto& field : Get$field$()) {\n");
125+
WITH_INDENT(Header) {
126+
Header->Print(vars, "out.Add$field$(field);\n");
127+
}
128+
Header->Print(vars, "}\n");
129+
}
130+
}
131+
}
132+
Header->Outdent();
133+
Header->Print(vars, "}\n");
134+
}
135+
}
136+
}
137+
138+
void GenerateSizeFields(TVars vars, const TSet<const FieldDescriptor*>& fields) {
139+
Header->Print(vars, "size_t $output$Size(const TProtoStringType& str) const {\n");
140+
WITH_INDENT(Header) {
141+
Header->Print(vars, "static std::map<TProtoStringType, size_t ($fqMessageClass$::*)() const> sizeHandlers{\n");
142+
WITH_INDENT(Header) {
143+
for (const auto* field : fields) {
144+
vars["field"] = field->name();
145+
Header->Print(vars, "{\"$field$\", &$fqMessageClass$::$field$Size},\n");
146+
}
147+
}
148+
Header->Print(vars, "};\n");
149+
150+
Header->Print(vars, "auto it = sizeHandlers.find(str);\n");
151+
Header->Print(vars, "return it == sizeHandlers.end() ? 0 : (this->*(it->second))();\n");
152+
}
153+
Header->Print(vars, "}\n");
154+
}
155+
156+
void GenerateAddFields(TVars vars, const TSet<const FieldDescriptor*>& fields) {
157+
Header->Print(vars, "$fqFieldClass$* Add$output$(const TProtoStringType& str) {\n");
158+
WITH_INDENT(Header) {
159+
Header->Print(vars, "static std::map<TProtoStringType, $fqFieldClass$* ($fqMessageClass$::*)()> addHandlers{\n");
160+
WITH_INDENT(Header) {
161+
for (const auto* field : fields) {
162+
vars["field"] = field->name();
163+
Header->Print(vars, "{\"$field$\", &$fqMessageClass$::Add$field$},\n");
164+
}
165+
}
166+
Header->Print(vars, "};\n");
167+
168+
Header->Print(vars, "return (this->*addHandlers.at(str))();\n");
169+
}
170+
Header->Print(vars, "}\n");
171+
}
172+
173+
void GenerateGetFields(TVars vars, const TSet<const FieldDescriptor*>& fields) {
174+
Header->Print(vars, "const ::google::protobuf::RepeatedPtrField<$fqFieldClass$>& Get$output$(const TProtoStringType& str) const {\n");
175+
WITH_INDENT(Header) {
176+
Header->Print(vars, "static std::map<TProtoStringType, const ::google::protobuf::RepeatedPtrField<$fqFieldClass$>& ($fqMessageClass$::*)() const> getHandlers{\n");
177+
WITH_INDENT(Header) {
178+
for (const auto* field : fields) {
179+
vars["field"] = field->name();
180+
Header->Print(vars, "{\"$field$\", &$fqMessageClass$::Get$field$},\n");
181+
}
182+
}
183+
Header->Print(vars, "};\n");
184+
185+
Header->Print(vars, "return (this->*getHandlers.at(str))();\n");
186+
}
187+
Header->Print(vars, "}\n");
188+
}
189+
190+
void GenerateMutableFields(TVars vars, const TSet<const FieldDescriptor*>& fields) {
191+
Header->Print(vars, "::google::protobuf::RepeatedPtrField<$fqFieldClass$>* Mutable$output$(const TProtoStringType& str) {\n");
192+
WITH_INDENT(Header) {
193+
Header->Print(vars, "static std::map<TProtoStringType, ::google::protobuf::RepeatedPtrField<$fqFieldClass$>* ($fqMessageClass$::*)()> mutableHandlers{\n");
194+
WITH_INDENT(Header) {
195+
for (const auto* field : fields) {
196+
vars["field"] = field->name();
197+
Header->Print(vars, "{\"$field$\", &$fqMessageClass$::Mutable$field$},\n");
198+
}
199+
}
200+
Header->Print(vars, "};\n");
201+
202+
Header->Print(vars, "return (this->*mutableHandlers.at(str))();\n");
203+
}
204+
Header->Print(vars, "}\n");
205+
}
206+
207+
void GenerateFieldsMap(TVars vars) {
208+
TMap<TString, TSet<const FieldDescriptor*>> outputs;
209+
for (auto i = 0; i < Message->field_count(); ++i) {
210+
const FieldDescriptor* field = Message->field(i);
211+
auto opts = field->options();
212+
for (int i = 0; i < opts.ExtensionSize(NKikimrConfig::NMarkers::AsMap); ++i) {
213+
outputs[opts.GetExtension(NKikimrConfig::NMarkers::AsMap, i)].insert(field);
214+
}
215+
}
216+
217+
if (!outputs) {
218+
return;
219+
}
220+
221+
WITH_PLUGIN_MARKUP(HeaderIncludes, PLUGIN_NAME) {
222+
HeaderIncludes->Print(Vars, "#include <map>\n");
223+
}
224+
225+
for (const auto& [output, fields] : outputs) {
226+
if (auto* fieldMessage = (*fields.begin())->message_type()) { // TODO implement for other classes
227+
for (const auto* field : fields) {
228+
if (fieldMessage->full_name() != field->message_type()->full_name()) {
229+
Cerr << "Messages in map MUST have the same type: "
230+
<< fieldMessage->full_name().c_str() << " != "
231+
<< field->message_type()->full_name().c_str() << Endl;
232+
Y_ABORT("Invariant failed");
233+
}
234+
235+
Y_ABORT_UNLESS(field->is_repeated(), "Only repeated fields are supported");
236+
}
237+
vars["fqFieldClass"] = FullyQualifiedClassName(fieldMessage);
238+
} else {
239+
Y_ABORT("Types other than Message is not supported yet.");
240+
}
241+
242+
vars["output"] = output;
243+
244+
WITH_PLUGIN_MARKUP(Header, PLUGIN_NAME) {
245+
GenerateSizeFields(vars, fields);
246+
GenerateAddFields(vars, fields);
247+
GenerateGetFields(vars, fields);
248+
GenerateMutableFields(vars, fields);
249+
}
250+
}
251+
}
252+
253+
public:
254+
255+
using TBaseMessageGenerator::TBaseMessageGenerator;
256+
257+
void Generate() {
258+
if (Message->options().GetExtension(NKikimrConfig::NMarkers::Root)) {
259+
GenerateConfigRoot(Vars);
260+
}
261+
262+
if (Message->options().GetExtension(NKikimrConfig::NMarkers::CombinedType)) {
263+
GenerateCombinedType(Vars);
264+
}
265+
266+
if (Message->options().GetExtension(NKikimrConfig::NMarkers::WithMapType)) {
267+
GenerateFieldsMap(Vars);
268+
}
269+
}
270+
};
271+
272+
class TCodeGenerator: public CodeGenerator {
273+
bool Generate(
274+
const FileDescriptor* file,
275+
const TProtoStringType&,
276+
OutputDirectory* output,
277+
TProtoStringType*) const override final
278+
{
279+
280+
for (auto i = 0; i < file->message_type_count(); ++i) {
281+
TMessageGenerator mg(file->message_type(i), output);
282+
mg.Generate();
283+
}
284+
285+
return true;
286+
}
287+
288+
uint64_t GetSupportedFeatures() const override
289+
{
290+
return FEATURE_PROTO3_OPTIONAL;
291+
}
292+
293+
}; // TCodeGenerator
294+
295+
int main(int argc, char* argv[]) {
296+
TCodeGenerator generator;
297+
return google::protobuf::compiler::PluginMain(argc, argv, &generator);
298+
}

0 commit comments

Comments
 (0)