Skip to content

Commit c0fe140

Browse files
EnjectionEnjection
authored
Merge fd5d302 into 0aeddde 
2 parents 0aeddde + fd5d302 commit c0fe140

File tree

18 files changed

+448
-16
lines changed

18 files changed

+448
-16
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

+279
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,279 @@
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/subst.h>
12+
13+
#include <google/protobuf/compiler/code_generator.h>
14+
#include <google/protobuf/compiler/plugin.h>
15+
#include <google/protobuf/io/printer.h>
16+
#include <google/protobuf/io/zero_copy_stream.h>
17+
#include <ydb/core/config/protos/marker.pb.h>
18+
19+
#include <ydb/public/lib/protobuf/helpers.h>
20+
21+
using TVariables = std::map<TString, TString>;
22+
23+
using namespace google::protobuf::compiler;
24+
using namespace google::protobuf;
25+
using namespace NKikimr::NProtobuf;
26+
27+
class TPrinter {
28+
public:
29+
explicit TPrinter(OutputDirectory* output, const TString& fileName, const TString& scope)
30+
: Output(output)
31+
, FileName(fileName)
32+
, Scope(scope)
33+
{
34+
}
35+
36+
io::Printer* operator->() {
37+
if (!Printer) {
38+
Stream.Reset(Output->OpenForInsert(FileName, Scope));
39+
Printer.ConstructInPlace(Stream.Get(), '$');
40+
}
41+
42+
return Printer.Get();
43+
}
44+
45+
private:
46+
OutputDirectory* Output;
47+
const TString FileName;
48+
const TString Scope;
49+
50+
THolder<io::ZeroCopyOutputStream> Stream;
51+
TMaybe<io::Printer> Printer;
52+
53+
}; // TPrinter
54+
55+
56+
class TMessageGenerator {
57+
public:
58+
explicit TMessageGenerator(const Descriptor* message, OutputDirectory* output)
59+
: Message(message)
60+
, Header(output, HeaderFileName(message), ClassScope(message))
61+
, HeaderIncludes(output, HeaderFileName(message), IncludesScope())
62+
, Vars({
63+
{"class", ClassName(message)},
64+
{"fqMessageClass", FullyQualifiedClassName(Message)},
65+
})
66+
{
67+
}
68+
69+
void Generate() {
70+
if (Message->options().GetExtension(NKikimrConfig::NMarkers::Root)) {
71+
for (auto i = 0; i < Message->field_count(); ++i) {
72+
const FieldDescriptor* field = Message->field(i);
73+
if (field->is_repeated()) {
74+
continue;
75+
}
76+
if (auto* fieldMessage = field->message_type()) {
77+
Vars["field"] = field->name();
78+
Vars["fqFieldClass"] = FullyQualifiedClassName(fieldMessage);
79+
Vars["fieldNumber"] = std::to_string(field->number());
80+
Header->Print(Vars, "/* BEGIN generated by config extension */\n");
81+
Header->Print(Vars, "struct T$field$FieldTag {};\n");
82+
Header->Print(Vars, "constexpr inline static std::tuple<\n");
83+
Header->Indent();
84+
Header->Print(Vars, "bool ($fqMessageClass$::*)() const,\n");
85+
Header->Print(Vars, "const $fqFieldClass$& ($fqMessageClass$::*)() const,\n");
86+
Header->Print(Vars, "$fqFieldClass$* ($fqMessageClass$::*)()\n");
87+
Header->Outdent();
88+
Header->Print(Vars, "> GetFieldAccessorsByFieldTag($fqMessageClass$::T$field$FieldTag) {\n");
89+
Header->Indent();
90+
Header->Print(Vars, "return std::tuple{\n");
91+
Header->Indent();
92+
Header->Print(Vars, "&$fqMessageClass$::Has$field$,\n");
93+
Header->Print(Vars, "&$fqMessageClass$::Get$field$,\n");
94+
Header->Print(Vars, "&$fqMessageClass$::Mutable$field$\n");
95+
Header->Outdent();
96+
Header->Print(Vars, "};\n");
97+
Header->Outdent();
98+
Header->Print(Vars, "}\n");
99+
Header->Print(Vars, "constexpr inline static ::NProtoBuf::uint32 GetFieldIdByFieldTag($fqMessageClass$::T$field$FieldTag) {\n");
100+
Header->Indent();
101+
Header->Print(Vars, "return $fieldNumber$;\n");
102+
Header->Outdent();
103+
Header->Print(Vars, "}\n");
104+
Header->Print(Vars, "/* END generated by config extension */\n");
105+
}
106+
}
107+
}
108+
109+
110+
if (Message->options().GetExtension(NKikimrConfig::NMarkers::CombinedType)) {
111+
TMap<TString, TSet<const FieldDescriptor*>> outputs;
112+
for (auto i = 0; i < Message->field_count(); ++i) {
113+
const FieldDescriptor* field = Message->field(i);
114+
auto opts = field->options();
115+
for (int i = 0; i < opts.ExtensionSize(NKikimrConfig::NMarkers::CopyTo); ++i) {
116+
outputs[opts.GetExtension(NKikimrConfig::NMarkers::CopyTo, i)].insert(field);
117+
}
118+
}
119+
120+
for (const auto& [output, fields] : outputs) {
121+
Vars["output"] = output;
122+
Header->Print(Vars, "/* BEGIN generated by config extension */\n");
123+
Header->Print(Vars, "template <class TOut>\n");
124+
Header->Print(Vars, "void CopyTo$output$(TOut& out) const {\n");
125+
Header->Indent();
126+
for (const auto* field : fields) {
127+
Vars["field"] = field->name();
128+
129+
if (!field->is_repeated()) {
130+
if (field->message_type()) {
131+
Header->Print(Vars, "if (Has$field$()) {\n");
132+
Header->Indent();
133+
Header->Print(Vars, "out.Mutable$field$()->CopyFrom(Get$field$());\n");
134+
Header->Outdent();
135+
Header->Print(Vars, "}\n");
136+
} else if (field->is_optional()) {
137+
Header->Print(Vars, "if (Has$field$()) {\n");
138+
Header->Indent();
139+
Header->Print(Vars, "out.Set$field$(Get$field$());\n");
140+
Header->Outdent();
141+
Header->Print(Vars, "}\n");
142+
} else {
143+
Header->Print(Vars, "out.Set$field$(Get$field$());\n");
144+
}
145+
} else {
146+
if (field->message_type()) {
147+
Header->Print(Vars, "for (size_t i = 0; i < $field$Size(); ++i) {\n");
148+
Header->Indent();
149+
Header->Print(Vars, "out.Add$field$()->CopyFrom(Get$field$(i));\n");
150+
Header->Outdent();
151+
Header->Print(Vars, "}\n");
152+
} else {
153+
Header->Print(Vars, "for (const auto& field : Get$field$()) {\n");
154+
Header->Indent();
155+
Header->Print(Vars, "out.Add$field$(field);\n");
156+
Header->Outdent();
157+
Header->Print(Vars, "}\n");
158+
}
159+
}
160+
}
161+
Header->Outdent();
162+
Header->Print(Vars, "}\n");
163+
Header->Print(Vars, "/* END generated by config extension */\n");
164+
}
165+
}
166+
167+
168+
if (Message->options().GetExtension(NKikimrConfig::NMarkers::WithMapType)) {
169+
HeaderIncludes->Print(Vars, "#include <map>\n");
170+
TMap<TString, TSet<const FieldDescriptor*>> outputs;
171+
for (auto i = 0; i < Message->field_count(); ++i) {
172+
const FieldDescriptor* field = Message->field(i);
173+
auto opts = field->options();
174+
for (int i = 0; i < opts.ExtensionSize(NKikimrConfig::NMarkers::AsMap); ++i) {
175+
outputs[opts.GetExtension(NKikimrConfig::NMarkers::AsMap, i)].insert(field);
176+
}
177+
}
178+
179+
// FIXME validate they have same type
180+
181+
for (const auto& [output, fields] : outputs) {
182+
if (auto* fieldMessage = (*fields.begin())->message_type()) { // FIXME for other classes
183+
Vars["fqFieldClass"] = FullyQualifiedClassName(fieldMessage);
184+
}
185+
Vars["output"] = output;
186+
Header->Print(Vars, "/* BEGIN generated by config extension */\n");
187+
Header->Print(Vars, "size_t $output$Size(const TProtoStringType& str) const {\n");
188+
Header->Indent();
189+
Header->Print(Vars, "static std::map<TProtoStringType, size_t ($fqMessageClass$::*)() const> sizeHandlers{\n");
190+
Header->Indent();
191+
for (const auto* field : fields) {
192+
Vars["field"] = field->name();
193+
Header->Print(Vars, "{\"$field$\", &$fqMessageClass$::$field$Size},\n");
194+
}
195+
Header->Outdent();
196+
Header->Print(Vars, "};\n");
197+
Header->Print(Vars, "auto it = sizeHandlers.find(str);\n");
198+
Header->Print(Vars, "return it == sizeHandlers.end() ? 0 : (this->*(it->second))();\n");
199+
Header->Outdent();
200+
Header->Print(Vars, "}\n");
201+
Header->Print(Vars, "$fqFieldClass$* Add$output$(const TProtoStringType& str) {\n");
202+
Header->Indent();
203+
Header->Print(Vars, "static std::map<TProtoStringType, $fqFieldClass$* ($fqMessageClass$::*)()> addHandlers{\n");
204+
Header->Indent();
205+
for (const auto* field : fields) {
206+
Vars["field"] = field->name();
207+
Header->Print(Vars, "{\"$field$\", &$fqMessageClass$::Add$field$},\n");
208+
}
209+
Header->Outdent();
210+
Header->Print(Vars, "};\n");
211+
Header->Print(Vars, "auto it = addHandlers.find(str);\n");
212+
Header->Print(Vars, "return it == addHandlers.end() ? nullptr : (this->*(it->second))();\n");
213+
Header->Outdent();
214+
Header->Print(Vars, "}\n");
215+
Header->Print(Vars, "const ::google::protobuf::RepeatedPtrField<$fqFieldClass$>& Get$output$(const TProtoStringType& str) const {\n");
216+
Header->Indent();
217+
Header->Print(Vars, "static std::map<TProtoStringType, const ::google::protobuf::RepeatedPtrField<$fqFieldClass$>& ($fqMessageClass$::*)() const> getHandlers{\n");
218+
Header->Indent();
219+
for (const auto* field : fields) {
220+
Vars["field"] = field->name();
221+
Header->Print(Vars, "{\"$field$\", &$fqMessageClass$::Get$field$},\n");
222+
}
223+
Header->Outdent();
224+
Header->Print(Vars, "};\n");
225+
Header->Print(Vars, "return (this->*getHandlers.at(str))();\n");
226+
Header->Outdent();
227+
Header->Print(Vars, "}\n");
228+
Header->Print(Vars, "::google::protobuf::RepeatedPtrField<$fqFieldClass$>* Mutable$output$(const TProtoStringType& str) {\n");
229+
Header->Indent();
230+
Header->Print(Vars, "static std::map<TProtoStringType, ::google::protobuf::RepeatedPtrField<$fqFieldClass$>* ($fqMessageClass$::*)()> mutableHandlers{\n");
231+
Header->Indent();
232+
for (const auto* field : fields) {
233+
Vars["field"] = field->name();
234+
Header->Print(Vars, "{\"$field$\", &$fqMessageClass$::Mutable$field$},\n");
235+
}
236+
Header->Outdent();
237+
Header->Print(Vars, "};\n");
238+
Header->Print(Vars, "return (this->*mutableHandlers.at(str))();\n");
239+
Header->Outdent();
240+
Header->Print(Vars, "}\n");
241+
Header->Print(Vars, "/* END generated by config extension */\n");
242+
}
243+
244+
}
245+
}
246+
private:
247+
const Descriptor* Message;
248+
TPrinter Header;
249+
TPrinter HeaderIncludes;
250+
TVariables Vars;
251+
};
252+
253+
class TCodeGenerator: public CodeGenerator {
254+
bool Generate(
255+
const FileDescriptor* file,
256+
const TProtoStringType&,
257+
OutputDirectory* output,
258+
TProtoStringType*) const override final
259+
{
260+
261+
for (auto i = 0; i < file->message_type_count(); ++i) {
262+
TMessageGenerator mg(file->message_type(i), output);
263+
mg.Generate();
264+
}
265+
266+
return true;
267+
}
268+
269+
uint64_t GetSupportedFeatures() const override
270+
{
271+
return FEATURE_PROTO3_OPTIONAL;
272+
}
273+
274+
}; // TCodeGenerator
275+
276+
int main(int argc, char* argv[]) {
277+
TCodeGenerator generator;
278+
return google::protobuf::compiler::PluginMain(argc, argv, &generator);
279+
}

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

