4 files changed, 22 insertions, 11 deletions

diff --git a/lld/lib/ReaderWriter/MachO/ArchHandler.cpp b/lld/lib/ReaderWriter/MachO/ArchHandler.cpp
index 3fa1fb3b91f..3ccb1712064 100644
--- a/lld/lib/ReaderWriter/MachO/ArchHandler.cpp
+++ b/lld/lib/ReaderWriter/MachO/ArchHandler.cpp
@@ -153,6 +153,18 @@ bool ArchHandler::isDwarfCIE(bool swap, const DefinedAtom *atom) {
   return read32(swap, *(uint32_t *)(atom->rawContent().data() + idOffset)) == 0;
 }
 
+const Atom *ArchHandler::fdeTargetFunction(const DefinedAtom *fde) {
+  for (auto ref : *fde) {
+    if (ref->kindNamespace() == Reference::KindNamespace::mach_o &&
+        ref->kindValue() == unwindRefToFunctionKind()) {
+      assert(ref->kindArch() == kindArch() && "unexpected Reference arch");
+      return ref->target();
+    }
+  }
+
+  return nullptr;
+}
+
 } // namespace mach_o
 } // namespace lld
 
diff --git a/lld/lib/ReaderWriter/MachO/ArchHandler.h b/lld/lib/ReaderWriter/MachO/ArchHandler.h
index f4e7bf7a691..9613f7a29bc 100644
--- a/lld/lib/ReaderWriter/MachO/ArchHandler.h
+++ b/lld/lib/ReaderWriter/MachO/ArchHandler.h
@@ -88,6 +88,8 @@ public:
   /// __eh_frame.
   virtual Reference::KindValue unwindRefToEhFrameKind() = 0;
 
+  virtual const Atom *fdeTargetFunction(const DefinedAtom *fde);
+
   /// Used by normalizedFromAtoms() to know where to generated rebasing and 
   /// binding info in final executables.
   virtual bool isPointer(const Reference &) = 0;
diff --git a/lld/lib/ReaderWriter/MachO/CompactUnwindPass.cpp b/lld/lib/ReaderWriter/MachO/CompactUnwindPass.cpp
index 59fb3195f8d..77bd3e9bc43 100644
--- a/lld/lib/ReaderWriter/MachO/CompactUnwindPass.cpp
+++ b/lld/lib/ReaderWriter/MachO/CompactUnwindPass.cpp
@@ -411,18 +411,13 @@ private:
   collectDwarfFrameEntries(std::unique_ptr<MutableFile> &mergedFile,
                            std::map<const Atom *, const Atom *> &dwarfFrames) {
     for (const DefinedAtom *ehFrameAtom : mergedFile->defined()) {
-      if (ehFrameAtom->contentType() != DefinedAtom::typeCFI ||
-          ArchHandler::isDwarfCIE(_swap, ehFrameAtom))
+      if (ehFrameAtom->contentType() != DefinedAtom::typeCFI)
+        continue;
+      if (ArchHandler::isDwarfCIE(_swap, ehFrameAtom))
         continue;
 
-      DefinedAtom::reference_iterator ref = ehFrameAtom->begin();
-      for (; ref != ehFrameAtom->end(); ++ref)
-        if (ref->kindNamespace() == Reference::KindNamespace::mach_o &&
-            ref->kindArch() == _archHandler.kindArch() &&
-            ref->kindValue() == _archHandler.unwindRefToFunctionKind()) {
-          dwarfFrames.insert(std::make_pair(ref->target(), ehFrameAtom));
-          break;
-        }
+      if (const Atom *function = _archHandler.fdeTargetFunction(ehFrameAtom))
+        dwarfFrames[function] = ehFrameAtom;
     }
   }
 
diff --git a/lld/lib/ReaderWriter/MachO/MachONormalizedFileToAtoms.cpp b/lld/lib/ReaderWriter/MachO/MachONormalizedFileToAtoms.cpp
index d52af892b94..d0536904bac 100644
--- a/lld/lib/ReaderWriter/MachO/MachONormalizedFileToAtoms.cpp
+++ b/lld/lib/ReaderWriter/MachO/MachONormalizedFileToAtoms.cpp
@@ -677,7 +677,9 @@ std::error_code addEHFrameReferences(const NormalizedFile &normalizedFile,
                        addend, handler.kindArch());
 
     // Linker needs to fixup reference from the FDE to the function it's
-    // describing.
+    // describing. FIXME: there are actually different ways to do this, and the
+    // particular method used is specified in the CIE's augmentation fields
+    // (hopefully)
     uint64_t rangeFieldInFDE = cieFieldInFDE + sizeof(uint32_t);
 
     int64_t functionFromFDE = readSPtr(is64, swap, frameData + rangeFieldInFDE);