From 05e0d2da35d822e0d3cad368ceb3848c6e035e4a Mon Sep 17 00:00:00 2001
From: Muxi Yan <mxyan@google.com>
Date: Mon, 18 Sep 2017 17:30:20 -0700
Subject: [PATCH] Generate forward declaration in pbrpc.h

---
 src/compiler/objective_c_generator.cc | 21 +++++++++++++++++++++
 src/compiler/objective_c_generator.h  |  4 ++++
 src/compiler/objective_c_plugin.cc    | 18 ++++++++++--------
 3 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/src/compiler/objective_c_generator.cc b/src/compiler/objective_c_generator.cc
index c05d15375b3..33b5fedd5c8 100644
--- a/src/compiler/objective_c_generator.cc
+++ b/src/compiler/objective_c_generator.cc
@@ -17,6 +17,7 @@
  */
 
 #include <map>
+#include <set>
 #include <sstream>
 
 #include "src/compiler/config.h"
@@ -29,7 +30,9 @@ using ::google::protobuf::compiler::objectivec::ClassName;
 using ::grpc::protobuf::io::Printer;
 using ::grpc::protobuf::MethodDescriptor;
 using ::grpc::protobuf::ServiceDescriptor;
+using ::grpc::protobuf::FileDescriptor;
 using ::std::map;
+using ::std::set;
 
 namespace grpc_objective_c_generator {
 namespace {
@@ -190,6 +193,24 @@ void PrintMethodImplementations(Printer *printer,
 
 }  // namespace
 
+::grpc::string GetAllMessageClasses(const FileDescriptor *file) {
+  ::grpc::string output;
+  set< ::grpc::string> classes;
+  for (int i = 0; i < file->service_count(); i++) {
+    const auto service = file->service(i);
+    for (int i = 0; i < service->method_count(); i++) {
+      const auto method = service->method(i);
+      classes.insert(ClassName(method->input_type()));
+      classes.insert(ClassName(method->output_type()));
+    }
+  }
+  for (auto one_class : classes) {
+    output += "  @class " + one_class + ";\n";
+  }
+
+  return output;
+}
+
 ::grpc::string GetHeader(const ServiceDescriptor *service) {
   ::grpc::string output;
   {
diff --git a/src/compiler/objective_c_generator.h b/src/compiler/objective_c_generator.h
index edbee7ff524..e912a52415f 100644
--- a/src/compiler/objective_c_generator.h
+++ b/src/compiler/objective_c_generator.h
@@ -24,8 +24,12 @@
 namespace grpc_objective_c_generator {
 
 using ::grpc::protobuf::ServiceDescriptor;
+using ::grpc::protobuf::FileDescriptor;
 using ::grpc::string;
 
+// Returns forward declaration of classes in the generated header file.
+string GetAllMessageClasses(const FileDescriptor *file);
+
 // Returns the content to be included in the "global_scope" insertion point of
 // the generated header file.
 string GetHeader(const ServiceDescriptor *service);
diff --git a/src/compiler/objective_c_plugin.cc b/src/compiler/objective_c_plugin.cc
index 96a3375e962..de878cf143b 100644
--- a/src/compiler/objective_c_plugin.cc
+++ b/src/compiler/objective_c_plugin.cc
@@ -58,9 +58,10 @@ class ObjectiveCGrpcGenerator : public grpc::protobuf::compiler::CodeGenerator {
                                "#import <RxLibrary/GRXWriteable.h>\n"
                                "#import <RxLibrary/GRXWriter.h>\n";
 
-      // TODO(jcanizales): Instead forward-declare the input and output types
-      // and import the files in the .pbrpc.m
       ::grpc::string proto_imports;
+      proto_imports += "#if GPB_GRPC_FORWARD_DECLARE_MESSAGE_PROTO\n" +
+                       grpc_objective_c_generator::GetAllMessageClasses(file) +
+                       "#else\n";
       for (int i = 0; i < file->dependency_count(); i++) {
         ::grpc::string header =
             grpc_objective_c_generator::MessageHeaderName(file->dependency(i));
@@ -70,19 +71,20 @@ class ObjectiveCGrpcGenerator : public grpc::protobuf::compiler::CodeGenerator {
           grpc_generator::StripPrefix(&base_name, "google/protobuf/");
           // create the import code snippet
           proto_imports +=
-              "#if GPB_USE_PROTOBUF_FRAMEWORK_IMPORTS\n"
-              "  #import <" +
+              "  #if GPB_USE_PROTOBUF_FRAMEWORK_IMPORTS\n"
+              "    #import <" +
               ::grpc::string(ProtobufLibraryFrameworkName) + "/" + base_name +
               ">\n"
-              "#else\n"
-              "  #import \"" +
+              "  #else\n"
+              "    #import \"" +
               header +
               "\"\n"
-              "#endif\n";
+              "  #endif\n";
         } else {
-          proto_imports += ::grpc::string("#import \"") + header + "\"\n";
+          proto_imports += ::grpc::string("  #import \"") + header + "\"\n";
         }
       }
+      proto_imports += "#endif\n";
 
       ::grpc::string declarations;
       for (int i = 0; i < file->service_count(); i++) {