1#ifndef PROTOBUF_BENCHMARKS_UTIL_SCHEMA_PROTO2_TO_PROTO3_UTIL_H_
2#define PROTOBUF_BENCHMARKS_UTIL_SCHEMA_PROTO2_TO_PROTO3_UTIL_H_
3
4#include "google/protobuf/message.h"
5#include "google/protobuf/descriptor.h"
6#include "google/protobuf/descriptor.pb.h"
7
8#include <sstream>
9#include <algorithm>
10
11using google::protobuf::Descriptor;
12using google::protobuf::DescriptorProto;
13using google::protobuf::FileDescriptorProto;
14using google::protobuf::FieldDescriptorProto;
15using google::protobuf::Message;
16using google::protobuf::EnumValueDescriptorProto;
17
18namespace google {
19namespace protobuf {
20namespace util {
21
22class SchemaGroupStripper {
23
24 public:
25  static void StripFile(const FileDescriptor* old_file,
26                        FileDescriptorProto *file) {
27    for (int i = file->mutable_message_type()->size() - 1; i >= 0; i--) {
28      if (IsMessageSet(old_file->message_type(i))) {
29        file->mutable_message_type()->DeleteSubrange(i, 1);
30        continue;
31      }
32      StripMessage(old_file->message_type(i), file->mutable_message_type(i));
33    }
34    for (int i = file->mutable_extension()->size() - 1; i >= 0; i--) {
35      auto field = old_file->extension(i);
36      if (field->type() == FieldDescriptor::TYPE_GROUP ||
37          IsMessageSet(field->message_type()) ||
38          IsMessageSet(field->containing_type())) {
39        file->mutable_extension()->DeleteSubrange(i, 1);
40      }
41    }
42  }
43
44 private:
45  static bool IsMessageSet(const Descriptor *descriptor) {
46    if (descriptor != nullptr
47        && descriptor->options().message_set_wire_format()) {
48      return true;
49    }
50    return false;
51  }
52
53  static void StripMessage(const Descriptor *old_message,
54                           DescriptorProto *new_message) {
55    for (int i = new_message->mutable_field()->size() - 1; i >= 0; i--) {
56      if (old_message->field(i)->type() == FieldDescriptor::TYPE_GROUP ||
57          IsMessageSet(old_message->field(i)->message_type())) {
58        new_message->mutable_field()->DeleteSubrange(i, 1);
59      }
60    }
61    for (int i = new_message->mutable_extension()->size() - 1; i >= 0; i--) {
62      auto field_type_name = new_message->mutable_extension(i)->type_name();
63      if (old_message->extension(i)->type() == FieldDescriptor::TYPE_GROUP ||
64          IsMessageSet(old_message->extension(i)->containing_type()) ||
65          IsMessageSet(old_message->extension(i)->message_type())) {
66        new_message->mutable_extension()->DeleteSubrange(i, 1);
67      }
68    }
69    for (int i = 0; i < new_message->mutable_nested_type()->size(); i++) {
70      StripMessage(old_message->nested_type(i),
71                   new_message->mutable_nested_type(i));
72    }
73  }
74
75};
76
77class EnumScrubber {
78
79 public:
80  EnumScrubber()
81      : total_added_(0) {
82  }
83
84  void ScrubFile(FileDescriptorProto *file) {
85    for (int i = 0; i < file->enum_type_size(); i++) {
86      ScrubEnum(file->mutable_enum_type(i));
87    }
88    for (int i = 0; i < file->mutable_message_type()->size(); i++) {
89      ScrubMessage(file->mutable_message_type(i));
90    }
91  }
92
93 private:
94  void ScrubEnum(EnumDescriptorProto *enum_type) {
95    if (enum_type->value(0).number() != 0) {
96      bool has_zero = false;
97      for (int j = 0; j < enum_type->value().size(); j++) {
98        if (enum_type->value(j).number() == 0) {
99          EnumValueDescriptorProto temp_enum_value;
100          temp_enum_value.CopyFrom(enum_type->value(j));
101          enum_type->mutable_value(j)->CopyFrom(enum_type->value(0));
102          enum_type->mutable_value(0)->CopyFrom(temp_enum_value);
103          has_zero = true;
104          break;
105        }
106      }
107      if (!has_zero) {
108        enum_type->mutable_value()->Add();
109        for (int i = enum_type->mutable_value()->size() - 1; i > 0; i--) {
110          enum_type->mutable_value(i)->CopyFrom(
111              *enum_type->mutable_value(i - 1));
112        }
113        enum_type->mutable_value(0)->set_number(0);
114        enum_type->mutable_value(0)->set_name("ADDED_ZERO_VALUE_" +
115                                              std::to_string(total_added_++));
116      }
117    }
118
119  }
120
121  void ScrubMessage(DescriptorProto *message_type) {
122    for (int i = 0; i < message_type->mutable_enum_type()->size(); i++) {
123      ScrubEnum(message_type->mutable_enum_type(i));
124    }
125    for (int i = 0; i < message_type->mutable_nested_type()->size(); i++) {
126      ScrubMessage(message_type->mutable_nested_type(i));
127    }
128  }
129
130  int total_added_;
131};
132
133class ExtensionStripper {
134 public:
135  static void StripFile(FileDescriptorProto *file) {
136    for (int i = 0; i < file->mutable_message_type()->size(); i++) {
137      StripMessage(file->mutable_message_type(i));
138    }
139    file->mutable_extension()->Clear();
140  }
141 private:
142  static void StripMessage(DescriptorProto *message_type) {
143    message_type->mutable_extension()->Clear();
144    message_type->clear_extension_range();
145    for (int i = 0; i < message_type->mutable_nested_type()->size(); i++) {
146      StripMessage(message_type->mutable_nested_type(i));
147    }
148  }
149};
150
151
152class FieldScrubber {
153 public:
154  static void ScrubFile(FileDescriptorProto *file) {
155    for (int i = 0; i < file->mutable_message_type()->size(); i++) {
156      ScrubMessage(file->mutable_message_type(i));
157    }
158    for (int i = 0; i < file->mutable_extension()->size(); i++) {
159      file->mutable_extension(i)->clear_default_value();
160      if (ShouldClearLabel(file->mutable_extension(i))) {
161        file->mutable_extension(i)->clear_label();
162      }
163    }
164  }
165 private:
166  static bool ShouldClearLabel(const FieldDescriptorProto *field) {
167    return field->label() == FieldDescriptorProto::LABEL_REQUIRED;
168  }
169
170  static void ScrubMessage(DescriptorProto *message_type) {
171    message_type->mutable_extension()->Clear();
172    for (int i = 0; i < message_type->mutable_extension()->size(); i++) {
173      message_type->mutable_extension(i)->clear_default_value();
174      if (ShouldClearLabel(message_type->mutable_extension(i))) {
175        message_type->mutable_extension(i)->clear_label();
176      }
177    }
178    for (int i = 0; i < message_type->mutable_field()->size(); i++) {
179      message_type->mutable_field(i)->clear_default_value();
180      if (ShouldClearLabel(message_type->mutable_field(i))) {
181        message_type->mutable_field(i)->clear_label();
182      }
183    }
184    for (int i = 0; i < message_type->mutable_nested_type()->size(); i++) {
185      ScrubMessage(message_type->mutable_nested_type(i));
186    }
187  }
188};
189
190}  // namespace util
191}  // namespace protobuf
192}  // namespace google
193
194#endif  // PROTOBUF_BENCHMARKS_UTIL_SCHEMA_PROTO2_TO_PROTO3_UTIL_H_
195