From 7824a4f0cd294bd0ad45c5fe73451eed613cd624 Mon Sep 17 00:00:00 2001
From: Thomas Van Lenten <thomasvl@google.com>
Date: Fri, 22 Sep 2023 11:21:32 -0700
Subject: [PATCH] [ObjC] Handle transtive public imports in minimal imports
 mode

If using minimal imports, a non direct import might be needed because
it is a public import of an other wise not required header; ensure
the import makes it into the generated code to type definitions can
be found.

PiperOrigin-RevId: 567672195
---
 .../protobuf/compiler/objectivec/file.cc      | 20 +++++++++++++++++++
 .../protobuf/compiler/objectivec/file.h       |  8 +++++++-
 .../protobuf/compiler/objectivec/generator.cc |  5 ++++-
 3 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/src/google/protobuf/compiler/objectivec/file.cc b/src/google/protobuf/compiler/objectivec/file.cc
index 4eb7edbbb7..9c67ec0d87 100644
--- a/src/google/protobuf/compiler/objectivec/file.cc
+++ b/src/google/protobuf/compiler/objectivec/file.cc
@@ -19,6 +19,7 @@
 #include "absl/container/btree_set.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
+#include "absl/log/absl_check.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
 #include "google/protobuf/compiler/objectivec/enum.h"
@@ -526,8 +527,27 @@ void FileGenerator::GenerateFile(io::Printer* p, GeneratedFileType file_type,
       const FileDescriptor* dep = file_->dependency(i);
       if (file_imports.contains(dep)) {
         import_writer.AddFile(file_->dependency(i), header_extension);
+        file_imports.erase(dep);
       }
     }
+    if (!file_imports.empty()) {
+      // If there are still things in file_imports, then there were files that
+      // were public imports into the non public imports, add those files are
+      // needed to define the types also.
+      //
+      // Sort them (to get stable generation), and add them to the extra files
+      // to imports.
+
+      // This can really only happen in minimal imports mode, every other case,
+      // it shouldn't happen.
+      ABSL_CHECK(generation_options_.generate_minimal_imports);
+      std::vector<const FileDescriptor*> still_needed(file_imports.begin(),
+                                                      file_imports.end());
+      std::sort(still_needed.begin(), still_needed.end(),
+                FileDescriptorsOrderedByName());
+      extra_files_to_import.insert(extra_files_to_import.end(),
+                                   still_needed.begin(), still_needed.end());
+    }
   }
 
   for (const auto& dep : extra_files_to_import) {
diff --git a/src/google/protobuf/compiler/objectivec/file.h b/src/google/protobuf/compiler/objectivec/file.h
index 66cbd13911..d560a93e58 100644
--- a/src/google/protobuf/compiler/objectivec/file.h
+++ b/src/google/protobuf/compiler/objectivec/file.h
@@ -98,7 +98,13 @@ class FileGenerator {
       const std::vector<const FileDescriptor*>& deps_with_extensions) const;
   void EmitFileDescription(io::Printer* p) const;
 
-  enum class PublicDepsHandling : int { kAsUsed, kForceInclude, kExclude };
+  enum class PublicDepsHandling : int {
+    kAsUsed,        // No special handing, require references to import then.
+    kForceInclude,  // Always treat them as needed.
+    kExclude,       // Never treat them as needed.
+  };
+  // `public_deps_handling` controls how the public imports in this file should
+  // be handed.
   void DetermineNeededDeps(absl::flat_hash_set<const FileDescriptor*>* deps,
                            PublicDepsHandling public_deps_handling) const;
 
diff --git a/src/google/protobuf/compiler/objectivec/generator.cc b/src/google/protobuf/compiler/objectivec/generator.cc
index 7e09747769..4f74006c76 100644
--- a/src/google/protobuf/compiler/objectivec/generator.cc
+++ b/src/google/protobuf/compiler/objectivec/generator.cc
@@ -257,7 +257,10 @@ bool ObjectiveCGenerator::GenerateAll(
       // Controls if minimal imports should be generated from a files imports.
       // Since custom options require imports, they current cause generated
       // imports even though there is nothing captured in the generated code,
-      // this provides smaller imports only for the things referenced.
+      // this provides smaller imports only for the things referenced. This
+      // could break code in complex cases where code uses types via long
+      // import chains with public imports mixed through the way, as things
+      // that aren't really needed for the local usages could be pruned.
       if (!StringToBool(options[i].second,
                         &generation_options.generate_minimal_imports)) {
         *error =