1 files changed, 39 insertions, 16 deletions

diff --git a/src/compiler/objective_c_plugin.cc b/src/compiler/objective_c_plugin.cc
index 3ccfd5b037..8de0997ebe 100644
--- a/src/compiler/objective_c_plugin.cc
+++ b/src/compiler/objective_c_plugin.cc
@@ -39,6 +39,12 @@
 #include "src/compiler/objective_c_generator.h"
 #include "src/compiler/objective_c_generator_helpers.h"
 
+#include <google/protobuf/compiler/objectivec/objectivec_helpers.h>
+
+using ::google::protobuf::compiler::objectivec::ProtobufLibraryFrameworkName;
+using ::google::protobuf::compiler::objectivec::
+    IsProtobufLibraryBundledProtoFile;
+
 class ObjectiveCGrpcGenerator : public grpc::protobuf::compiler::CodeGenerator {
  public:
   ObjectiveCGrpcGenerator() {}
@@ -48,7 +54,6 @@ class ObjectiveCGrpcGenerator : public grpc::protobuf::compiler::CodeGenerator {
                         const ::grpc::string &parameter,
                         grpc::protobuf::compiler::GeneratorContext *context,
                         ::grpc::string *error) const {
-
     if (file->service_count() == 0) {
       // No services.  Do nothing.
       return true;
@@ -61,18 +66,35 @@ class ObjectiveCGrpcGenerator : public grpc::protobuf::compiler::CodeGenerator {
       // Generate .pbrpc.h
 
       ::grpc::string imports = ::grpc::string("#import \"") + file_name +
-        ".pbobjc.h\"\n\n"
-        "#import <ProtoRPC/ProtoService.h>\n"
-        "#import <RxLibrary/GRXWriteable.h>\n"
-        "#import <RxLibrary/GRXWriter.h>\n";
+                               ".pbobjc.h\"\n\n"
+                               "#import <ProtoRPC/ProtoService.h>\n"
+                               "#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;
       for (int i = 0; i < file->dependency_count(); i++) {
-        ::grpc::string header = grpc_objective_c_generator::MessageHeaderName(
-            file->dependency(i));
-        proto_imports += ::grpc::string("#import \"") + header + "\"\n";
+        ::grpc::string header =
+            grpc_objective_c_generator::MessageHeaderName(file->dependency(i));
+        const grpc::protobuf::FileDescriptor *dependency = file->dependency(i);
+        if (IsProtobufLibraryBundledProtoFile(dependency)) {
+          ::grpc::string base_name = header;
+          grpc_generator::StripPrefix(&base_name, "google/protobuf/");
+          // create the import code snippet
+          proto_imports +=
+              "#if GPB_USE_PROTOBUF_FRAMEWORK_IMPORTS\n"
+              "  #import <" +
+              ::grpc::string(ProtobufLibraryFrameworkName) + "/" + base_name +
+              ">\n"
+              "#else\n"
+              "  #import \"" +
+              header +
+              "\"\n"
+              "#endif\n";
+        } else {
+          proto_imports += ::grpc::string("#import \"") + header + "\"\n";
+        }
       }
 
       ::grpc::string declarations;
@@ -81,21 +103,22 @@ class ObjectiveCGrpcGenerator : public grpc::protobuf::compiler::CodeGenerator {
         declarations += grpc_objective_c_generator::GetHeader(service);
       }
 
-      static const ::grpc::string kNonNullBegin = "\nNS_ASSUME_NONNULL_BEGIN\n\n";
+      static const ::grpc::string kNonNullBegin =
+          "\nNS_ASSUME_NONNULL_BEGIN\n\n";
       static const ::grpc::string kNonNullEnd = "\nNS_ASSUME_NONNULL_END\n";
 
-      Write(context, file_name + ".pbrpc.h",
-          imports + '\n' + proto_imports + '\n' + kNonNullBegin + 
-          declarations + kNonNullEnd);
+      Write(context, file_name + ".pbrpc.h", imports + '\n' + proto_imports +
+                                                 '\n' + kNonNullBegin +
+                                                 declarations + kNonNullEnd);
     }
 
     {
       // Generate .pbrpc.m
 
       ::grpc::string imports = ::grpc::string("#import \"") + file_name +
-        ".pbrpc.h\"\n\n"
-        "#import <ProtoRPC/ProtoRPC.h>\n"
-        "#import <RxLibrary/GRXWriter+Immediate.h>\n";
+                               ".pbrpc.h\"\n\n"
+                               "#import <ProtoRPC/ProtoRPC.h>\n"
+                               "#import <RxLibrary/GRXWriter+Immediate.h>\n";
 
       ::grpc::string definitions;
       for (int i = 0; i < file->service_count(); i++) {
@@ -112,7 +135,7 @@ class ObjectiveCGrpcGenerator : public grpc::protobuf::compiler::CodeGenerator {
  private:
   // Write the given code into the given file.
   void Write(grpc::protobuf::compiler::GeneratorContext *context,
-              const ::grpc::string &filename, const ::grpc::string &code) const {
+             const ::grpc::string &filename, const ::grpc::string &code) const {
     std::unique_ptr<grpc::protobuf::io::ZeroCopyOutputStream> output(
         context->Open(filename));
     grpc::protobuf::io::CodedOutputStream coded_out(output.get());