use equivalence classes to simplify and handle more cases

Author: pranavk

lld/ELF/ICF.cpp

@@ -80,6 +80,7 @@
 #include "SymbolTable.h"
 #include "Symbols.h"
 #include "SyntheticSections.h"
+#include "llvm/ADT/EquivalenceClasses.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/Object/ELF.h"
 #include "llvm/Support/Parallel.h"
@@ -558,7 +559,7 @@ template <class ELFT> void ICF<ELFT>::run() {
     return {ctx, ctx.arg.printIcfSections ? DiagLevel::Msg : DiagLevel::None};
   };
 
-  DenseMap<Symbol *, Symbol *> symbolMap;
+  EquivalenceClasses<Symbol *> symbolEquivalence;
   // Merge sections by the equivalence class.
   // Merge symbols identified as equivalent during ICF
   forEachClassRange(0, sections.size(), [&](size_t begin, size_t end) {
@@ -577,14 +578,7 @@ template <class ELFT> void ICF<ELFT>::run() {
         if (syms[i] == replacedSyms[i] || !syms[i]->isGlobal() ||
             !replacedSyms[i]->isGlobal())
           continue;
-        auto [it, inserted] =
-            symbolMap.insert(std::make_pair(replacedSyms[i], syms[i]));
-        print() << "  selected symbol: " << syms[i]->getName().data()
-                << "; replaced symbol: " << replacedSyms[i]->getName().data();
-        if (!inserted) {
-          print() << "    replacement already exists: "
-                  << it->getSecond()->getName().data();
-        }
+        symbolEquivalence.unionSets(syms[i], replacedSyms[i]);
       }
 
       // At this point we know sections merged are fully identical and hence
@@ -606,15 +600,23 @@ template <class ELFT> void ICF<ELFT>::run() {
   };
   for (Symbol *sym : ctx.symtab->getSymbols()) {
     fold(sym);
-    if (Symbol *s = symbolMap.lookup(sym))
-      ctx.symtab->redirect(sym, s);
+    auto it = symbolEquivalence.findLeader(sym);
+    if (it != symbolEquivalence.member_end() && *it != sym) {
+      print() << "Redirecting " << sym->getName() << " to " << (*it)->getName();
+      ctx.symtab->redirect(sym, *it);
+    }
   }
   parallelForEach(ctx.objectFiles, [&](ELFFileBase *file) {
     for (Symbol *sym : file->getLocalSymbols())
       fold(sym);
-    for (Symbol *&sym : file->getMutableGlobalSymbols())
-      if (Symbol *s = symbolMap.lookup(sym))
-        sym = s;
+    for (Symbol *&sym : file->getMutableGlobalSymbols()) {
+      auto it = symbolEquivalence.findLeader(sym);
+      if (it != symbolEquivalence.member_end() && *it != sym) {
+        print() << "Redirecting " << sym->getName() << " to "
+                << (*it)->getName();
+        sym = *it;
+      }
+    }
   });
 
   // InputSectionDescription::sections is populated by processSectionCommands().

lld/ELF/SymbolTable.cpp

@@ -32,7 +32,6 @@ using namespace lld::elf;
 void SymbolTable::redirect(Symbol *from, Symbol *to) {
   int &fromIdx = symMap[CachedHashStringRef(from->getName())];
   const int toIdx = symMap[CachedHashStringRef(to->getName())];
-
   fromIdx = toIdx;
 }
 

lld/test/ELF/aarch64-got-merging-icf.s

@@ -11,7 +11,7 @@
 # EXE: {{.*}}.got 00000010{{.*}}
 
 ## When symbols are preemptible in DSO mode, GOT entries wouldn't be merged
-# DSO: {{.*}}.got 00000020{{.*}}
+# DSO: {{.*}}.got 00000028{{.*}}
 
 .addrsig
 
@@ -50,11 +50,36 @@ bl f1_\index
 
 .endm
 
+# another set of sections merging: g1 <- g2
+
+.section .text.t1_0,"ax",@progbits
+t1_0:
+adrp x2, :got:g1
+mov x3, #1
+b t2_0
+
+.section .text.t2_0,"ax",@progbits
+t2_0:
+ldr x2, [x2, :got_lo12:g1]
+b callee
+
+.section .text.t1_1,"ax",@progbits
+t1_1:
+adrp x2, :got:g2
+mov x3, #2
+b t2_1
+
+.section .text.t2_1,"ax",@progbits
+t2_1:
+ldr x2, [x2, :got_lo12:g2]
+b callee
+
 .section .text._start,"ax",@progbits
 .globl _start
 _start:
 
 f 0 1
 f 1 1
 f 2 1
-f 3
+f 3 1
+f 4

lld/test/ELF/icf-preemptible.s

@@ -11,16 +11,15 @@
 # EXE-NOT:  {{.}}
 # EXE:      selected section {{.*}}:(.text.g1)
 # EXE-NEXT:   removing identical section {{.*}}:(.text.g2)
-# EXE-NEXT:   selected symbol: f1; replaced symbol: f2
 # EXE-NEXT:   removing identical section {{.*}}:(.text.g3)
 # EXE-NEXT: selected section {{.*}}:(.text.f1)
 # EXE-NEXT:   removing identical section {{.*}}:(.text.f2)
 # EXE-NEXT: selected section {{.*}}:(.text.h1)
 # EXE-NEXT:   removing identical section {{.*}}:(.text.h2)
-# EXE-NEXT:   selected symbol: g1; replaced symbol: g2
 # EXE-NEXT:   removing identical section {{.*}}:(.text.h3)
-# EXE-NEXT:   selected symbol: g1; replaced symbol: g3
-# EXE-NOT:  {{.}}
+# EXE-NEXT: Redirecting f2 to f1
+# EXE-NEXT: Redirecting g2 to g1
+# EXE-NEXT: Redirecting g3 to g1
 
 ## Definitions are preemptible in a DSO. Only leaf functions can be folded.
 # DSO-NOT:  {{.}}