gogo_data_scrubber.cc 2.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. #include "benchmarks.pb.h"
  2. #include "datasets/google_message1/proto2/benchmark_message1_proto2.pb.h"
  3. #include "datasets/google_message1/proto3/benchmark_message1_proto3.pb.h"
  4. #include "datasets/google_message2/benchmark_message2.pb.h"
  5. #include "datasets/google_message3/benchmark_message3.pb.h"
  6. #include "datasets/google_message4/benchmark_message4.pb.h"
  7. #include "data_proto2_to_proto3_util.h"
  8. #include <fstream>
  9. using google::protobuf::util::GogoDataStripper;
  10. std::string ReadFile(const std::string& name) {
  11. std::ifstream file(name.c_str());
  12. GOOGLE_CHECK(file.is_open()) << "Couldn't find file '"
  13. << name
  14. << "', please make sure you are running this command from the benchmarks"
  15. << " directory.\n";
  16. return std::string((std::istreambuf_iterator<char>(file)),
  17. std::istreambuf_iterator<char>());
  18. }
  19. int main(int argc, char *argv[]) {
  20. if (argc % 2 == 0 || argc == 1) {
  21. std::cerr << "Usage: [input_files] [output_file_names] where " <<
  22. "input_files are one to one mapping to output_file_names." <<
  23. std::endl;
  24. return 1;
  25. }
  26. for (int i = argc / 2; i > 0; i--) {
  27. const std::string &input_file = argv[i];
  28. const std::string &output_file = argv[i + argc / 2];
  29. std::cerr << "Generating " << input_file
  30. << " to " << output_file << std::endl;
  31. benchmarks::BenchmarkDataset dataset;
  32. Message* message;
  33. std::string dataset_payload = ReadFile(input_file);
  34. GOOGLE_CHECK(dataset.ParseFromString(dataset_payload))
  35. << "Can' t parse data file " << input_file;
  36. if (dataset.message_name() == "benchmarks.proto3.GoogleMessage1") {
  37. message = new benchmarks::proto3::GoogleMessage1;
  38. } else if (dataset.message_name() == "benchmarks.proto2.GoogleMessage1") {
  39. message = new benchmarks::proto2::GoogleMessage1;
  40. } else if (dataset.message_name() == "benchmarks.proto2.GoogleMessage2") {
  41. message = new benchmarks::proto2::GoogleMessage2;
  42. } else if (dataset.message_name() ==
  43. "benchmarks.google_message3.GoogleMessage3") {
  44. message = new benchmarks::google_message3::GoogleMessage3;
  45. } else if (dataset.message_name() ==
  46. "benchmarks.google_message4.GoogleMessage4") {
  47. message = new benchmarks::google_message4::GoogleMessage4;
  48. } else {
  49. std::cerr << "Unknown message type: " << dataset.message_name();
  50. exit(1);
  51. }
  52. for (int i = 0; i < dataset.payload_size(); i++) {
  53. message->ParseFromString(dataset.payload(i));
  54. GogoDataStripper stripper;
  55. stripper.StripMessage(message);
  56. dataset.set_payload(i, message->SerializeAsString());
  57. }
  58. std::ofstream ofs(output_file);
  59. ofs << dataset.SerializeAsString();
  60. ofs.close();
  61. }
  62. return 0;
  63. }