+24
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
#include <ydb/core/config/tools/protobuf_plugin/ut/protos/config_root_test.pb.h>
2+
3+
#include <library/cpp/testing/unittest/registar.h>
4+
5+
Y_UNIT_TEST_SUITE(ValidationTests) {
6+
Y_UNIT_TEST(CanDispatchByTag) {
7+
NKikimrConfig::ActualConfigMessage msg;
8+
9+
msg.MutableField21();
10+
11+
auto [has1, get1, mut1] = NKikimrConfig::ActualConfigMessage::GetFieldAccessorsByFieldTag(NKikimrConfig::ActualConfigMessage::TField1FieldTag{});
12+
auto [has21, get21, mut21] = NKikimrConfig::ActualConfigMessage::GetFieldAccessorsByFieldTag(NKikimrConfig::ActualConfigMessage::TField21FieldTag{});
13+
14+
Y_UNUSED(get21, get1);
15+
16+
UNIT_ASSERT(!(msg.*has1)());
17+
UNIT_ASSERT((msg.*has21)());
18+
19+
(msg.*mut1)();
20+
21+
UNIT_ASSERT(msg.HasField1());
22+
UNIT_ASSERT((msg.*has1)());
23+
}
24+
}

‎ydb/core/config/tools/protobuf_plugin/ut/protos/config_root_test.proto

