schema_proto2_to_proto3_util.h 6.4 KB

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