Initial Protozero extension support

In the initial version, both extended message and its extensions have to
reside in the same file; this limitation will be addressed in further
CLs.

The extension support is not used at the moment, but conformance test
has been added.

Bug: 156900028
Change-Id: I1a6801fae1751c16b3bcc8ac3c3fe8f9d1aa9ce9
diff --git a/Android.bp b/Android.bp
index 703adff..756ff40 100644
--- a/Android.bp
+++ b/Android.bp
@@ -6321,6 +6321,7 @@
 genrule {
   name: "perfetto_src_protozero_testing_messages_cpp_gen",
   srcs: [
+    "src/protozero/test/example_proto/extensions.proto",
     "src/protozero/test/example_proto/library.proto",
     "src/protozero/test/example_proto/library_internals/galaxies.proto",
     "src/protozero/test/example_proto/test_messages.proto",
@@ -6332,6 +6333,7 @@
   ],
   cmd: "mkdir -p $(genDir)/external/perfetto/ && $(location aprotoc) --proto_path=external/perfetto --plugin=protoc-gen-plugin=$(location perfetto_src_protozero_protoc_plugin_cppgen_plugin) --plugin_out=wrapper_namespace=gen:$(genDir)/external/perfetto/ $(in)",
   out: [
+    "external/perfetto/src/protozero/test/example_proto/extensions.gen.cc",
     "external/perfetto/src/protozero/test/example_proto/library.gen.cc",
     "external/perfetto/src/protozero/test/example_proto/library_internals/galaxies.gen.cc",
     "external/perfetto/src/protozero/test/example_proto/test_messages.gen.cc",
@@ -6343,6 +6345,7 @@
 genrule {
   name: "perfetto_src_protozero_testing_messages_cpp_gen_headers",
   srcs: [
+    "src/protozero/test/example_proto/extensions.proto",
     "src/protozero/test/example_proto/library.proto",
     "src/protozero/test/example_proto/library_internals/galaxies.proto",
     "src/protozero/test/example_proto/test_messages.proto",
@@ -6354,6 +6357,7 @@
   ],
   cmd: "mkdir -p $(genDir)/external/perfetto/ && $(location aprotoc) --proto_path=external/perfetto --plugin=protoc-gen-plugin=$(location perfetto_src_protozero_protoc_plugin_cppgen_plugin) --plugin_out=wrapper_namespace=gen:$(genDir)/external/perfetto/ $(in)",
   out: [
+    "external/perfetto/src/protozero/test/example_proto/extensions.gen.h",
     "external/perfetto/src/protozero/test/example_proto/library.gen.h",
     "external/perfetto/src/protozero/test/example_proto/library_internals/galaxies.gen.h",
     "external/perfetto/src/protozero/test/example_proto/test_messages.gen.h",
@@ -6369,6 +6373,7 @@
 genrule {
   name: "perfetto_src_protozero_testing_messages_lite_gen",
   srcs: [
+    "src/protozero/test/example_proto/extensions.proto",
     "src/protozero/test/example_proto/library.proto",
     "src/protozero/test/example_proto/library_internals/galaxies.proto",
     "src/protozero/test/example_proto/test_messages.proto",
@@ -6379,6 +6384,7 @@
   ],
   cmd: "mkdir -p $(genDir)/external/perfetto/ && $(location aprotoc) --proto_path=external/perfetto --cpp_out=lite=true:$(genDir)/external/perfetto/ $(in)",
   out: [
+    "external/perfetto/src/protozero/test/example_proto/extensions.pb.cc",
     "external/perfetto/src/protozero/test/example_proto/library.pb.cc",
     "external/perfetto/src/protozero/test/example_proto/library_internals/galaxies.pb.cc",
     "external/perfetto/src/protozero/test/example_proto/test_messages.pb.cc",
@@ -6390,6 +6396,7 @@
 genrule {
   name: "perfetto_src_protozero_testing_messages_lite_gen_headers",
   srcs: [
+    "src/protozero/test/example_proto/extensions.proto",
     "src/protozero/test/example_proto/library.proto",
     "src/protozero/test/example_proto/library_internals/galaxies.proto",
     "src/protozero/test/example_proto/test_messages.proto",
@@ -6400,6 +6407,7 @@
   ],
   cmd: "mkdir -p $(genDir)/external/perfetto/ && $(location aprotoc) --proto_path=external/perfetto --cpp_out=lite=true:$(genDir)/external/perfetto/ $(in)",
   out: [
+    "external/perfetto/src/protozero/test/example_proto/extensions.pb.h",
     "external/perfetto/src/protozero/test/example_proto/library.pb.h",
     "external/perfetto/src/protozero/test/example_proto/library_internals/galaxies.pb.h",
     "external/perfetto/src/protozero/test/example_proto/test_messages.pb.h",
@@ -6415,6 +6423,7 @@
 genrule {
   name: "perfetto_src_protozero_testing_messages_zero_gen",
   srcs: [
+    "src/protozero/test/example_proto/extensions.proto",
     "src/protozero/test/example_proto/library.proto",
     "src/protozero/test/example_proto/library_internals/galaxies.proto",
     "src/protozero/test/example_proto/test_messages.proto",
@@ -6426,6 +6435,7 @@
   ],
   cmd: "mkdir -p $(genDir)/external/perfetto/ && $(location aprotoc) --proto_path=external/perfetto --plugin=protoc-gen-plugin=$(location protozero_plugin) --plugin_out=wrapper_namespace=pbzero:$(genDir)/external/perfetto/ $(in)",
   out: [
+    "external/perfetto/src/protozero/test/example_proto/extensions.pbzero.cc",
     "external/perfetto/src/protozero/test/example_proto/library.pbzero.cc",
     "external/perfetto/src/protozero/test/example_proto/library_internals/galaxies.pbzero.cc",
     "external/perfetto/src/protozero/test/example_proto/test_messages.pbzero.cc",
@@ -6437,6 +6447,7 @@
 genrule {
   name: "perfetto_src_protozero_testing_messages_zero_gen_headers",
   srcs: [
+    "src/protozero/test/example_proto/extensions.proto",
     "src/protozero/test/example_proto/library.proto",
     "src/protozero/test/example_proto/library_internals/galaxies.proto",
     "src/protozero/test/example_proto/test_messages.proto",
@@ -6448,6 +6459,7 @@
   ],
   cmd: "mkdir -p $(genDir)/external/perfetto/ && $(location aprotoc) --proto_path=external/perfetto --plugin=protoc-gen-plugin=$(location protozero_plugin) --plugin_out=wrapper_namespace=pbzero:$(genDir)/external/perfetto/ $(in)",
   out: [
+    "external/perfetto/src/protozero/test/example_proto/extensions.pbzero.h",
     "external/perfetto/src/protozero/test/example_proto/library.pbzero.h",
     "external/perfetto/src/protozero/test/example_proto/library_internals/galaxies.pbzero.h",
     "external/perfetto/src/protozero/test/example_proto/test_messages.pbzero.h",
diff --git a/src/protozero/BUILD.gn b/src/protozero/BUILD.gn
index 18f8e7e..de6c11c 100644
--- a/src/protozero/BUILD.gn
+++ b/src/protozero/BUILD.gn
@@ -76,6 +76,7 @@
 
 perfetto_proto_library("testing_messages_@TYPE@") {
   sources = [
+    "test/example_proto/extensions.proto",
     "test/example_proto/library.proto",
     "test/example_proto/library_internals/galaxies.proto",
     "test/example_proto/test_messages.proto",
diff --git a/src/protozero/protoc_plugin/protozero_plugin.cc b/src/protozero/protoc_plugin/protozero_plugin.cc
index dd49d45..874ee61 100644
--- a/src/protozero/protoc_plugin/protozero_plugin.cc
+++ b/src/protozero/protoc_plugin/protozero_plugin.cc
@@ -97,6 +97,8 @@
       GenerateEnumDescriptor(enumeration);
     for (const Descriptor* message : messages_)
       GenerateMessageDescriptor(message);
+    for (auto key_value : extensions_)
+      GenerateExtension(key_value.first, key_value.second);
     GenerateEpilogue();
     return error_.empty();
   }
@@ -248,9 +250,29 @@
     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));
+
+      if (message->extension_count() > 0) {
+        if (message->field_count() > 0 || message->nested_type_count() > 0 ||
+            message->enum_type_count() > 0) {
+          Abort("message with extend blocks shouldn't contain anything else");
+        }
+
+        // Iterate over all fields in "extend" blocks.
+        for (int i = 0; i < message->extension_count(); ++i) {
+          const FieldDescriptor* extension = message->extension(i);
+
+          // Protoc plugin API does not group fields in "extend" blocks.
+          // As the support for extensions in protozero is limited, the code
+          // assumes that extend blocks are located inside a wrapper message and
+          // name of this message is used to group them.
+          std::string extension_name = extension->extension_scope()->name();
+          extensions_[extension_name].push_back(extension);
+        }
+      } else {
+        messages_.push_back(message);
+        for (int i = 0; i < message->nested_type_count(); ++i) {
+          stack.push_back(message->nested_type(i));
+        }
       }
     }
 
@@ -258,6 +280,9 @@
     for (int i = 0; i < source_->enum_type_count(); ++i)
       enums_.push_back(source_->enum_type(i));
 
+    if (source_->extension_count() > 0)
+      Abort("top-level extension blocks are not supported");
+
     for (const Descriptor* message : messages_) {
       for (int i = 0; i < message->enum_type_count(); ++i) {
         enums_.push_back(message->enum_type(i));
@@ -764,20 +789,68 @@
 
     // Field descriptors.
     for (int i = 0; i < message->field_count(); ++i) {
-      const FieldDescriptor* field = message->field(i);
-      if (field->is_packed()) {
-        GeneratePackedRepeatedFieldDescriptor(field);
-      } else if (field->type() != FieldDescriptor::TYPE_MESSAGE) {
-        GenerateSimpleFieldDescriptor(field);
-      } else {
-        GenerateNestedMessageFieldDescriptor(field);
-      }
+      GenerateFieldDescriptor(message->field(i));
     }
 
     stub_h_->Outdent();
     stub_h_->Print("};\n\n");
   }
 
+  void GenerateFieldDescriptor(const FieldDescriptor* field) {
+    if (field->is_packed()) {
+      GeneratePackedRepeatedFieldDescriptor(field);
+    } else if (field->type() != FieldDescriptor::TYPE_MESSAGE) {
+      GenerateSimpleFieldDescriptor(field);
+    } else {
+      GenerateNestedMessageFieldDescriptor(field);
+    }
+  }
+
+  // Generate extension class for a group of FieldDescriptor instances
+  // representing one "extend" block in proto definition. For example:
+  //
+  //   message SpecificExtension {
+  //     extend GeneralThing {
+  //       optional Fizz fizz = 101;
+  //       optional Buzz buzz = 102;
+  //     }
+  //   }
+  //
+  // This is going to be passed as a vector of two elements, "fizz" and
+  // "buzz". Wrapping message is used to provide a name for generated
+  // extension class.
+  //
+  // In the example above, generated code is going to look like:
+  //
+  //   class SpecificExtension : public GeneralThing {
+  //     Fizz* set_fizz();
+  //     Buzz* set_buzz();
+  //   }
+  void GenerateExtension(
+      const std::string& extension_name,
+      const std::vector<const FieldDescriptor*>& descriptors) {
+    // Use an arbitrary descriptor in order to get generic information not
+    // specific to any of them.
+    const FieldDescriptor* descriptor = descriptors[0];
+    const Descriptor* base_message = descriptor->containing_type();
+
+    // TODO(ddrone): ensure that this code works when containing_type located in
+    // other file or namespace.
+    stub_h_->Print("class $name$ : public $extendee$ {\n", "name",
+                   extension_name, "extendee", GetCppClassName(base_message));
+    stub_h_->Print(" public:\n");
+    stub_h_->Indent();
+    for (const FieldDescriptor* field : descriptors) {
+      if (field->containing_type() != base_message) {
+        Abort("one wrapper should extend only one message");
+        return;
+      }
+      GenerateFieldDescriptor(field);
+    }
+    stub_h_->Outdent();
+    stub_h_->Print("};\n");
+  }
+
   void GenerateEpilogue() {
     for (unsigned i = 0; i < namespaces_.size(); ++i) {
       stub_h_->Print("} // Namespace.\n");
@@ -795,6 +868,7 @@
   std::string full_namespace_prefix_;
   std::vector<const Descriptor*> messages_;
   std::vector<const EnumDescriptor*> enums_;
+  std::map<std::string, std::vector<const FieldDescriptor*>> extensions_;
 
   // The custom *Comp comparators are to ensure determinism of the generator.
   std::set<const FileDescriptor*, FileDescriptorComp> public_imports_;
diff --git a/src/protozero/test/example_proto/extensions.proto b/src/protozero/test/example_proto/extensions.proto
new file mode 100644
index 0000000..bbc344d
--- /dev/null
+++ b/src/protozero/test/example_proto/extensions.proto
@@ -0,0 +1,44 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+syntax = "proto2";
+
+package protozero.test.protos;
+
+message RealFakeEvent {
+  optional uint32 base_int = 1;
+  optional string base_string = 2;
+
+  extensions 10 to 99;
+}
+
+message SystemA {
+  optional uint32 int_a = 1;
+  optional string string_a = 2;
+}
+
+message SystemB {
+  optional uint32 int_b = 1;
+  optional string string_b = 2;
+}
+
+// Dummy message used for naming purposes.
+message BrowserExtension {
+  extend RealFakeEvent {
+    optional SystemA extension_a = 10;
+    optional SystemB extension_b = 11;
+  }
+}
diff --git a/src/protozero/test/protozero_conformance_unittest.cc b/src/protozero/test/protozero_conformance_unittest.cc
index 85eaf43..05efa8f 100644
--- a/src/protozero/test/protozero_conformance_unittest.cc
+++ b/src/protozero/test/protozero_conformance_unittest.cc
@@ -24,6 +24,8 @@
 #include "test/gtest_and_gmock.h"
 
 // Autogenerated headers in out/*/gen/
+#include "src/protozero/test/example_proto/extensions.pb.h"
+#include "src/protozero/test/example_proto/extensions.pbzero.h"
 #include "src/protozero/test/example_proto/library.pbzero.h"
 #include "src/protozero/test/example_proto/test_messages.pb.h"
 #include "src/protozero/test/example_proto/test_messages.pbzero.h"
@@ -116,6 +118,39 @@
   EXPECT_EQ(1000, gold_msg_a.super_nested().value_c());
 }
 
+TEST(ProtoZeroConformanceTest, Extensions) {
+  HeapBuffered<pbtest::BrowserExtension> msg_a{kChunkSize, kChunkSize};
+
+  msg_a->set_base_int(4);
+
+  pbtest::SystemA* msg_b = msg_a->set_extension_a();
+  msg_b->set_int_a(3);
+  msg_b->set_string_a("string a");
+
+  pbtest::SystemB* msg_c = msg_a->set_extension_b();
+  msg_c->set_int_b(10);
+  msg_c->set_string_b("string b");
+
+  msg_a->set_base_string("base string");
+
+  std::string serialized = msg_a.SerializeAsString();
+  pbgold::RealFakeEvent gold_msg_a;
+  gold_msg_a.ParseFromString(serialized);
+
+  EXPECT_EQ(gold_msg_a.base_int(), 4u);
+  EXPECT_EQ(gold_msg_a.base_string(), "base string");
+
+  pbgold::SystemA gold_msg_b =
+      gold_msg_a.GetExtension(pbgold::BrowserExtension::extension_a);
+  EXPECT_EQ(gold_msg_b.int_a(), 3u);
+  EXPECT_EQ(gold_msg_b.string_a(), "string a");
+
+  pbgold::SystemB gold_msg_c =
+      gold_msg_a.GetExtension(pbgold::BrowserExtension::extension_b);
+  EXPECT_EQ(gold_msg_c.int_b(), 10u);
+  EXPECT_EQ(gold_msg_c.string_b(), "string b");
+}
+
 TEST(ProtoZeroConformanceTest, Import) {
   // Test the includes for indirect public import: library.pbzero.h ->
   // library_internals/galaxies.pbzero.h -> upper_import.pbzero.h .