1 #include "google/protobuf/compiler/code_generator.h"
2 #include "google/protobuf/io/zero_copy_stream.h"
3 #include "google/protobuf/io/printer.h"
4 #include "google/protobuf/descriptor.h"
5 #include "google/protobuf/descriptor.pb.h"
6 #include "schema_proto2_to_proto3_util.h"
7 
8 #include "google/protobuf/compiler/plugin.h"
9 
10 using google::protobuf::FileDescriptorProto;
11 using google::protobuf::FileDescriptor;
12 using google::protobuf::DescriptorPool;
13 using google::protobuf::io::Printer;
14 using google::protobuf::util::SchemaGroupStripper;
15 using google::protobuf::util::EnumScrubber;
16 using google::protobuf::util::ExtensionStripper;
17 using google::protobuf::util::FieldScrubber;
18 
19 namespace google {
20 namespace protobuf {
21 namespace compiler {
22 
23 namespace {
24 
StripProto(string filename)25 string StripProto(string filename) {
26   return filename.substr(0, filename.rfind(".proto"));
27 }
28 
GetPool()29 DescriptorPool* GetPool() {
30   static DescriptorPool *pool = new DescriptorPool();
31   return pool;
32 }
33 
34 }  // namespace
35 
36 class Proto2ToProto3Generator final : public CodeGenerator {
37  public:
GenerateAll(const std::vector<const FileDescriptor * > & files,const string & parameter,GeneratorContext * context,string * error) const38   bool GenerateAll(const std::vector<const FileDescriptor*>& files,
39                            const string& parameter,
40                            GeneratorContext* context,
41                            string* error) const {
42     for (int i = 0; i < files.size(); i++) {
43       for (auto file : files) {
44         if (CanGenerate(file)) {
45           Generate(file, parameter, context, error);
46           break;
47         }
48       }
49     }
50 
51     return true;
52   }
53 
Generate(const FileDescriptor * file,const string & parameter,GeneratorContext * context,string * error) const54   bool Generate(const FileDescriptor* file,
55                         const string& parameter,
56                         GeneratorContext* context,
57                         string* error) const {
58     FileDescriptorProto new_file;
59     file->CopyTo(&new_file);
60     SchemaGroupStripper::StripFile(file, &new_file);
61 
62     EnumScrubber enum_scrubber;
63     enum_scrubber.ScrubFile(&new_file);
64     ExtensionStripper::StripFile(&new_file);
65     FieldScrubber::ScrubFile(&new_file);
66     new_file.set_syntax("proto3");
67 
68     string filename = file->name();
69     string basename = StripProto(filename);
70 
71     std::vector<std::pair<string,string>> option_pairs;
72     ParseGeneratorParameter(parameter, &option_pairs);
73 
74     std::unique_ptr<google::protobuf::io::ZeroCopyOutputStream> output(
75         context->Open(basename + ".proto"));
76     string content = GetPool()->BuildFile(new_file)->DebugString();
77     Printer printer(output.get(), '$');
78     printer.WriteRaw(content.c_str(), content.size());
79 
80     return true;
81   }
82  private:
CanGenerate(const FileDescriptor * file) const83   bool CanGenerate(const FileDescriptor* file) const {
84     if (GetPool()->FindFileByName(file->name()) != nullptr) {
85       return false;
86     }
87     for (int j = 0; j < file->dependency_count(); j++) {
88       if (GetPool()->FindFileByName(file->dependency(j)->name()) == nullptr) {
89         return false;
90       }
91     }
92     for (int j = 0; j < file->public_dependency_count(); j++) {
93       if (GetPool()->FindFileByName(
94           file->public_dependency(j)->name()) == nullptr) {
95         return false;
96       }
97     }
98     for (int j = 0; j < file->weak_dependency_count(); j++) {
99       if (GetPool()->FindFileByName(
100           file->weak_dependency(j)->name()) == nullptr) {
101         return false;
102       }
103     }
104     return true;
105   }
106 };
107 
108 }  // namespace compiler
109 }  // namespace protobuf
110 }  // namespace google
111 
main(int argc,char * argv[])112 int main(int argc, char* argv[]) {
113   google::protobuf::compiler::Proto2ToProto3Generator generator;
114   return google::protobuf::compiler::PluginMain(argc, argv, &generator);
115 }
116