+19
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
import "ydb/core/config/protos/marker.proto";
2+
3+
package NKikimrConfig;
4+
option java_package = "ru.yandex.kikimr.proto";
5+
6+
message AdditionalMessage {
7+
}
8+
9+
message AnotherMessage {
10+
optional string Content = 1;
11+
optional AdditionalMessage Field1 = 2;
12+
}
13+
14+
message ActualConfigMessage {
15+
option (NMarkers.Root) = true;
16+
optional AdditionalMessage Field1 = 1;
17+
optional AdditionalMessage Field21 = 21;
18+
optional string Field201 = 201;
19+
}

‎ydb/core/config/tools/protobuf_plugin/ut/protos/ya.make

+15
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
PROTO_LIBRARY()
2+
3+
PEERDIR(
4+
ydb/core/config/protos
5+
)
6+
7+
SRCS(
8+
config_root_test.proto
9+
)
10+
11+
CPP_PROTO_PLUGIN0(config_proto_plugin ydb/core/config/tools/protobuf_plugin)
12+
13+
EXCLUDE_TAGS(GO_PROTO)
14+
15+
END()

‎ydb/core/config/tools/protobuf_plugin/ut/ya.make

+18
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
UNITTEST_FOR(ydb/core/config/tools/protobuf_plugin)
2+
3+
FORK_SUBTESTS()
4+
5+
IF (SANITIZER_TYPE OR WITH_VALGRIND)
6+
SIZE(MEDIUM)
7+
ENDIF()
8+
9+
PEERDIR(
10+
library/cpp/testing/unittest
11+
ydb/core/config/tools/protobuf_plugin/ut/protos
12+
)
13+
14+
SRCS(
15+
ut.cpp
16+
)
17+
18+
END()

0 commit comments

Comments
 (0)