ProtoZero plugin implementation.

components/tracing/tools/proto_zero_plugin/proto_zero_generator.cc

 // Copyright 2016 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "proto_zero_generator.h"
 
+#include <map>
 #include <memory>
 #include <string>
 
 #include "third_party/protobuf/src/google/protobuf/descriptor.h"
 #include "third_party/protobuf/src/google/protobuf/io/printer.h"
 #include "third_party/protobuf/src/google/protobuf/io/zero_copy_stream.h"
 #include "third_party/protobuf/src/google/protobuf/stubs/strutil.h"
 
 namespace tracing {
 namespace proto {
 
+using google::protobuf::Descriptor;  // Message descriptor.
+using google::protobuf::EnumDescriptor;
+using google::protobuf::EnumValueDescriptor;
+using google::protobuf::FieldDescriptor;
 using google::protobuf::FileDescriptor;
-using google::protobuf::StripSuffixString;
 using google::protobuf::compiler::GeneratorContext;
 using google::protobuf::io::Printer;
 using google::protobuf::io::ZeroCopyOutputStream;
 
+using google::protobuf::Split;
+using google::protobuf::StripPrefixString;
+using google::protobuf::StripString;
+using google::protobuf::StripSuffixString;
+using google::protobuf::UpperString;
+
 namespace {
 
 class GeneratorJob {
  public:
   GeneratorJob(const FileDescriptor *file,
                Printer* stub_h_printer,
                Printer* stub_cc_printer)
-    : file_(file),
+    : source_(file),
       stub_h_(stub_h_printer),
       stub_cc_(stub_cc_printer) {}
 
   bool GenerateStubs() {
-    stub_h_->Print(
-        "// Autogenerated. DO NOT EDIT.\n"
-        "// Generated by: //components/tracing/proto_zero_plugin.\n\n"
-        "// Package: $package$\n",
-        "package", file_->package());
-    stub_cc_->Print(
-        "// Autogenerated. DO NOT EDIT.\n"
-        "// Generated by: //components/tracing/proto_zero_plugin.\n\n"
-        "// This file intentionally left blank.\n");
-    // TODO(kraynov) Implement in the next CL (crbug.com/608721).
-    return true;
+    Preprocess();
+    GeneratePrologue();
+    for (const EnumDescriptor* enumeration: enums_)
+      GenerateEnumDescriptor(enumeration);
+    for (const Descriptor* message : messages_)
+      GenerateMessageDescriptor(message);
+    GenerateEpilogue();
+    return error_.empty();
   }
 
   // If generator fails to produce stubs for a particular proto definitions
   // it finishes with undefined output and writes the first error occured.
   const std::string& GetFirstError() const {
     return error_;
   }
 
  private:
   // Only the first error will be recorded.
   void Abort(const std::string& reason) {
-    if (error_.empty()) {
+    if (error_.empty())
       error_ = reason;
-    }
-  }
-
-  const FileDescriptor* const file_;
+  }
+
+  // Get full name (including outer descriptors) of proto descriptor.
+  template <class T>
+  inline std::string GetDescriptorName(const T* descriptor) {
+    if (!package_.empty()) {
+      return StripPrefixString(descriptor->full_name(), package_ + ".");
+    } else {
+      return descriptor->full_name();
+    }
+  }
+
+  // Get C++ class name corresponding to proto descriptor.
+  // Nested names are splitted by underscores. Underscores in type names aren't
+  // prohibited but not recommended in order to avoid name collisions.
+  template <class T>
+  inline std::string GetCppClassName(const T* descriptor, bool full = false) {
+    std::string name = GetDescriptorName(descriptor);
+    StripString(&name, ".", '_');
+    if (full)
+      name = full_namespace_prefix_ + name;
+    return name;
+  }
+
+  // Small enums can be written faster without involving VarInt encoder.
+  inline bool IsTinyEnumField(const FieldDescriptor* field) {
+    if (field->type() != FieldDescriptor::TYPE_ENUM)
+      return false;
+    const EnumDescriptor* enumeration = field->enum_type();
+
+    for (int i = 0; i < enumeration->value_count(); ++i) {
+      int32_t value = enumeration->value(i)->number();
+      if (value < 0 || value > 0x7F)
+        return false;
+    }
+    return true;
+  }
+
+  void Preprocess() {
+    // Package name maps to a series of namespaces.
+    package_ = source_->package();
+    namespaces_ = Split(package_, ".");
+    full_namespace_prefix_ = "::";
+    for (const std::string& ns : namespaces_)
+      full_namespace_prefix_ += ns + "::";
+
+    // Collect message descriptors in DFS order.
+    std::vector<const Descriptor*> stack;
+    for (int i = 0; i < source_->message_type_count(); ++i)
+      stack.push_back(source_->message_type(i));
+
+    while (!stack.empty()) {
+      const Descriptor* message = stack.back();
+      stack.pop_back();
+      messages_.push_back(message);
+      for (int i = 0; i < message->nested_type_count(); ++i) {
+        stack.push_back(message->nested_type(i));
+      }
+    }
+
+    // Collect enums.
+    for (int i = 0; i < source_->enum_type_count(); ++i)
+      enums_.push_back(source_->enum_type(i));
+
+    for (const Descriptor* message : messages_) {
+      for (int i = 0; i < message->enum_type_count(); ++i) {
+        enums_.push_back(message->enum_type(i));
+      }
+    }
+  }
+
+  // Print top header, namespaces and forward declarations.
+  void GeneratePrologue() {
+    std::string greeting =
+        "// Autogenerated. DO NOT EDIT.\n"
+        "// Protobuf compiler (protoc) has generated these stubs with\n"
+        "// //components/tracing/tools/proto_zero_plugin.\n";
+    std::string guard = package_ + "_" + source_->name() + "_H_";
+    UpperString(&guard);
+    StripString(&guard, ".-", '_');
+
+    stub_h_->Print(
+        "$greeting$\n"
+        "#ifndef $guard$\n"
+        "#define $guard$\n\n"
+        "#include <stddef.h>\n"
+        "#include <stdint.h>\n\n"
+        "#include \"components/tracing/core/proto_zero_message.h\"\n\n",
+        "greeting", greeting,
+        "guard", guard);
+    stub_cc_->Print(
+        "$greeting$\n"
+        "// This file intentionally left blank.\n",
+        "greeting", greeting);
+
+    // Print namespaces.
+    for (const std::string& ns : namespaces_)
+      stub_h_->Print("namespace $ns$ {\n", "ns", ns);
+    stub_h_->Print("\n");
+    // Print forward declarations.
+    for (const Descriptor* message : messages_) {
+      stub_h_->Print(
+          "class $class$;\n",
+          "class", GetCppClassName(message));
+    }
+    stub_h_->Print("\n");
+  }
+
+  void GenerateEnumDescriptor(const EnumDescriptor* enumeration) {
+    stub_h_->Print(
+        "enum $class$ {\n",
+        "class", GetCppClassName(enumeration));
+    stub_h_->Indent();
+
+    std::string value_name_prefix;
+    if (enumeration->containing_type() != nullptr)
+        value_name_prefix = GetCppClassName(enumeration) + "_";
+
+    for (int i = 0; i < enumeration->value_count(); ++i) {
+      const EnumValueDescriptor* value = enumeration->value(i);
+      stub_h_->Print(
+          "$name$ = $number$,\n",
+          "name", value_name_prefix + value->name(),
+          "number", std::to_string(value->number()));
+    }
+
+    stub_h_->Outdent();
+    stub_h_->Print("};\n\n");
+  }
+
+  void GenerateSimpleFieldDescriptor(const FieldDescriptor* field) {
+    std::map<std::string, std::string> setter;
+    setter["id"] = std::to_string(field->number());
+    setter["name"] = field->name();
+    setter["action"] = field->is_repeated() ? "add" : "set";
+
+    std::string appender;
+    std::string cpp_type;
+
+    switch (field->type()) {
+      case FieldDescriptor::TYPE_BOOL: {
+        appender = "AppendBool";
+        cpp_type = "bool";
+        break;
+      }
+      case FieldDescriptor::TYPE_INT32: {
+        appender = "AppendInt32";
+        cpp_type = "int32_t";
+        break;
+      }
+      case FieldDescriptor::TYPE_INT64: {
+        appender = "AppendInt64";
+        cpp_type = "int64_t";
+        break;
+      }
+      case FieldDescriptor::TYPE_UINT32: {
+        appender = "AppendUint32";
+        cpp_type = "uint32_t";
+        break;
+      }
+      case FieldDescriptor::TYPE_UINT64: {
+        appender = "AppendUint64";
+        cpp_type = "uint64_t";
+        break;
+      }
+      case FieldDescriptor::TYPE_SINT32: {
+        appender = "AppendSint32";
+        cpp_type = "int32_t";
+        break;
+      }
+      case FieldDescriptor::TYPE_SINT64: {
+        appender = "AppendSint64";
+        cpp_type = "int64_t";
+        break;
+      }
+      case FieldDescriptor::TYPE_FIXED32: {
+        appender = "AppendFixed32";
+        cpp_type = "uint32_t";
+        break;
+      }
+      case FieldDescriptor::TYPE_FIXED64: {
+        appender = "AppendFixed64";
+        cpp_type = "uint64_t";
+        break;
+      }
+      case FieldDescriptor::TYPE_SFIXED32: {
+        appender = "AppendSfixed32";
+        cpp_type = "int32_t";
+        break;
+      }
+      case FieldDescriptor::TYPE_SFIXED64: {
+        appender = "AppendSfixed64";
+        cpp_type = "int64_t";
+        break;
+      }
+      case FieldDescriptor::TYPE_FLOAT: {
+        appender = "AppendFloat";
+        cpp_type = "float";
+        break;
+      }
+      case FieldDescriptor::TYPE_DOUBLE: {
+        appender = "AppendDouble";
+        cpp_type = "double";
+        break;
+      }
+      case FieldDescriptor::TYPE_ENUM: {
+        appender = IsTinyEnumField(field) ? "AppendTinyNumber" : "AppendInt32";
+        cpp_type = GetCppClassName(field->enum_type(), true);
+        break;
+      }
+      case FieldDescriptor::TYPE_STRING: {
+        appender = "AppendString";
+        cpp_type = "const char*";
+        break;
+      }
+      case FieldDescriptor::TYPE_BYTES: {
+        stub_h_->Print(
+            setter,
+            "void $action$_$name$(const uint8_t* data, size_t size) {\n"
+            "  // AppendBytes($id$, data, size);\n"
+            "}\n");
+        return;
+      }
+      default: {
+        Abort("Unsupported field type.");
+        return;
+      }
+    }
+    setter["appender"] = appender;
+    setter["cpp_type"] = cpp_type;
+    stub_h_->Print(
+        setter,
+        "void $action$_$name$($cpp_type$ value) {\n"
+        "  // $appender$($id$, value);\n"
+        "}\n");
+  }
+
+  void GenerateNestedMessageFieldDescriptor(const FieldDescriptor* field) {
+    stub_h_->Print(
+        "$class$* $action$_$name$() {\n"
+        "  return BeginNestedMessage<$class$>($id$);\n"
+        "}\n",
+        "id", std::to_string(field->number()),
+        "name", field->name(),
+        "action", field->is_repeated() ? "add" : "set",
+        "class", GetCppClassName(field->message_type()));
+  }
+
+  void GenerateMessageDescriptor(const Descriptor* message) {
+    stub_h_->Print(
+        "class $name$ : public ::tracing::v2::ProtoZeroMessage {\n"
+        " public:\n",
+        "name", GetCppClassName(message));
+    stub_h_->Indent();
+
+    // Using statements for nested messages.
+    for (int i = 0; i < message->nested_type_count(); ++i) {
+      const Descriptor* nested_message = message->nested_type(i);
+      stub_h_->Print(
+          "using $local_name$ = $global_name$;\n",
+          "local_name", nested_message->name(),
+          "global_name", GetCppClassName(nested_message, true));
+    }
+
+    // Using statements for nested enums.
+    for (int i = 0; i < message->enum_type_count(); ++i) {
+      const EnumDescriptor* nested_enum = message->enum_type(i);
+      stub_h_->Print(
+          "using $local_name$ = $global_name$;\n",
+          "local_name", nested_enum->name(),
+          "global_name", GetCppClassName(nested_enum, true));
+    }
+
+    // Values of nested enums.
+    for (int i = 0; i < message->enum_type_count(); ++i) {
+      const EnumDescriptor* nested_enum = message->enum_type(i);
+      std::string value_name_prefix = GetCppClassName(nested_enum) + "_";
+
+      for (int j = 0; j < nested_enum->value_count(); ++j) {
+        const EnumValueDescriptor* value = nested_enum->value(j);
+        stub_h_->Print(
+            "static const $class$ $name$ = $full_name$;\n",
+            "class", nested_enum->name(),
+            "name", value->name(),
+            "full_name", value_name_prefix + value->name());
+      }
+    }
+
+    // Fields descriptors.
+    for (int i = 0; i < message->field_count(); ++i) {
+      const FieldDescriptor* field = message->field(i);
+      if (field->is_packed()) {
+        Abort("Packed repeated fields are not supported.");
+        return;
+      }
+      if (field->type() != FieldDescriptor::TYPE_MESSAGE) {
+        GenerateSimpleFieldDescriptor(field);
+      } else {
+        GenerateNestedMessageFieldDescriptor(field);
+      }
+    }
+
+    stub_h_->Outdent();
+    stub_h_->Print("};\n\n");
+  }
+
+  void GenerateEpilogue() {
+    for (unsigned i = 0; i < namespaces_.size(); ++i) {
+      stub_h_->Print("} // Namespace.\n");
+    }
+    stub_h_->Print("#endif  // Include guard.\n");
+  }
+
+  const FileDescriptor* const source_;
   Printer* const stub_h_;
   Printer* const stub_cc_;
   std::string error_;
+
+  std::string package_;
+  std::vector<std::string> namespaces_;
+  std::string full_namespace_prefix_;
+  std::vector<const Descriptor*> messages_;
+  std::vector<const EnumDescriptor*> enums_;
 };
 
 }  // namespace
 
 ProtoZeroGenerator::ProtoZeroGenerator() {
 }
 
 ProtoZeroGenerator::~ProtoZeroGenerator() {
 }
 
@@ skipping 17 matching lines @@
   GeneratorJob job(file, &stub_h_printer, &stub_cc_printer);
   if (!job.GenerateStubs()) {
     *error = job.GetFirstError();
     return false;
   }
   return true;
 }
 
 }  // namespace proto
 }  // namespace tracing