From 3bc507d15fe5d510dd0f48b935605275a7af2b52 Mon Sep 17 00:00:00 2001
From: Thomas Van Lenten <thomasvl@google.com>
Date: Wed, 23 Aug 2023 07:34:51 -0700
Subject: [PATCH] [ObjC] Fix minimal imports for root generation.

Imports where the types aren't used, but they do vend extensions
also need to be imported into the generated source.

PiperOrigin-RevId: 559419899
---
 .../protobuf/compiler/objectivec/file.cc      | 35 ++++++++-----------
 .../protobuf/compiler/objectivec/file.h       |  2 +-
 2 files changed, 15 insertions(+), 22 deletions(-)

diff --git a/src/google/protobuf/compiler/objectivec/file.cc b/src/google/protobuf/compiler/objectivec/file.cc
index 34df0646e5..ea899a43b5 100644
--- a/src/google/protobuf/compiler/objectivec/file.cc
+++ b/src/google/protobuf/compiler/objectivec/file.cc
@@ -376,16 +376,7 @@ void FileGenerator::GenerateSource(io::Printer* p) const {
   std::vector<const FileDescriptor*> deps_with_extensions =
       common_state_->CollectMinimalFileDepsContainingExtensions(file_);
   GeneratedFileOptions file_options;
-
-  // If any indirect dependency provided extensions, it needs to be directly
-  // imported so it can get merged into the root's extensions registry.
-  // See the Note by CollectMinimalFileDepsContainingExtensions before
-  // changing this.
-  for (auto& dep : deps_with_extensions) {
-    if (!IsDirectDependency(dep, file_)) {
-      file_options.extra_files_to_import.push_back(dep);
-    }
-  }
+  file_options.forced_files_to_import = deps_with_extensions;
 
   absl::btree_set<std::string> fwd_decls;
   for (const auto& generator : message_generators_) {
@@ -432,16 +423,7 @@ void FileGenerator::GenerateGlobalSource(io::Printer* p) const {
   std::vector<const FileDescriptor*> deps_with_extensions =
       common_state_->CollectMinimalFileDepsContainingExtensions(file_);
   GeneratedFileOptions file_options;
-
-  // If any indirect dependency provided extensions, it needs to be directly
-  // imported so it can get merged into the root's extensions registry.
-  // See the Note by CollectMinimalFileDepsContainingExtensions before
-  // changing this.
-  for (auto& dep : deps_with_extensions) {
-    if (!IsDirectDependency(dep, file_)) {
-      file_options.extra_files_to_import.push_back(dep);
-    }
-  }
+  file_options.forced_files_to_import = deps_with_extensions;
 
   absl::btree_set<std::string> fwd_decls;
   for (const auto& generator : extension_generators_) {
@@ -551,7 +533,18 @@ void FileGenerator::GenerateFile(io::Printer* p, GeneratedFileType file_type,
       break;
   }
 
+  // If a forced file was a direct dep, move it into the file_imports.
+  std::vector<const FileDescriptor*> extra_files_to_import;
+  for (const auto& dep : file_options.forced_files_to_import) {
+    if (IsDirectDependency(dep, file_)) {
+      file_imports.insert(dep);
+    } else {
+      extra_files_to_import.push_back(dep);
+    }
+  }
+
   if (!file_imports.empty()) {
+    // Output the file_imports in the order they were listed as dependencies.
     for (int i = 0; i < file_->dependency_count(); i++) {
       const FileDescriptor* dep = file_->dependency(i);
       if (file_imports.contains(dep)) {
@@ -560,7 +553,7 @@ void FileGenerator::GenerateFile(io::Printer* p, GeneratedFileType file_type,
     }
   }
 
-  for (const auto& dep : file_options.extra_files_to_import) {
+  for (const auto& dep : extra_files_to_import) {
     import_writer.AddFile(dep, header_extension);
   }
 
diff --git a/src/google/protobuf/compiler/objectivec/file.h b/src/google/protobuf/compiler/objectivec/file.h
index 7fb3a6785a..cab51cfad5 100644
--- a/src/google/protobuf/compiler/objectivec/file.h
+++ b/src/google/protobuf/compiler/objectivec/file.h
@@ -100,7 +100,7 @@ class FileGenerator {
   enum class GeneratedFileType : int { kHeader, kSource };
   struct GeneratedFileOptions {
     std::vector<std::string> ignored_warnings;
-    std::vector<const FileDescriptor*> extra_files_to_import;
+    std::vector<const FileDescriptor*> forced_files_to_import;
     std::vector<std::string> extra_system_headers;
   };