[lld] [lld-macho] Include ICF safe thunks in balanced partitioning (PR #212096)

[Ellis Hoag via llvm-commits]

================
@@ -119,36 +119,76 @@ DenseMap<const InputSection *, int> lld::macho::runBalancedPartitioning(
     bool compressionSortStartupFunctions, bool verbose) {
   // Collect candidate sections and associated symbols.
   SmallVector<InputSection *> sections;
+  DenseMap<const InputSection *, unsigned> sectionToIdx;
   DenseMap<CachedHashStringRef, std::set<unsigned>> rootSymbolToSectionIdxs;
+  auto addSectionForName = [&](StringRef name, unsigned idx) {
+    auto rootName = lld::utils::getRootSymbol(name);
+    rootSymbolToSectionIdxs[CachedHashStringRef(rootName)].insert(idx);
+    if (auto linkageName = BPOrdererMachO::getResolvedLinkageName(rootName))
+      rootSymbolToSectionIdxs[CachedHashStringRef(*linkageName)].insert(idx);
+  };
+  auto addSection = [&](InputSection *isec) {
+    if (!isec || isec->data.empty() || !isec->data.data())
+      return;
+    // CString section order is handled by
+    // {Deduplicated}CStringSection::finalizeContents()
+    if (isa<CStringInputSection>(isec) || isec->isFinal)
+      return;
+    // ConcatInputSections are entirely live or dead, so the offset is
+    // irrelevant.
+    if (isa<ConcatInputSection>(isec) && !isec->isLive(0))
+      return;
+    unsigned idx = sections.size();
+    sections.emplace_back(isec);
+    sectionToIdx.try_emplace(isec, idx);
+    for (auto *sym : isec->symbols)
+      addSectionForName(sym->getName(), idx);
+  };
   for (const auto *file : inputFiles) {
     for (auto *sec : file->sections) {
       if (sec->name == section_names::ehFrame &&
           sec->segname == segment_names::text)
         continue;
-      for (auto &subsec : sec->subsections) {
-        auto *isec = subsec.isec;
-        if (!isec || isec->data.empty() || !isec->data.data())
-          continue;
-        // CString section order is handled by
-        // {Deduplicated}CStringSection::finalizeContents()
-        if (isa<CStringInputSection>(isec) || isec->isFinal)
-          continue;
-        // ConcatInputSections are entirely live or dead, so the offset is
-        // irrelevant.
-        if (isa<ConcatInputSection>(isec) && !isec->isLive(0))
-          continue;
-        size_t idx = sections.size();
-        sections.emplace_back(isec);
-        for (auto *sym : BPOrdererMachO::getSymbols(*isec)) {
-          auto rootName = lld::utils::getRootSymbol(sym->getName());
-          rootSymbolToSectionIdxs[CachedHashStringRef(rootName)].insert(idx);
-          if (auto linkageName =
-                  BPOrdererMachO::getResolvedLinkageName(rootName))
-            rootSymbolToSectionIdxs[CachedHashStringRef(*linkageName)].insert(
-                idx);
+      for (auto &subsec : sec->subsections)
+        addSection(subsec.isec);
+    }
+  }
+  // ICF safe thunks are linker-created after the input-file section graph is
+  // built, so they do not appear in file->sections. Include them through the
+  // same path as input-file sections so only live BP candidates are added.
+  for (auto *isec : inputSections) {
+    for (auto *sym : isec->symbols) {
+      if (sym->identicalCodeFoldingKind == Symbol::ICFFoldKind::Thunk) {
+        addSection(isec);
+        break;
+      }
+    }
+  }
+
+  // A temporal profile naming an ICF thunk describes execution of both the
+  // thunk and the shared body it branches to. Add the body to each name that
+  // resolves to a thunk, following the thunk relocation to find it.
+  for (auto &entry : rootSymbolToSectionIdxs) {
----------------
ellishg wrote:

Does this compile? I remember we sometimes need to use `const auto`, but that might not be possible since we are appending to `sectionIdxs`.
```suggestion
  for (auto &[symbol, sectionIdxs] : rootSymbolToSectionIdxs) {
```

https://github.com/llvm/llvm-project/pull/212096