1 files changed, 27 insertions, 31 deletions

diff --git a/lld/lib/ReaderWriter/MachO/MachONormalizedFileToAtoms.cpp b/lld/lib/ReaderWriter/MachO/MachONormalizedFileToAtoms.cpp
index 44e22ab6491..32ad88fe396 100644
--- a/lld/lib/ReaderWriter/MachO/MachONormalizedFileToAtoms.cpp
+++ b/lld/lib/ReaderWriter/MachO/MachONormalizedFileToAtoms.cpp
@@ -88,6 +88,7 @@ const MachORelocatableSectionToAtomType sectsToAtomType[] = {
   ENTRY("__DATA", "__thread_data", S_THREAD_LOCAL_REGULAR, typeTLVInitialData),
   ENTRY("__DATA", "__thread_bss",     S_THREAD_LOCAL_ZEROFILL,
                                                         typeTLVInitialZeroFill),
+  ENTRY("__DATA", "__objc_imageinfo", S_REGULAR,          typeObjCImageInfo),
   ENTRY("",       "",                 S_INTERPOSING,      typeInterposingTuples),
   ENTRY("__LD",   "__compact_unwind", S_REGULAR,
                                                          typeCompactUnwindInfo),
@@ -867,52 +868,35 @@ std::error_code addEHFrameReferences(const NormalizedFile &normalizedFile,
   return ehFrameErr;
 }
 
-std::error_code parseObjCImageInfo(const NormalizedFile &normalizedFile,
+std::error_code parseObjCImageInfo(const Section &sect,
+                                   const NormalizedFile &normalizedFile,
                                    MachOFile &file) {
 
-  const Section *imageInfoSection = nullptr;
-  for (auto &section : normalizedFile.sections) {
-    if (section.segmentName == "__OBJC" &&
-        section.sectionName == "__image_info") {
-      imageInfoSection = &section;
-      break;
-    }
-    if (section.segmentName == "__DATA" &&
-        section.sectionName == "__objc_imageinfo") {
-      imageInfoSection = &section;
-      break;
-    }
-  }
-
-  // No image info section so nothing to do.
-  if (!imageInfoSection)
-    return std::error_code();
-
   //	struct objc_image_info  {
   //		uint32_t	version;	// initially 0
   //		uint32_t	flags;
   //	};
 
-  ArrayRef<uint8_t> content = imageInfoSection->content;
+  ArrayRef<uint8_t> content = sect.content;
   if (content.size() != 8)
-    return make_dynamic_error_code(imageInfoSection->segmentName + "/" +
-                                   imageInfoSection->sectionName +
+    return make_dynamic_error_code(sect.segmentName + "/" +
+                                   sect.sectionName +
                                    " in file " + file.path() +
                                    " should be 8 bytes in size");
 
   const bool isBig = MachOLinkingContext::isBigEndian(normalizedFile.arch);
   uint32_t version = read32(content.data(), isBig);
   if (version)
-    return make_dynamic_error_code(imageInfoSection->segmentName + "/" +
-                                   imageInfoSection->sectionName +
+    return make_dynamic_error_code(sect.segmentName + "/" +
+                                   sect.sectionName +
                                    " in file " + file.path() +
                                    " should have version=0");
 
   uint32_t flags = read32(content.data() + 4, isBig);
   if (flags & (MachOLinkingContext::objc_supports_gc |
                MachOLinkingContext::objc_gc_only))
-    return make_dynamic_error_code(imageInfoSection->segmentName + "/" +
-                                   imageInfoSection->sectionName +
+    return make_dynamic_error_code(sect.segmentName + "/" +
+                                   sect.sectionName +
                                    " in file " + file.path() +
                                    " uses GC.  This is not supported");
 
@@ -951,6 +935,11 @@ dylibToAtoms(const NormalizedFile &normalizedFile, StringRef path,
 
 namespace normalized {
 
+static bool isObjCImageInfo(const Section &sect) {
+  return (sect.segmentName == "__OBJC" && sect.sectionName == "__image_info") ||
+    (sect.segmentName == "__DATA" && sect.sectionName == "__objc_imageinfo");
+}
+
 std::error_code
 normalizedObjectToAtoms(MachOFile *file,
                         const NormalizedFile &normalizedFile,
@@ -964,6 +953,18 @@ normalizedObjectToAtoms(MachOFile *file,
     DEBUG(llvm::dbgs() << "Creating atoms: "; sect.dump());
     if (isDebugInfoSection(sect))
       continue;
+
+
+    // If the file contains an objc_image_info struct, then we should parse the
+    // ObjC flags and Swift version.
+    if (isObjCImageInfo(sect)) {
+      if (std::error_code ec = parseObjCImageInfo(sect, normalizedFile, *file))
+        return ec;
+      // We then skip adding atoms for this section as we use the ObjCPass to
+      // re-emit this data after it has been aggregated for all files.
+      continue;
+    }
+
     bool customSectionName;
     DefinedAtom::ContentType atomType = atomTypeFromSection(sect,
                                                             customSectionName);
@@ -1007,11 +1008,6 @@ normalizedObjectToAtoms(MachOFile *file,
   if (std::error_code ec = addEHFrameReferences(normalizedFile, *file, *handler))
     return ec;
 
-  // If the file contains an objc_image_info struct, then we should parse the
-  // ObjC flags and Swift version.
-  if (std::error_code ec = parseObjCImageInfo(normalizedFile, *file))
-    return ec;
-
   // Process mach-o data-in-code regions array. That information is encoded in
   // atoms as References at each transition point.
   unsigned nextIndex = 0;