[lld-macho] Fix compatibility between --icf=safe_thunks and --keep-icf-stabs

lld/MachO/Arch/ARM64.cpp

@@ -44,6 +44,7 @@ struct ARM64 : ARM64Common {
 
   void initICFSafeThunkBody(InputSection *thunk,
                             InputSection *branchTarget) const override;
+  InputSection *getThunkBranchTarget(InputSection *thunk) const override;
   uint32_t getICFSafeThunkSize() const override;
 };
 
@@ -197,6 +198,16 @@ void ARM64::initICFSafeThunkBody(InputSection *thunk,
                              /*referent=*/branchTarget);
 }
 
+InputSection *ARM64::getThunkBranchTarget(InputSection *thunk) const {
+  assert(thunk->relocs.size() == 1 &&
+         "expected a single reloc on ARM64 ICF thunk");
+  auto &reloc = thunk->relocs[0];
+  assert(reloc.referent.is<InputSection *>() &&
+         "ARM64 thunk reloc is expected to point to an InputSection");
+
+  return reloc.referent.dyn_cast<InputSection *>();
+}
+
 uint32_t ARM64::getICFSafeThunkSize() const { return sizeof(icfSafeThunkCode); }
 
 ARM64::ARM64() : ARM64Common(LP64()) {

lld/MachO/ICF.cpp

@@ -481,6 +481,33 @@ void macho::markAddrSigSymbols() {
   }
 }
 
