Tracing V2: Import support for protobuf plugin.

components/tracing/tools/proto_zero_plugin/proto_zero_generator.cc

Index: components/tracing/tools/proto_zero_plugin/proto_zero_generator.cc
diff --git a/components/tracing/tools/proto_zero_plugin/proto_zero_generator.cc b/components/tracing/tools/proto_zero_plugin/proto_zero_generator.cc
index 12a32ab3783581e3665d5da37145670dc847255d..2165673b362eeefc3bbb110ca2e0be0fa6aac01c 100644
--- a/components/tracing/tools/proto_zero_plugin/proto_zero_generator.cc
+++ b/components/tracing/tools/proto_zero_plugin/proto_zero_generator.cc
@@ -6,6 +6,7 @@
 
 #include <map>
 #include <memory>
+#include <set>
 #include <string>
 
 #include "third_party/protobuf/src/google/protobuf/descriptor.h"
@@ -33,6 +34,10 @@ 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,
@@ -102,14 +107,7 @@ class GeneratorJob {
     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)
@@ -135,6 +133,74 @@ class GeneratorJob {
     }
   }
 
+  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 =
@@ -151,30 +217,65 @@ class GeneratorJob {
         "#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();
 
@@ -302,14 +403,24 @@ class GeneratorJob {
   }
 
   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) {
@@ -373,6 +484,7 @@ class GeneratorJob {
   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");
   }
@@ -387,6 +499,11 @@ class GeneratorJob {
   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
@@ -402,13 +519,10 @@ bool ProtoZeroGenerator::Generate(const FileDescriptor* file,
                                   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(), '$');