Yilun Chong
7 years ago
committed by
GitHub
commit
320d56c833
6 changed files with 584 additions and 53 deletions
@ -0,0 +1,105 @@ |
||||
#include "benchmarks.pb.h" |
||||
#include "datasets/google_message1/proto2/benchmark_message1_proto2.pb.h" |
||||
#include "datasets/google_message1/proto3/benchmark_message1_proto3.pb.h" |
||||
#include "datasets/google_message2/benchmark_message2.pb.h" |
||||
#include "datasets/google_message3/benchmark_message3.pb.h" |
||||
#include "datasets/google_message4/benchmark_message4.pb.h" |
||||
|
||||
#include "google/protobuf/message.h" |
||||
#include "google/protobuf/descriptor.h" |
||||
|
||||
#include <fstream> |
||||
|
||||
using google::protobuf::FieldDescriptor; |
||||
using google::protobuf::Message; |
||||
using google::protobuf::Reflection; |
||||
|
||||
|
||||
class DataGroupStripper { |
||||
public: |
||||
static void StripMessage(Message *message) { |
||||
std::vector<const FieldDescriptor*> set_fields; |
||||
const Reflection* reflection = message->GetReflection(); |
||||
reflection->ListFields(*message, &set_fields); |
||||
|
||||
for (size_t i = 0; i < set_fields.size(); i++) { |
||||
const FieldDescriptor* field = set_fields[i]; |
||||
if (field->type() == FieldDescriptor::TYPE_GROUP) { |
||||
reflection->ClearField(message, field); |
||||
} |
||||
if (field->type() == FieldDescriptor::TYPE_MESSAGE) { |
||||
if (field->is_repeated()) { |
||||
for (int j = 0; j < reflection->FieldSize(*message, field); j++) { |
||||
StripMessage(reflection->MutableRepeatedMessage(message, field, j)); |
||||
} |
||||
} else { |
||||
StripMessage(reflection->MutableMessage(message, field)); |
||||
} |
||||
} |
||||
} |
||||
|
||||
reflection->MutableUnknownFields(message)->Clear(); |
||||
} |
||||
}; |
||||
|
||||
std::string ReadFile(const std::string& name) { |
||||
std::ifstream file(name.c_str()); |
||||
GOOGLE_CHECK(file.is_open()) << "Couldn't find file '" |
||||
<< name |
||||
<< "', please make sure you are running this command from the benchmarks" |
||||
<< " directory.\n"; |
||||
return std::string((std::istreambuf_iterator<char>(file)), |
||||
std::istreambuf_iterator<char>()); |
||||
} |
||||
|
||||
int main(int argc, char *argv[]) { |
||||
if (argc % 2 == 0 || argc == 1) { |
||||
std::cerr << "Usage: [input_files] [output_file_names] where " << |
||||
"input_files are one to one mapping to output_file_names." << |
||||
std::endl; |
||||
return 1; |
||||
} |
||||
|
||||
for (int i = argc / 2; i > 0; i--) { |
||||
const std::string &input_file = argv[i]; |
||||
const std::string &output_file = argv[i + argc / 2]; |
||||
|
||||
std::cerr << "Generating " << input_file |
||||
<< " to " << output_file << std::endl; |
||||
benchmarks::BenchmarkDataset dataset; |
||||
Message* message; |
||||
std::string dataset_payload = ReadFile(input_file); |
||||
GOOGLE_CHECK(dataset.ParseFromString(dataset_payload)) |
||||
<< "Can' t parse data file " << input_file; |
||||
|
||||
if (dataset.message_name() == "benchmarks.proto3.GoogleMessage1") { |
||||
message = new benchmarks::proto3::GoogleMessage1; |
||||
} else if (dataset.message_name() == "benchmarks.proto2.GoogleMessage1") { |
||||
message = new benchmarks::proto2::GoogleMessage1; |
||||
} else if (dataset.message_name() == "benchmarks.proto2.GoogleMessage2") { |
||||
message = new benchmarks::proto2::GoogleMessage2; |
||||
} else if (dataset.message_name() == |
||||
"benchmarks.google_message3.GoogleMessage3") { |
||||
message = new benchmarks::google_message3::GoogleMessage3; |
||||
} else if (dataset.message_name() == |
||||
"benchmarks.google_message4.GoogleMessage4") { |
||||
message = new benchmarks::google_message4::GoogleMessage4; |
||||
} else { |
||||
std::cerr << "Unknown message type: " << dataset.message_name(); |
||||
exit(1); |
||||
} |
||||
|
||||
for (int i = 0; i < dataset.payload_size(); i++) { |
||||
message->ParseFromString(dataset.payload(i)); |
||||
DataGroupStripper::StripMessage(message); |
||||
dataset.set_payload(i, message->SerializeAsString()); |
||||
} |
||||
|
||||
std::ofstream ofs(output_file); |
||||
ofs << dataset.SerializeAsString(); |
||||
ofs.close(); |
||||
} |
||||
|
||||
|
||||
return 0; |
||||
} |
||||
@ -0,0 +1,103 @@ |
||||
#include "google/protobuf/compiler/code_generator.h" |
||||
#include "google/protobuf/io/zero_copy_stream.h" |
||||
#include "google/protobuf/io/printer.h" |
||||
#include "google/protobuf/descriptor.h" |
||||
#include "google/protobuf/descriptor.pb.h" |
||||
#include "schema_proto2_to_proto3_util.h" |
||||
|
||||
#include "google/protobuf/compiler/plugin.h" |
||||
|
||||
using google::protobuf::FileDescriptorProto; |
||||
using google::protobuf::FileDescriptor; |
||||
using google::protobuf::DescriptorPool; |
||||
using google::protobuf::io::Printer; |
||||
using google::protobuf::util::SchemaGroupStripper; |
||||
using google::protobuf::util::SchemaAddZeroEnumValue; |
||||
|
||||
namespace google { |
||||
namespace protobuf { |
||||
namespace compiler { |
||||
|
||||
namespace { |
||||
|
||||
string StripProto(string filename) { |
||||
if (filename.substr(filename.size() - 11) == ".protodevel") { |
||||
// .protodevel
|
||||
return filename.substr(0, filename.size() - 11); |
||||
} else { |
||||
// .proto
|
||||
return filename.substr(0, filename.size() - 6); |
||||
} |
||||
} |
||||
|
||||
DescriptorPool new_pool_; |
||||
|
||||
} // namespace
|
||||
|
||||
class GoGoProtoGenerator : public CodeGenerator { |
||||
public: |
||||
virtual bool GenerateAll(const std::vector<const FileDescriptor*>& files, |
||||
const string& parameter, |
||||
GeneratorContext* context, |
||||
string* error) const { |
||||
for (int i = 0; i < files.size(); i++) { |
||||
for (auto file : files) { |
||||
bool can_generate = |
||||
(new_pool_.FindFileByName(file->name()) == nullptr); |
||||
for (int j = 0; j < file->dependency_count(); j++) { |
||||
can_generate &= (new_pool_.FindFileByName( |
||||
file->dependency(j)->name()) != nullptr); |
||||
} |
||||
for (int j = 0; j < file->public_dependency_count(); j++) { |
||||
can_generate &= (new_pool_.FindFileByName( |
||||
file->public_dependency(j)->name()) != nullptr); |
||||
} |
||||
for (int j = 0; j < file->weak_dependency_count(); j++) { |
||||
can_generate &= (new_pool_.FindFileByName( |
||||
file->weak_dependency(j)->name()) != nullptr); |
||||
} |
||||
if (can_generate) { |
||||
Generate(file, parameter, context, error); |
||||
break; |
||||
} |
||||
} |
||||
} |
||||
|
||||
return true; |
||||
} |
||||
|
||||
virtual bool Generate(const FileDescriptor* file, |
||||
const string& parameter, |
||||
GeneratorContext* context, |
||||
string* error) const { |
||||
FileDescriptorProto new_file; |
||||
file->CopyTo(&new_file); |
||||
SchemaGroupStripper::StripFile(file, &new_file); |
||||
|
||||
SchemaAddZeroEnumValue enum_scrubber; |
||||
enum_scrubber.ScrubFile(&new_file); |
||||
|
||||
string filename = file->name(); |
||||
string basename = StripProto(filename); |
||||
|
||||
std::vector<std::pair<string,string>> option_pairs; |
||||
ParseGeneratorParameter(parameter, &option_pairs); |
||||
|
||||
std::unique_ptr<google::protobuf::io::ZeroCopyOutputStream> output( |
||||
context->Open(basename + ".proto")); |
||||
string content = new_pool_.BuildFile(new_file)->DebugString(); |
||||
Printer printer(output.get(), '$'); |
||||
printer.WriteRaw(content.c_str(), content.size()); |
||||
|
||||
return true; |
||||
} |
||||
}; |
||||
|
||||
} // namespace compiler
|
||||
} // namespace protobuf
|
||||
} // namespace google
|
||||
|
||||
int main(int argc, char* argv[]) { |
||||
google::protobuf::compiler::GoGoProtoGenerator generator; |
||||
return google::protobuf::compiler::PluginMain(argc, argv, &generator); |
||||
} |
||||
@ -0,0 +1,137 @@ |
||||
#ifndef PROTOBUF_BENCHMARKS_UTIL_SCHEMA_PROTO2_TO_PROTO3_UTIL_H_ |
||||
#define PROTOBUF_BENCHMARKS_UTIL_SCHEMA_PROTO2_TO_PROTO3_UTIL_H_ |
||||
|
||||
#include "google/protobuf/message.h" |
||||
#include "google/protobuf/descriptor.h" |
||||
#include "google/protobuf/descriptor.pb.h" |
||||
|
||||
#include <sstream> |
||||
#include <algorithm> |
||||
|
||||
using google::protobuf::Descriptor; |
||||
using google::protobuf::DescriptorProto; |
||||
using google::protobuf::FileDescriptorProto; |
||||
using google::protobuf::FieldDescriptorProto; |
||||
using google::protobuf::Message; |
||||
using google::protobuf::EnumValueDescriptorProto; |
||||
|
||||
namespace google { |
||||
namespace protobuf { |
||||
namespace util { |
||||
|
||||
class SchemaGroupStripper { |
||||
|
||||
public: |
||||
static void StripFile(const FileDescriptor* old_file, |
||||
FileDescriptorProto *file) { |
||||
for (int i = file->mutable_message_type()->size() - 1; i >= 0; i--) { |
||||
if (IsMessageSet(old_file->message_type(i))) { |
||||
file->mutable_message_type()->DeleteSubrange(i, 1); |
||||
continue; |
||||
} |
||||
StripMessage(old_file->message_type(i), file->mutable_message_type(i)); |
||||
} |
||||
for (int i = file->mutable_extension()->size() - 1; i >= 0; i--) { |
||||
auto field = old_file->extension(i); |
||||
if (field->type() == FieldDescriptor::TYPE_GROUP || |
||||
IsMessageSet(field->message_type()) || |
||||
IsMessageSet(field->containing_type())) { |
||||
file->mutable_extension()->DeleteSubrange(i, 1); |
||||
} |
||||
} |
||||
} |
||||
|
||||
private: |
||||
static bool IsMessageSet(const Descriptor *descriptor) { |
||||
if (descriptor != nullptr |
||||
&& descriptor->options().message_set_wire_format()) { |
||||
return true; |
||||
} |
||||
return false; |
||||
} |
||||
|
||||
static void StripMessage(const Descriptor *old_message, |
||||
DescriptorProto *new_message) { |
||||
for (int i = new_message->mutable_field()->size() - 1; i >= 0; i--) { |
||||
if (old_message->field(i)->type() == FieldDescriptor::TYPE_GROUP || |
||||
IsMessageSet(old_message->field(i)->message_type())) { |
||||
new_message->mutable_field()->DeleteSubrange(i, 1); |
||||
} |
||||
} |
||||
for (int i = new_message->mutable_extension()->size() - 1; i >= 0; i--) { |
||||
auto field_type_name = new_message->mutable_extension(i)->type_name(); |
||||
if (old_message->extension(i)->type() == FieldDescriptor::TYPE_GROUP || |
||||
IsMessageSet(old_message->extension(i)->containing_type()) || |
||||
IsMessageSet(old_message->extension(i)->message_type())) { |
||||
new_message->mutable_extension()->DeleteSubrange(i, 1); |
||||
} |
||||
} |
||||
for (int i = 0; i < new_message->mutable_nested_type()->size(); i++) { |
||||
StripMessage(old_message->nested_type(i), |
||||
new_message->mutable_nested_type(i)); |
||||
} |
||||
} |
||||
|
||||
}; |
||||
|
||||
class SchemaAddZeroEnumValue { |
||||
|
||||
public: |
||||
SchemaAddZeroEnumValue() |
||||
: total_added_(0) { |
||||
} |
||||
|
||||
void ScrubFile(FileDescriptorProto *file) { |
||||
for (int i = 0; i < file->enum_type_size(); i++) { |
||||
ScrubEnum(file->mutable_enum_type(i)); |
||||
} |
||||
for (int i = 0; i < file->mutable_message_type()->size(); i++) { |
||||
ScrubMessage(file->mutable_message_type(i)); |
||||
} |
||||
} |
||||
|
||||
private: |
||||
void ScrubEnum(EnumDescriptorProto *enum_type) { |
||||
if (enum_type->value(0).number() != 0) { |
||||
bool has_zero = false; |
||||
for (int j = 0; j < enum_type->value().size(); j++) { |
||||
if (enum_type->value(j).number() == 0) { |
||||
EnumValueDescriptorProto temp_enum_value; |
||||
temp_enum_value.CopyFrom(enum_type->value(j)); |
||||
enum_type->mutable_value(j)->CopyFrom(enum_type->value(0)); |
||||
enum_type->mutable_value(0)->CopyFrom(temp_enum_value); |
||||
has_zero = true; |
||||
break; |
||||
} |
||||
} |
||||
if (!has_zero) { |
||||
enum_type->mutable_value()->Add(); |
||||
for (int i = enum_type->mutable_value()->size() - 1; i > 0; i--) { |
||||
enum_type->mutable_value(i)->CopyFrom( |
||||
*enum_type->mutable_value(i - 1)); |
||||
} |
||||
enum_type->mutable_value(0)->set_number(0); |
||||
enum_type->mutable_value(0)->set_name("ADDED_ZERO_VALUE_" + |
||||
std::to_string(total_added_++)); |
||||
} |
||||
} |
||||
|
||||
} |
||||
|
||||
void ScrubMessage(DescriptorProto *message_type) { |
||||
for (int i = 0; i < message_type->mutable_enum_type()->size(); i++) { |
||||
ScrubEnum(message_type->mutable_enum_type(i)); |
||||
} |
||||
for (int i = 0; i < message_type->mutable_nested_type()->size(); i++) { |
||||
ScrubMessage(message_type->mutable_nested_type(i)); |
||||
} |
||||
} |
||||
|
||||
int total_added_; |
||||
}; |
||||
|
||||
} // namespace util
|
||||
} // namespace protobuf
|
||||
} // namespace google
|
||||
|
||||
#endif // PROTOBUF_BENCHMARKS_UTIL_SCHEMA_PROTO2_TO_PROTO3_UTIL_H_
|
||||
Loading…
Reference in new issue