Tracing V2: Import support for protobuf plugin.

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 <set>
 #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::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 {
 
+inline std::string ProtoStubName(const FileDescriptor* proto) {
+  return StripSuffixString(proto->name(), ".proto") + ".pbzero";
+}
+
 class GeneratorJob {
  public:
   GeneratorJob(const FileDescriptor *file,
                Printer* stub_h_printer,
                Printer* stub_cc_printer)
     : source_(file),
       stub_h_(stub_h_printer),
       stub_cc_(stub_cc_printer) {}
 
   bool GenerateStubs() {
@@ skipping 49 matching lines @@
     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 + "::";
-
+  void CollectDescriptors() {
     // 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));
       }
     }
   }
 
+  void CollectDependencies() {
+    // Public import basically means that callers only need to import this
+    // proto in order to use the stuff publicly imported by this proto.
+    for (int i = 0; i < source_->public_dependency_count(); ++i)
+      public_imports_.insert(source_->public_dependency(i));
+
+    if (source_->weak_dependency_count() > 0)
+      Abort("Weak imports are not supported.");
+
+    // Sanity check. Collect public imports (of collected imports) in DFS order.
+    // Visibilty for current proto:
+    // - all imports listed in current proto,
+    // - public imports of everything imported (recursive).
+    std::vector<const FileDescriptor*> stack;
+    for (int i = 0; i < source_->dependency_count(); ++i) {
+      const FileDescriptor* import = source_->dependency(i);
+      stack.push_back(import);
+      if (public_imports_.count(import) == 0) {
+        private_imports_.insert(import);
+      }
+    }
+
+    while (!stack.empty()) {
+      const FileDescriptor* import = stack.back();
+      stack.pop_back();
+      // Having imports under different packages leads to unnecessary
+      // complexity with namespaces.
+      if (import->package() != package_)
+        Abort("Imported proto must be in the same package.");
+
+      for (int i = 0; i < import->public_dependency_count(); ++i) {
+        stack.push_back(import->public_dependency(i));
+      }
+    }
+
+    // Collect descriptors of messages and enums used in current proto.
+    // It will be used to generate necessary forward declarations and performed
+    // sanity check guarantees that everything lays in the same namespace.
+    for (const Descriptor* message : messages_) {
+      for (int i = 0; i < message->field_count(); ++i) {
+        const FieldDescriptor* field = message->field(i);
+
+        if (field->type() == FieldDescriptor::TYPE_MESSAGE) {
+          if (public_imports_.count(field->message_type()->file()) == 0) {
+            // Avoid multiple forward declarations since
+            // public imports have been already included.
+            referenced_messages_.insert(field->message_type());
+          }
+        } else if (field->type() == FieldDescriptor::TYPE_ENUM) {
+          if (public_imports_.count(field->enum_type()->file()) == 0) {
+            referenced_enums_.insert(field->enum_type());
+          }
+        }
+      }
+    }
+  }
+
+  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 + "::";
+    CollectDescriptors();
+    CollectDependencies();
+  }
+
   // 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",
+        "#include \"components/tracing/core/proto_zero_message.h\"\n",
         "greeting", greeting,
         "guard", guard);
     stub_cc_->Print(
         "$greeting$\n"
-        "// This file intentionally left blank.\n",
-        "greeting", greeting);
+        "#include \"$name$.h\"\n",
+        "greeting", greeting,
+        "name", ProtoStubName(source_));
+
+    // Print includes for public imports.
+    for (const FileDescriptor* dependency : public_imports_) {
+      // Dependency name could contatin slashes but importing from upper-level
+      // directories is not possible anyway since build system process each
+      // proto file individually. Hence proto lookup path always equal to the
+      // directory where particular proto file is located and protoc does not
+      // allow reference to upper directory (aka ..) in import path.
+      //
+      // Laconically said:
+      // - source_->name() may never have slashes,
+      // - dependency->name() may have slashes but always reffers to inner path.
+      stub_h_->Print(
+          "#include \"$name$.h\"\n",
+          "name", ProtoStubName(dependency));
+    }
+    stub_h_->Print("\n");
+
+    // Print includes for private imports to .cc file.
+    for (const FileDescriptor* dependency : private_imports_) {
+      stub_cc_->Print(
+         "#include \"$name$.h\"\n",
+         "name", ProtoStubName(dependency));
+    }
+    stub_cc_->Print("\n");
 
     // Print namespaces.
