[lld-macho] Fix compatibility between --icf=safe_thunks and --keep-icf-stabs (#116687)

Currently when `--icf=safe_thunks` is used, `STABS` entries cannot be
generated for ICF'ed functions. This is because if ICF converts a full
function into a thunk and then we generate a `STABS` entry for the
thunk, `dsymutil` will expect to find the entire function body at the
location of the thunk. Because just a thunk will be present at the
location of the `STABS` entry - dsymutil will generate invalid debug
info for such scenarios.

With this change, if `--icf=safe_thunks` is used and `--keep-icf-stabs`
is also specified, STABS entries will be created for all functions, even
merged ones. However, the STABS entries will point at the actual (full)
function body while having the name of the thunk. This way we still get
program correctness as well as correct DWARF data. When doing this, the
debug data will be identical to the scenario where we're using
`--icf=all` and `--keep-icf-stabs`, but the actual program will also
contain thunks, which won't show up in the DWARF data.

7 files changed, 125 insertions, 17 deletions

diff --git a/lld/MachO/Arch/ARM64.cpp b/lld/MachO/Arch/ARM64.cpp
index 195a8f0..882873a 100644
--- a/lld/MachO/Arch/ARM64.cpp
+++ b/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()) {
diff --git a/lld/MachO/ICF.cpp b/lld/MachO/ICF.cpp
index aedaecf..32dd44a 100644
--- a/lld/MachO/ICF.cpp
+++ b/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
diff --git a/lld/MachO/ICF.h b/lld/MachO/ICF.h
index 34ceb1c..e382fd6 100644
--- a/lld/MachO/ICF.h
+++ b/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
diff --git a/lld/MachO/SyntheticSections.cpp b/lld/MachO/SyntheticSections.cpp
index eee87b3..24844c2 100644
--- a/lld/MachO/SyntheticSections.cpp
+++ b/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));
diff --git a/lld/MachO/SyntheticSections.h b/lld/MachO/SyntheticSections.h
index a4c7f58..af99f22 100644
--- a/lld/MachO/SyntheticSections.h
+++ b/lld/MachO/SyntheticSections.h
@@ -485,6 +485,7 @@ private:
   void emitEndSourceStab();
   void emitObjectFileStab(ObjFile *);
   void emitEndFunStab(Defined *);
+  Defined *getFuncBodySym(Defined *);
   void emitStabs();
 
 protected:
diff --git a/lld/MachO/Target.h b/lld/MachO/Target.h
index eaa0336..b5b80e0 100644
--- a/lld/MachO/Target.h
+++ b/lld/MachO/Target.h
@@ -80,6 +80,12 @@ public:
     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");
   }
diff --git a/lld/test/MachO/icf-safe-thunks-dwarf.ll b/lld/test/MachO/icf-safe-thunks-dwarf.ll
index 1e4422a..9edad60 100644
--- a/lld/test/MachO/icf-safe-thunks-dwarf.ll
+++ b/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;