[BOLT] Keep folded functions in BinaryFunctions map. NFC (#180392)

In relocation mode, keep folded functions in the BinaryFunctions map
instead of erasing them. Mark them as folded using setFolded() and skip
emitting them.

bolt/include/bolt/Core/BinaryContext.h

@@ -234,9 +234,6 @@ class BinaryContext {
   /// Functions to be considered for the output in a sorted order.
   BinaryFunctionListType OutputFunctions;
 
-  /// A mutex that is used to control parallel accesses to BinaryFunctions.
-  mutable llvm::sys::RWMutex BinaryFunctionsMutex;
-
   /// Functions injected by BOLT.
   BinaryFunctionListType InjectedBinaryFunctions;
 

bolt/include/bolt/Core/BinaryFunction.h

@@ -771,6 +771,9 @@ private:
   /// disassembled state was later invalidated.
   void clearDisasmState();
 
+  /// Reset the function state into Empty state, i.e. pre-disassembly form.
+  void resetState();
+
   /// Release memory allocated for CFG and instructions.
   /// We still keep basic blocks for address translation/mapping purposes.
   void releaseCFG() {

bolt/lib/Core/BinaryContext.cpp

@@ -1556,39 +1556,26 @@ void BinaryContext::foldFunction(BinaryFunction &ChildBF,
   ChildBF.Aliases.clear();
 
   if (HasRelocations) {
-    // Merge execution counts of ChildBF into those of ParentBF.
-    // Without relocations, we cannot reliably merge profiles as both functions
-    // continue to exist and either one can be executed.
+    // Merge execution counts of ChildBF into those of ParentBF. We require
+    // relocations as without relocations we cannot reliably merge profiles as
+    // both functions continue to exist and either one can be executed.
     ChildBF.mergeProfileDataInto(ParentBF);
 
-    std::shared_lock<llvm::sys::RWMutex> ReadBfsLock(BinaryFunctionsMutex,
-                                                     std::defer_lock);
-    std::unique_lock<llvm::sys::RWMutex> WriteBfsLock(BinaryFunctionsMutex,
-                                                      std::defer_lock);
-    // Remove ChildBF from the global set of functions in relocs mode.
-    ReadBfsLock.lock();
-    auto FI = BinaryFunctions.find(ChildBF.getAddress());
-    ReadBfsLock.unlock();
-
-    assert(FI != BinaryFunctions.end() && "function not found");
-    assert(&ChildBF == &FI->second && "function mismatch");
-
-    WriteBfsLock.lock();
-    ChildBF.clearDisasmState();
-    FI = BinaryFunctions.erase(FI);
-    WriteBfsLock.unlock();
-
-  } else {
-    // In non-relocation mode we keep the function, but rename it.
-    std::string NewName = "__ICF_" + ChildName.str();
-
-    WriteCtxLock.lock();
-    ChildBF.getSymbols().push_back(Ctx->getOrCreateSymbol(NewName));
-    WriteCtxLock.unlock();
-
-    ChildBF.setFolded(&ParentBF);
+    // Clear CFG state to free memory, but keep function in map.
+    // The function is marked as folded and will not be emitted.
+    ChildBF.resetState();
   }
 
+  // Add a new symbol to the function. In relocation mode, this is a
+  // placeholder so that getSymbol() doesn't crash. In non-relocation mode,
+  // this effectively renames the function.
+  WriteCtxLock.lock();
+  ChildBF.getSymbols().push_back(
+      Ctx->getOrCreateSymbol("__ICF_" + ChildName.str()));
+  WriteCtxLock.unlock();
+
+  ChildBF.setFolded(&ParentBF);
+
   ParentBF.setHasFunctionsFoldedInto();
 }
 
@@ -1967,6 +1954,12 @@ bool BinaryContext::shouldEmit(const BinaryFunction &Function) const {
   if (Function.isPseudo())
     return false;
 
+  // In relocation mode, folded functions should not be emitted - their code
+  // is part of the parent. In non-relocation mode, folded functions are still
+  // emitted at their original location.
+  if (HasRelocations && Function.isFolded())
+    return false;
+
   if (opts::processAllFunctions())
     return true;
 

bolt/lib/Core/BinaryFunction.cpp

@@ -3256,6 +3256,30 @@ void BinaryFunction::clearDisasmState() {
   clearList(TakenBranches);
 }
 
+void BinaryFunction::resetState() {
+  clearDisasmState();
+
+  // Clear CFG state too.
+  if (hasCFG()) {
+    releaseCFG();
+
+    for (BinaryBasicBlock *BB : BasicBlocks)
+      delete BB;
+    clearList(BasicBlocks);
+
+    for (BinaryBasicBlock *BB : DeletedBasicBlocks)
+      delete BB;
+    clearList(DeletedBasicBlocks);
+
+    Layout.clear();
+  }
+
+  IsSimple = false;
+  IsIgnored = true;
+
+  CurrentState = State::Empty;
+}
+
 void BinaryFunction::setTrapOnEntry() {
   clearDisasmState();
 
@@ -3290,24 +3314,7 @@ void BinaryFunction::setIgnored() {
   if (CurrentState == State::Empty)
     return;
 
-  clearDisasmState();
-
-  // Clear CFG state too.
-  if (hasCFG()) {
-    releaseCFG();
-
-    for (BinaryBasicBlock *BB : BasicBlocks)
-      delete BB;
-    clearList(BasicBlocks);
-
-    for (BinaryBasicBlock *BB : DeletedBasicBlocks)
-      delete BB;
-    clearList(DeletedBasicBlocks);
-
-    Layout.clear();
-  }
-
-  CurrentState = State::Empty;
+  resetState();
 
   // Fix external references in the original function body.
   if (BC.HasRelocations) {

bolt/lib/Passes/PatchEntries.cpp

@@ -36,14 +36,20 @@ Error PatchEntries::runOnFunctions(BinaryContext &BC) {
   if (!opts::ForcePatch) {
     // Mark the binary for patching if we did not create external references
     // for original code in any of functions we are not going to emit.
-    bool NeedsPatching = llvm::any_of(
-        llvm::make_second_range(BC.getBinaryFunctions()),
-        [&](BinaryFunction &BF) {
-          return (!BC.shouldEmit(BF) && !BF.hasExternalRefRelocations()) ||
-                 BF.needsPatch();
-        });
+    auto needsPatching = [&](const BinaryFunction &BF) {
+      // FIXME: keep compatibility for NFC testing.
+      if (BF.isFolded())
+        return false;
 
-    if (!NeedsPatching)
+      // Patching is always needed if explicitly requested.
+      if (BF.needsPatch())
+        return true;
+
+      return !BC.shouldEmit(BF) && !BF.hasExternalRefRelocations();
+    };
+
+    if (!llvm::any_of(llvm::make_second_range(BC.getBinaryFunctions()),
+                      needsPatching))
       return Error::success();
   }
 

bolt/lib/Rewrite/RewriteInstance.cpp

@@ -5276,6 +5276,13 @@ void RewriteInstance::updateELFSymbolTable(
 
     const BinaryFunction *Function =
         BC->getBinaryFunctionAtAddress(Symbol.st_value);
+    // In relocation mode, if this is a folded function, use the parent function
+    // instead so that the symbol gets updated to the parent's output address.
+    // In non-relocation mode, folded functions are emitted at their original
+    // location, so we keep the original function reference.
+    // Follow the chain of folded functions to get the final parent.
+    while (BC->HasRelocations && Function && Function->isFolded())
+      Function = Function->getFoldedIntoFunction();
     // Ignore false function references, e.g. when the section address matches
     // the address of the function.
     if (Function && Symbol.getType() == ELF::STT_SECTION)
@@ -6076,6 +6083,11 @@ uint64_t RewriteInstance::getNewFunctionAddress(uint64_t OldAddress) {
   if (!Function)
     return 0;
 
+  // If this function was folded, its output address is 0 since it wasn't
+  // emitted. Follow the chain to get the parent function's address.
+  while (Function->isFolded())
+    Function = Function->getFoldedIntoFunction();
+
   return Function->getOutputAddress();
 }