+// Given a symbol that was folded into a thunk, return the symbol pointing to
+// the actual body of the function. We use this approach rather than storing the
+// needed info in the Defined itself in order to minimize memory usage.
+Defined *macho::getBodyForThunkFoldedSym(Defined *foldedSym) {
+  assert(isa<ConcatInputSection>(foldedSym->originalIsec) &&
+         "thunk-folded ICF symbol expected to be on a ConcatInputSection");
+  // foldedSec is the InputSection that was marked as deleted upon fold
+  ConcatInputSection *foldedSec =
+      cast<ConcatInputSection>(foldedSym->originalIsec);
+
+  // thunkBody is the actual live thunk, containing the code that branches to
+  // the actual body of the function.
+  InputSection *thunkBody = foldedSec->replacement;
+
+  // The actual (merged) body of the function that the thunk jumps to. This will
+  // end up in the final binary.
+  InputSection *functionBody = target->getThunkBranchTarget(thunkBody);
+
+  for (Symbol *sym : functionBody->symbols) {
+    Defined *d = dyn_cast<Defined>(sym);
+    // The symbol needs to be at the start of the InputSection
+    if (d && d->value == 0)
+      return d;
+  }
+
+  llvm_unreachable("could not find body symbol for ICF-generated thunk");
+}
 void macho::foldIdenticalSections(bool onlyCfStrings) {
   TimeTraceScope timeScope("Fold Identical Code Sections");
   // The ICF equivalence-class segregation algorithm relies on pre-computed

lld/MachO/ICF.h

@@ -15,11 +15,17 @@
 
 namespace lld::macho {
 class Symbol;
+class Defined;
 
 void markAddrSigSymbols();
 void markSymAsAddrSig(Symbol *s);
 void foldIdenticalSections(bool onlyCfStrings);
 
+// Given a symbol that was folded into a thunk, return the symbol pointing to
+// the actual body of the function. We expose this function to allow getting the
+// main function body for a symbol that was folded via a thunk.
+Defined *getBodyForThunkFoldedSym(Defined *foldedSym);
+
 } // namespace lld::macho
 
 #endif

lld/MachO/SyntheticSections.cpp

@@ -10,6 +10,7 @@
 #include "ConcatOutputSection.h"
 #include "Config.h"
 #include "ExportTrie.h"
+#include "ICF.h"
 #include "InputFiles.h"
 #include "MachOStructs.h"
 #include "ObjC.h"
@@ -1204,6 +1205,18 @@ void SymtabSection::emitEndFunStab(Defined *defined) {
   stabs.emplace_back(std::move(stab));
 }
 
+// Given a pointer to a function symbol, return the symbol that points to the
+// actual function body that will go in the final binary. Generally this is the
+// symbol itself, but if the symbol was folded using a thunk, we retrieve the
+// target function body from the thunk.
+Defined *SymtabSection::getFuncBodySym(Defined *originalSym) {
+  if (originalSym->identicalCodeFoldingKind == Symbol::ICFFoldKind::None ||
+      originalSym->identicalCodeFoldingKind == Symbol::ICFFoldKind::Body)
+    return originalSym;
+
+  return macho::getBodyForThunkFoldedSym(originalSym);
+}
+
 void SymtabSection::emitStabs() {
   if (config->omitDebugInfo)
     return;
@@ -1229,20 +1242,10 @@ void SymtabSection::emitStabs() {
       if (defined->isAbsolute())
         continue;
 
-      // Never generate a STABS entry for a symbol that has been ICF'ed using a
-      // thunk, just as we do for fully ICF'ed functions. Otherwise, we end up
-      // generating invalid DWARF as dsymutil will assume the entire function
-      // body is at that location, when, in reality, only the thunk is
-      // present. This will end up causing overlapping DWARF entries.
-      // TODO: Find an implementation that works in combination with
-      // `--keep-icf-stabs`.
-      if (defined->identicalCodeFoldingKind == Symbol::ICFFoldKind::Thunk)
-        continue;
-
       // Constant-folded symbols go in the executable's symbol table, but don't
-      // get a stabs entry unless --keep-icf-stabs flag is specified
+      // get a stabs entry unless --keep-icf-stabs flag is specified.
       if (!config->keepICFStabs &&
-          defined->identicalCodeFoldingKind == Symbol::ICFFoldKind::Body)
+          defined->identicalCodeFoldingKind != Symbol::ICFFoldKind::None)
         continue;
 
       ObjFile *file = defined->getObjectFile();
@@ -1251,8 +1254,8 @@ void SymtabSection::emitStabs() {
 
       // We use 'originalIsec' to get the file id of the symbol since 'isec()'
       // might point to the merged ICF symbol's file
-      symbolsNeedingStabs.emplace_back(defined,
-                                       defined->originalIsec->getFile()->id);
+      symbolsNeedingStabs.emplace_back(
+          defined, getFuncBodySym(defined)->originalIsec->getFile()->id);
     }
   }
 
@@ -1269,7 +1272,8 @@ void SymtabSection::emitStabs() {
     Defined *defined = pair.first;
     // We use 'originalIsec' of the symbol since we care about the actual origin
     // of the symbol, not the canonical location returned by `isec()`.
-    InputSection *isec = defined->originalIsec;
+    Defined *funcBodySym = getFuncBodySym(defined);
+    InputSection *isec = funcBodySym->originalIsec;
     ObjFile *file = cast<ObjFile>(isec->getFile());
 
     if (lastFile == nullptr || lastFile != file) {
@@ -1284,12 +1288,12 @@ void SymtabSection::emitStabs() {
     StabsEntry symStab;
     symStab.sect = isec->parent->index;
     symStab.strx = stringTableSection.addString(defined->getName());
-    symStab.value = defined->getVA();
+    symStab.value = funcBodySym->getVA();
 
     if (isCodeSection(isec)) {
       symStab.type = N_FUN;
       stabs.emplace_back(std::move(symStab));
-      emitEndFunStab(defined);
+      emitEndFunStab(funcBodySym);
     } else {
       symStab.type = defined->isExternal() ? N_GSYM : N_STSYM;
       stabs.emplace_back(std::move(symStab));

lld/MachO/SyntheticSections.h

@@ -485,6 +485,7 @@ class SymtabSection : public LinkEditSection {
   void emitEndSourceStab();
   void emitObjectFileStab(ObjFile *);
   void emitEndFunStab(Defined *);
+  Defined *getFuncBodySym(Defined *);
   void emitStabs();
 
 protected:

lld/MachO/Target.h

@@ -80,6 +80,12 @@ class TargetInfo {
     llvm_unreachable("target does not support ICF safe thunks");
   }
 
+  // Given a thunk for which `initICFSafeThunkBody` was called, return the
+  // branchTarget it was initialized with.
+  virtual InputSection *getThunkBranchTarget(InputSection *thunk) const {
+    llvm_unreachable("target does not support ICF safe thunks");
+  }
+
   virtual uint32_t getICFSafeThunkSize() const {
     llvm_unreachable("target does not support ICF safe thunks");
   }

lld/test/MachO/icf-safe-thunks-dwarf.ll

@@ -20,6 +20,59 @@
 ; VERIFY-STABS:  N_FUN{{.*}}_func_A
 ; VERIFY-STABS:  N_FUN{{.*}}_take_func_addr
 
+;;;;;;;;;;;;;;;;;;;;;;;;;;;; Check safe_thunks ICF + keeping STABS entries ;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+;;; Check scenario with where we do safe_thunks ICF and also generate STABS entries
+; RUN: %lld -arch arm64 -lSystem --icf=safe_thunks --keep-icf-stabs -dylib -o %t/a_thunks.dylib %t/a.o
+; RUN: dsymutil -s %t/a_thunks.dylib > %t/a_thunks.txt
+
+
+; RUN: dsymutil --flat --verify-dwarf=none %t/a_thunks.dylib -o %t/a_thunks.dSYM
+; RUN: dsymutil -s %t/a_thunks.dSYM >> %t/a_thunks.txt
+; RUN: llvm-dwarfdump -a %t/a_thunks.dSYM >> %t/a_thunks.txt
+
+; RUN: cat %t/a_thunks.txt | FileCheck %s --check-prefix=VERIFY-THUNK
+
+# VERIFY-THUNK-LABEL: Symbol table for: '{{.*}}/a_thunks.dylib'
+# Capture the 'n_value's for N_FUN entries of _func_A, _func_B, and _func_C
+# VERIFY-THUNK:  [[MERGED_FUN_ADDR:[0-9a-f]+]] '_func_A'
+# VERIFY-THUNK:  [[MERGED_FUN_ADDR]] '_func_B'
+# VERIFY-THUNK:  [[MERGED_FUN_ADDR]] '_func_C'
+
+# Capture the 'n_value's for SECT EXT entries in the first part
+# VERIFY-THUNK: SECT EXT{{.*}} [[SECT_EXT_A_NVALUE:[0-9a-f]+]] '_func_A'
+# VERIFY-THUNK: SECT EXT{{.*}} [[SECT_EXT_B_NVALUE:[0-9a-f]+]] '_func_B'
+# VERIFY-THUNK: SECT EXT{{.*}} [[SECT_EXT_C_NVALUE:[0-9a-f]+]] '_func_C'
+
+# VERIFY-THUNK: ----------------------------------------------------------------------
+# VERIFY-THUNK-LABEL: Symbol table for: '{{.*}}/a_thunks.dSYM'
+
+# Verify that the SECT EXT 'n_value's in the second part match the first part
+# VERIFY-THUNK: SECT EXT{{.*}} [[SECT_EXT_A_NVALUE]] '_func_A'
+# VERIFY-THUNK: SECT EXT{{.*}} [[SECT_EXT_B_NVALUE]] '_func_B'
+# VERIFY-THUNK: SECT EXT{{.*}} [[SECT_EXT_C_NVALUE]] '_func_C'
+
+# Ensure that N_FUN 'n_value's match the DW_TAG_subprogram's DW_AT_low_pc
+# and that the DW_AT_name is at a specific relative position
+
+# VERIFY-THUNK-LABEL: .debug_info contents:
+# VERIFY-THUNK: Compile Unit: length = {{.*}}
+
+# Match the subprogram for func_A
+# VERIFY-THUNK: :   DW_TAG_subprogram
+# VERIFY-THUNK-NEXT: {{ +}}DW_AT_low_pc	(0x[[MERGED_FUN_ADDR]])
+# VERIFY-THUNK-NEXT-NEXT-NEXT-NEXT-NEXT: {{ +}}DW_AT_name	("func_A")
+
+# Match the subprogram for func_B
+# VERIFY-THUNK: :   DW_TAG_subprogram
+# VERIFY-THUNK-NEXT: {{ +}}DW_AT_low_pc	(0x[[MERGED_FUN_ADDR]])
+# VERIFY-THUNK-NEXT-NEXT-NEXT-NEXT-NEXT: {{ +}}DW_AT_name	("func_B")
+
+# Match the subprogram for func_C
+# VERIFY-THUNK: :   DW_TAG_subprogram
+# VERIFY-THUNK-NEXT: {{ +}}DW_AT_low_pc	(0x[[MERGED_FUN_ADDR]])
+# VERIFY-THUNK-NEXT-NEXT-NEXT-NEXT-NEXT: {{ +}}DW_AT_name	("func_C")
+
 ;--- a.cpp
 #define ATTR __attribute__((noinline)) extern "C"
 typedef unsigned long long ULL;