-    for (const std::string& ns : namespaces_)
+    for (const std::string& ns : namespaces_) {
       stub_h_->Print("namespace $ns$ {\n", "ns", ns);
+      stub_cc_->Print("namespace $ns$ {\n", "ns", ns);
+    }
     stub_h_->Print("\n");
+    stub_cc_->Print("\n");
+
     // Print forward declarations.
-    for (const Descriptor* message : messages_) {
+    for (const Descriptor* message : referenced_messages_) {
       stub_h_->Print(
           "class $class$;\n",
           "class", GetCppClassName(message));
     }
+    for (const EnumDescriptor* enumeration : referenced_enums_) {
+      stub_h_->Print(
+          "enum $class$ : int32_t;\n",
+          "class", GetCppClassName(enumeration));
+    }
     stub_h_->Print("\n");
   }
 
   void GenerateEnumDescriptor(const EnumDescriptor* enumeration) {
     stub_h_->Print(
-        "enum $class$ {\n",
+        "enum $class$ : int32_t {\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(
@@ skipping 107 matching lines @@
     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) {
+    std::string action = field->is_repeated() ? "add" : "set";
+    std::string inner_class = GetCppClassName(field->message_type());
+    std::string outer_class = GetCppClassName(field->containing_type());
+
     stub_h_->Print(
-        "$class$* $action$_$name$() {\n"
-        "  return BeginNestedMessage<$class$>($id$);\n"
-        "}\n",
+        "$inner_class$* $action$_$name$();\n",
+        "name", field->name(),
+        "action", action,
+        "inner_class", inner_class);
+    stub_cc_->Print(
+        "$inner_class$* $outer_class$::$action$_$name$() {\n"
+        "  return BeginNestedMessage<$inner_class$>($id$);\n"
+        "}\n\n",
         "id", std::to_string(field->number()),
         "name", field->name(),
-        "action", field->is_repeated() ? "add" : "set",
-        "class", GetCppClassName(field->message_type()));
+        "action", action,
+        "inner_class", inner_class,
+        "outer_class", outer_class);
   }
 
   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.
@@ skipping 43 matching lines @@
       }
     }
 
     stub_h_->Outdent();
     stub_h_->Print("};\n\n");
   }
 
   void GenerateEpilogue() {
     for (unsigned i = 0; i < namespaces_.size(); ++i) {
       stub_h_->Print("} // Namespace.\n");
+      stub_cc_->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_;
+
+  std::set<const FileDescriptor*> public_imports_;
+  std::set<const FileDescriptor*> private_imports_;
+  std::set<const Descriptor*> referenced_messages_;
+  std::set<const EnumDescriptor*> referenced_enums_;
 };
 
 }  // namespace
 
 ProtoZeroGenerator::ProtoZeroGenerator() {
 }
 
 ProtoZeroGenerator::~ProtoZeroGenerator() {
 }
 
 bool ProtoZeroGenerator::Generate(const FileDescriptor* file,
                                   const std::string& options,
                                   GeneratorContext* context,
                                   std::string* error) const {
 
-  const std::string proto_stubs_name =
-      StripSuffixString(file->name(), ".proto") + ".pbzero";
-
   const std::unique_ptr<ZeroCopyOutputStream> stub_h_file_stream(
-      context->Open(proto_stubs_name + ".h"));
+      context->Open(ProtoStubName(file) + ".h"));
   const std::unique_ptr<ZeroCopyOutputStream> stub_cc_file_stream(
-      context->Open(proto_stubs_name + ".cc"));
+      context->Open(ProtoStubName(file) + ".cc"));
 
   // Variables are delimited by $.
   Printer stub_h_printer(stub_h_file_stream.get(), '$');
   Printer stub_cc_printer(stub_cc_file_stream.get(), '$');
 
   GeneratorJob job(file, &stub_h_printer, &stub_cc_printer);
   if (!job.GenerateStubs()) {
     *error = job.GetFirstError();
     return false;
   }
   return true;
 }
 
 }  // namespace proto
 }  // namespace tracing