[ObjC] Only include system headers in multi source where needed.

- Add a new option for extra system headers.
- Have the two places that generate the enum support add the extra header.

PiperOrigin-RevId: 508143253

src/google/protobuf/compiler/objectivec/file.cc

@@ -346,31 +346,36 @@ void FileGenerator::GenerateSource(io::Printer* p) const {
     file_options.ignored_warnings.push_back("dollar-in-identifier-extension");
   }
 
-  GenerateFile(
-      p, GeneratedFileType::kSource, file_options, [&] {
-        if (!fwd_decls.empty()) {
-          p->Print(
-              // clang-format off
-              "#pragma mark - Objective C Class declarations\n"
-              "// Forward declarations of Objective C classes that we can use as\n"
-              "// static values in struct initializers.\n"
-              "// We don't use [Foo class] because it is not a static value.\n"
-              "$fwd_decls$\n"
-              "\n",
-              // clang-format on
-              "fwd_decls", absl::StrJoin(fwd_decls, "\n"));
-        }
+  // Enum implementation uses atomic in the generated code, so add
+  // the system import as needed.
+  if (!enum_generators_.empty()) {
+    file_options.extra_system_headers.push_back("stdatomic.h");
+  }
 
-        PrintRootImplementation(p, deps_with_extensions);
-        PrintFileDescription(p);
+  GenerateFile(p, GeneratedFileType::kSource, file_options, [&] {
+    if (!fwd_decls.empty()) {
+      p->Print(
+          // clang-format off
+          "#pragma mark - Objective C Class declarations\n"
+          "// Forward declarations of Objective C classes that we can use as\n"
+          "// static values in struct initializers.\n"
+          "// We don't use [Foo class] because it is not a static value.\n"
+          "$fwd_decls$\n"
+          "\n",
+          // clang-format on
+          "fwd_decls", absl::StrJoin(fwd_decls, "\n"));
+    }
 
-        for (const auto& generator : enum_generators_) {
-          generator->GenerateSource(p);
-        }
-        for (const auto& generator : message_generators_) {
-          generator->GenerateSource(p);
-        }
-      });
+    PrintRootImplementation(p, deps_with_extensions);
+    PrintFileDescription(p);
+
+    for (const auto& generator : enum_generators_) {
+      generator->GenerateSource(p);
+    }
+    for (const auto& generator : message_generators_) {
+      generator->GenerateSource(p);
+    }
+  });
 }
 
 void FileGenerator::GenerateGlobalSource(io::Printer* p) const {
@@ -397,27 +402,30 @@ void FileGenerator::GenerateGlobalSource(io::Printer* p) const {
     file_options.ignored_warnings.push_back("dollar-in-identifier-extension");
   }
 
-  GenerateFile(
-      p, GeneratedFileType::kSource, file_options, [&] {
-        if (!fwd_decls.empty()) {
-          p->Print(
-              // clang-format off
+  GenerateFile(p, GeneratedFileType::kSource, file_options, [&] {
+    if (!fwd_decls.empty()) {
+      p->Print(
+          // clang-format off
               "#pragma mark - Objective C Class declarations\n"
               "// Forward declarations of Objective C classes that we can use as\n"
               "// static values in struct initializers.\n"
               "// We don't use [Foo class] because it is not a static value.\n"
               "$fwd_decls$\n"
               "\n",
-              // clang-format on
-              "fwd_decls", absl::StrJoin(fwd_decls, "\n"));
-        }
+          // clang-format on
+          "fwd_decls", absl::StrJoin(fwd_decls, "\n"));
+    }
 
-        PrintRootImplementation(p, deps_with_extensions);
-      });
+    PrintRootImplementation(p, deps_with_extensions);
+  });
 }
 
 void FileGenerator::GenerateSourceForEnums(io::Printer* p) const {
-  GenerateFile(p, GeneratedFileType::kSource, [&] {
+  // Enum implementation uses atomic in the generated code.
+  GeneratedFileOptions file_options;
+  file_options.extra_system_headers.push_back("stdatomic.h");
+
+  GenerateFile(p, GeneratedFileType::kSource, file_options, [&] {
     for (const auto& generator : enum_generators_) {
       generator->GenerateSource(p);
     }
@@ -545,10 +553,11 @@ void FileGenerator::GenerateFile(io::Printer* p, GeneratedFileType file_type,
       "google_protobuf_objc_version",
       absl::StrCat(GOOGLE_PROTOBUF_OBJC_VERSION));
 
-  // Enum implementation uses atomic in the generated code, so add
-  // the system import as needed.
-  if (file_type == GeneratedFileType::kSource && !enum_generators_.empty()) {
-    p->Print("#import <stdatomic.h>\n\n");
+  if (!file_options.extra_system_headers.empty()) {
+    for (const auto& system_header : file_options.extra_system_headers) {
+      p->Print("#import <$header$>\n", "header", system_header);
+    }
+    p->Print("\n");
   }
 
   import_writer.PrintFileImports(p);

src/google/protobuf/compiler/objectivec/file.h

@@ -96,6 +96,7 @@ class FileGenerator {
   struct GeneratedFileOptions {
     std::vector<std::string> ignored_warnings;
     std::vector<const FileDescriptor*> extra_files_to_import;
+    std::vector<std::string> extra_system_headers;
   };
 
   void GenerateFile(io::Printer* p, GeneratedFileType file_type,