1 files changed, 45 insertions, 13 deletions

diff --git a/src/google/protobuf/compiler/objectivec/objectivec_file.cc b/src/google/protobuf/compiler/objectivec/objectivec_file.cc
index 6f80a0de..f0d9b4d5 100644
--- a/src/google/protobuf/compiler/objectivec/objectivec_file.cc
+++ b/src/google/protobuf/compiler/objectivec/objectivec_file.cc
@@ -188,6 +188,7 @@ bool IsDirectDependency(const FileDescriptor* dep, const FileDescriptor* file) {
 FileGenerator::FileGenerator(const FileDescriptor *file, const Options& options)
     : file_(file),
       root_class_name_(FileClassName(file)),
+      is_bundled_proto_(IsProtobufLibraryBundledProtoFile(file)),
       options_(options) {
   for (int i = 0; i < file_->enum_type_count(); i++) {
     EnumGenerator *generator = new EnumGenerator(file_->enum_type(i));
@@ -214,7 +215,17 @@ FileGenerator::~FileGenerator() {
 }
 
 void FileGenerator::GenerateHeader(io::Printer *printer) {
-  PrintFileRuntimePreamble(printer, "GPBProtocolBuffers.h");
+  std::set<string> headers;
+  // Generated files bundled with the library get minimal imports, everything
+  // else gets the wrapper so everything is usable.
+  if (is_bundled_proto_) {
+    headers.insert("GPBRootObject.h");
+    headers.insert("GPBMessage.h");
+    headers.insert("GPBDescriptor.h");
+  } else {
+    headers.insert("GPBProtocolBuffers.h");
+  }
+  PrintFileRuntimePreamble(printer, headers);
 
   // Add some verification that the generated code matches the source the
   // code is being compiled with.
@@ -236,7 +247,8 @@ void FileGenerator::GenerateHeader(io::Printer *printer) {
   {
     ImportWriter import_writer(
         options_.generate_for_named_framework,
-        options_.named_framework_to_proto_path_mappings_path);
+        options_.named_framework_to_proto_path_mappings_path,
+        is_bundled_proto_);
     const string header_extension(kHeaderExtension);
     for (int i = 0; i < file_->public_dependency_count(); i++) {
       import_writer.AddFile(file_->public_dependency(i), header_extension);
@@ -337,7 +349,9 @@ void FileGenerator::GenerateHeader(io::Printer *printer) {
 
 void FileGenerator::GenerateSource(io::Printer *printer) {
   // #import the runtime support.
-  PrintFileRuntimePreamble(printer, "GPBProtocolBuffers_RuntimeSupport.h");
+  std::set<string> headers;
+  headers.insert("GPBProtocolBuffers_RuntimeSupport.h");
+  PrintFileRuntimePreamble(printer, headers);
 
   // Enums use atomic in the generated code, so add the system import as needed.
   if (FileContainsEnums(file_)) {
@@ -352,7 +366,8 @@ void FileGenerator::GenerateSource(io::Printer *printer) {
   {
     ImportWriter import_writer(
         options_.generate_for_named_framework,
-        options_.named_framework_to_proto_path_mappings_path);
+        options_.named_framework_to_proto_path_mappings_path,
+        is_bundled_proto_);
     const string header_extension(kHeaderExtension);
 
     // #import the header for this proto file.
@@ -566,7 +581,7 @@ void FileGenerator::GenerateSource(io::Printer *printer) {
 // files. This currently only supports the runtime coming from a framework
 // as defined by the official CocoaPod.
 void FileGenerator::PrintFileRuntimePreamble(
-    io::Printer* printer, const string& header_to_import) const {
+    io::Printer* printer, const std::set<string>& headers_to_import) const {
   printer->Print(
       "// Generated by the protocol buffer compiler.  DO NOT EDIT!\n"
       "// source: $filename$\n"
@@ -575,6 +590,7 @@ void FileGenerator::PrintFileRuntimePreamble(
 
   const string framework_name(ProtobufLibraryFrameworkName);
   const string cpp_symbol(ProtobufFrameworkImportSymbol(framework_name));
+
   printer->Print(
       "// This CPP symbol can be defined to use imports that match up to the framework\n"
       "// imports needed when using CocoaPods.\n"
@@ -582,15 +598,31 @@ void FileGenerator::PrintFileRuntimePreamble(
       " #define $cpp_symbol$ 0\n"
       "#endif\n"
       "\n"
-      "#if $cpp_symbol$\n"
-      " #import <$framework_name$/$header$>\n"
-      "#else\n"
-      " #import \"$header$\"\n"
+      "#if $cpp_symbol$\n",
+      "cpp_symbol", cpp_symbol);
+
+
+  for (std::set<string>::const_iterator iter = headers_to_import.begin();
+       iter != headers_to_import.end(); ++iter) {
+    printer->Print(
+        " #import <$framework_name$/$header$>\n",
+        "header", *iter,
+        "framework_name", framework_name);
+  }
+
+  printer->Print(
+      "#else\n");
+
+  for (std::set<string>::const_iterator iter = headers_to_import.begin();
+       iter != headers_to_import.end(); ++iter) {
+    printer->Print(
+        " #import \"$header$\"\n",
+        "header", *iter);
+  }
+
+  printer->Print(
       "#endif\n"
-      "\n",
-      "cpp_symbol", cpp_symbol,
-      "header", header_to_import,
-      "framework_name", framework_name);
+      "\n");
 }
 
 }  // namespace objectivec