[BOLT][NFC] Track fragment relationships using EquivalenceClasses

Three-way splitting can create references between split fragments (warm
to cold or vice versa) that are not handled by
`isChildOf/isParentOf/isChildOrParentOf`. Generalize fragment
relationships to allow checking if two functions belong to one group,
potentially in presence of ICF which can join multiple groups.

Test Plan: NFC for existing tests

Reviewers: maksfb, ayermolo, rafaelauler, dcci

Reviewed By: rafaelauler

Pull Request: https://github.com/llvm/llvm-project/pull/99979

bolt/include/bolt/Core/BinaryContext.h

@@ -23,6 +23,7 @@
 #include "bolt/RuntimeLibs/RuntimeLibrary.h"
 #include "llvm/ADT/AddressRanges.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/EquivalenceClasses.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/iterator.h"
 #include "llvm/BinaryFormat/Dwarf.h"
@@ -241,6 +242,10 @@ class BinaryContext {
   /// Function fragments to skip.
   std::unordered_set<BinaryFunction *> FragmentsToSkip;
 
+  /// Fragment equivalence classes to query belonging to the same "family" in
+  /// presence of multiple fragments/multiple parents.
+  EquivalenceClasses<const BinaryFunction *> FragmentClasses;
+
   /// The runtime library.
   std::unique_ptr<RuntimeLibrary> RtLibrary;
 
@@ -1032,7 +1037,15 @@ public:
   ///   fragment_name == parent_name.cold(.\d+)?
   /// True if the Function is registered, false if the check failed.
   bool registerFragment(BinaryFunction &TargetFunction,
-                        BinaryFunction &Function) const;
+                        BinaryFunction &Function);
+
+  /// Return true if two functions belong to the same "family": are fragments
+  /// of one another, or fragments of the same parent, or transitively fragment-
+  /// related.
+  bool areRelatedFragments(const BinaryFunction *LHS,
+                           const BinaryFunction *RHS) const {
+    return FragmentClasses.isEquivalent(LHS, RHS);
+  }
 
   /// Add interprocedural reference for \p Function to \p Address
   void addInterproceduralReference(BinaryFunction *Function, uint64_t Address) {

bolt/include/bolt/Core/BinaryFunction.h

@@ -1793,11 +1793,6 @@ public:
     return ParentFragments.contains(&Other);
   }
 
-  /// Returns if this function is a parent of \p Other function.
-  bool isParentOf(const BinaryFunction &Other) const {
-    return Fragments.contains(&Other);
-  }
-
   /// Return the child fragment form parent function
   iterator_range<FragmentsSetTy::const_iterator> getFragments() const {
     return iterator_range<FragmentsSetTy::const_iterator>(Fragments.begin(),
@@ -1807,11 +1802,6 @@ public:
   /// Return the parent function for split function fragments.
   FragmentsSetTy *getParentFragments() { return &ParentFragments; }
 
-  /// Returns if this function is a parent or child of \p Other function.
-  bool isParentOrChildOf(const BinaryFunction &Other) const {
-    return isChildOf(Other) || isParentOf(Other);
-  }
-
   /// Set the profile data for the number of times the function was called.
   BinaryFunction &setExecutionCount(uint64_t Count) {
     ExecutionCount = Count;

bolt/lib/Core/BinaryContext.cpp

@@ -646,7 +646,7 @@ bool BinaryContext::analyzeJumpTable(const uint64_t Address,
     const BinaryFunction *TargetBF = getBinaryFunctionContainingAddress(Value);
     const bool DoesBelongToFunction =
         BF.containsAddress(Value) ||
-        (TargetBF && TargetBF->isParentOrChildOf(BF));
+        (TargetBF && areRelatedFragments(TargetBF, &BF));
     if (!DoesBelongToFunction) {
       LLVM_DEBUG({
         if (!BF.containsAddress(Value)) {
@@ -841,7 +841,7 @@ BinaryContext::getOrCreateJumpTable(BinaryFunction &Function, uint64_t Address,
     // Prevent associating a jump table to a specific fragment twice.
     // This simple check arises from the assumption: no more than 2 fragments.
     if (JT->Parents.size() == 1 && JT->Parents[0] != &Function) {
-      assert(JT->Parents[0]->isParentOrChildOf(Function) &&
+      assert(areRelatedFragments(JT->Parents[0], &Function) &&
              "cannot re-use jump table of a different function");
       // Duplicate the entry for the parent function for easy access
       JT->Parents.push_back(&Function);
@@ -1209,12 +1209,13 @@ void BinaryContext::generateSymbolHashes() {
 }
 
 bool BinaryContext::registerFragment(BinaryFunction &TargetFunction,
-                                     BinaryFunction &Function) const {
+                                     BinaryFunction &Function) {
   assert(TargetFunction.isFragment() && "TargetFunction must be a fragment");
   if (TargetFunction.isChildOf(Function))
     return true;
   TargetFunction.addParentFragment(Function);
   Function.addFragment(TargetFunction);
+  FragmentClasses.unionSets(&TargetFunction, &Function);
   if (!HasRelocations) {
     TargetFunction.setSimple(false);
     Function.setSimple(false);
@@ -1336,7 +1337,7 @@ void BinaryContext::processInterproceduralReferences() {
 
     if (TargetFunction) {
       if (TargetFunction->isFragment() &&
-          !TargetFunction->isChildOf(Function)) {
+          !areRelatedFragments(TargetFunction, &Function)) {
         this->errs()
             << "BOLT-WARNING: interprocedural reference between unrelated "
                "fragments: "

bolt/lib/Core/Exceptions.cpp

@@ -207,7 +207,7 @@ Error BinaryFunction::parseLSDA(ArrayRef<uint8_t> LSDASectionData,
                "BOLT-ERROR: cannot find landing pad fragment");
         BC.addInterproceduralReference(this, Fragment->getAddress());
         BC.processInterproceduralReferences();
-        assert(isParentOrChildOf(*Fragment) &&
+        assert(BC.areRelatedFragments(this, Fragment) &&
                "BOLT-ERROR: cannot have landing pads in different functions");
         setHasIndirectTargetToSplitFragment(true);
         BC.addFragmentsToSkip(this);