[lld][elf] add safe-thunks in ELF

lld/ELF/Arch/AArch64.cpp

@@ -83,6 +83,8 @@ class AArch64 : public TargetInfo {
                 uint64_t val) const override;
   RelExpr adjustTlsExpr(RelType type, RelExpr expr) const override;
   void relocateAlloc(InputSectionBase &sec, uint8_t *buf) const override;
+  void initICFSafeThunkBody(InputSection *thunk, Symbol *target) const override;
+  uint32_t getICFSafeThunkSize() const override;
 
 private:
   void relaxTlsGdToLe(uint8_t *loc, const Relocation &rel, uint64_t val) const;
@@ -926,6 +928,18 @@ static bool needsGotForMemtag(const Relocation &rel) {
   return rel.sym->isTagged() && needsGot(rel.expr);
 }
 
+static constexpr uint8_t icfSafeThunkCode[] = {0x00, 0x00, 0x00, 0x14};
+
+void AArch64::initICFSafeThunkBody(InputSection *thunk, Symbol *target) const {
+  thunk->content_ = icfSafeThunkCode;
+  thunk->size = sizeof(icfSafeThunkCode);
+  thunk->relocations.push_back({R_PC, R_AARCH64_JUMP26, 0, 0, target});
+}
+
+uint32_t AArch64::getICFSafeThunkSize() const {
+  return sizeof(icfSafeThunkCode);
+}
+
 void AArch64::relocateAlloc(InputSectionBase &sec, uint8_t *buf) const {
   uint64_t secAddr = sec.getOutputSection()->addr;
   if (auto *s = dyn_cast<InputSection>(&sec))

lld/ELF/Config.h

@@ -97,7 +97,7 @@ enum class CGProfileSortKind { None, Hfsort, Cdsort };
 enum class DiscardPolicy { Default, All, Locals, None };
 
 // For --icf={none,safe,all}.
-enum class ICFLevel { None, Safe, All };
+enum class ICFLevel { None, Safe, SafeThunks, All };
 
 // For --strip-{all,debug}.
 enum class StripPolicy { None, All, Debug };

lld/ELF/Driver.cpp

@@ -808,11 +808,14 @@ static int getMemtagMode(Ctx &ctx, opt::InputArgList &args) {
 }
 
 static ICFLevel getICF(opt::InputArgList &args) {
-  auto *arg = args.getLastArg(OPT_icf_none, OPT_icf_safe, OPT_icf_all);
+  auto *arg = args.getLastArg(OPT_icf_none, OPT_icf_safe, OPT_icf_safe_thunks,
+                              OPT_icf_all);
   if (!arg || arg->getOption().getID() == OPT_icf_none)
     return ICFLevel::None;
   if (arg->getOption().getID() == OPT_icf_safe)
     return ICFLevel::Safe;
+  if (arg->getOption().getID() == OPT_icf_safe_thunks)
+    return ICFLevel::SafeThunks;
   return ICFLevel::All;
 }
 
@@ -2474,7 +2477,8 @@ static void findKeepUniqueSections(Ctx &ctx, opt::InputArgList &args) {
 
   // Symbols in the dynsym could be address-significant in other executables
   // or DSOs, so we conservatively mark them as address-significant.
-  bool icfSafe = ctx.arg.icf == ICFLevel::Safe;
+  bool icfSafe =
+      (ctx.arg.icf == ICFLevel::Safe || ctx.arg.icf == ICFLevel::SafeThunks);
   for (Symbol *sym : ctx.symtab->getSymbols())
     if (sym->isExported)
       markAddrsig(icfSafe, sym);

lld/ELF/ICF.cpp

@@ -77,9 +77,11 @@
 #include "InputFiles.h"
 #include "LinkerScript.h"
 #include "OutputSections.h"
+#include "Relocations.h"
 #include "SymbolTable.h"
 #include "Symbols.h"
 #include "SyntheticSections.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/Object/ELF.h"
 #include "llvm/Support/Parallel.h"
@@ -121,6 +123,8 @@ template <class ELFT> class ICF {
 
   void forEachClass(llvm::function_ref<void(size_t, size_t)> fn);
 
+  void applySafeThunksToRange(size_t begin, size_t end);
+
   Ctx &ctx;
   SmallVector<InputSection *, 0> sections;
 
@@ -160,10 +164,14 @@ template <class ELFT> class ICF {
 }
 
 // Returns true if section S is subject of ICF.
-static bool isEligible(InputSection *s) {
-  if (!s->isLive() || s->keepUnique || !(s->flags & SHF_ALLOC))
+static bool isEligible(InputSection *s, bool safeThunksMode) {
+  if (!s->isLive() || (s->keepUnique && !safeThunksMode) ||
+      !(s->flags & SHF_ALLOC))
     return false;
 
+  if (s->keepUnique)
+    return safeThunksMode && (s->flags & ELF::SHF_EXECINSTR);
+
   // Don't merge writable sections. .data.rel.ro sections are marked as writable
   // but are semantically read-only.
   if ((s->flags & SHF_WRITE) && s->name != ".data.rel.ro" &&
@@ -459,6 +467,58 @@ static void combineRelocHashes(unsigned cnt, InputSection *isec,
   isec->eqClass[(cnt + 1) % 2] = hash | (1U << 31);
 }
 
+// Given a range of identical icfInputs, replace address significant functions
+// with a thunk that is just a direct branch to the first function in the
+// series. This way we keep only one main body of the function but we still
+// retain the address uniqueness of relevant functions by having them be a
+// direct branch thunk rather than containing a full copy of the actual function
+// body.
+template <class ELFT>
+void ICF<ELFT>::applySafeThunksToRange(size_t begin, size_t end) {
+  InputSection *masterIsec = sections[begin];
+
+  uint32_t thunkSize = ctx.target->getICFSafeThunkSize();
+  // If the functions we're dealing with are smaller than the thunk size, then
+  // just leave them all as-is - creating thunks would be a net loss.
+  if (masterIsec->getSize() <= thunkSize)
+    return;
+
+  // Find the symbol to create the thunk for.
+  Symbol *masterSym = nullptr;
+  for (Symbol *sym : masterIsec->file->getSymbols()) {
+    if (auto *d = dyn_cast<Defined>(sym)) {
+      if (d->section == masterIsec) {
+        masterSym = sym;
+        break;
+      }
+    }
+  }
+
+  if (!masterSym)
+    return;
+
+  for (size_t i = begin + 1; i < end; ++i) {
+    InputSection *isec = sections[i];
+    if (!isec->keepUnique)
+      break;
+
+    auto *thunk = make<InputSection>(*isec);
+    ctx.target->initICFSafeThunkBody(thunk, masterSym);
+    thunk->markLive();
+    auto *osec = isec->getParent();
+    auto *isd = cast<InputSectionDescription>(osec->commands.back());
+    isd->sections.push_back(thunk);
+    osec->commitSection(thunk);
+    isec->repl = thunk;
+    isec->markDead();
+
+    for (Symbol *sym : thunk->file->getSymbols())
+      if (auto *d = dyn_cast<Defined>(sym))
+        if (d->section == isec)
+          d->size = thunkSize;
+  }
+}
+
 // The main function of ICF.
 template <class ELFT> void ICF<ELFT>::run() {
   // Two text sections may have identical content and relocations but different
@@ -475,10 +535,11 @@ template <class ELFT> void ICF<ELFT>::run() {
         [&](InputSection &s) { s.eqClass[0] = s.eqClass[1] = ++uniqueId; });
 
   // Collect sections to merge.
+  bool safeThunksMode = ctx.arg.icf == ICFLevel::SafeThunks;
   for (InputSectionBase *sec : ctx.inputSections) {
     auto *s = dyn_cast<InputSection>(sec);
     if (s && s->eqClass[0] == 0) {
-      if (isEligible(s))
+      if (isEligible(s, safeThunksMode))
         sections.push_back(s);
       else
         // Ineligible sections are assigned unique IDs, i.e. each section
@@ -510,9 +571,13 @@ template <class ELFT> void ICF<ELFT>::run() {
 
   // From now on, sections in Sections vector are ordered so that sections
   // in the same equivalence class are consecutive in the vector.
-  llvm::stable_sort(sections, [](const InputSection *a, const InputSection *b) {
-    return a->eqClass[0] < b->eqClass[0];
-  });
+  llvm::stable_sort(
+      sections, [safeThunksMode](const InputSection *a, const InputSection *b) {
+        if (safeThunksMode)
+          if (a->eqClass[0] == b->eqClass[0])
+            return a->keepUnique > b->keepUnique;
+        return a->eqClass[0] < b->eqClass[0];
+      });
 
   // Compare static contents and assign unique equivalence class IDs for each
   // static content. Use a base offset for these IDs to ensure no overlap with
@@ -535,12 +600,19 @@ template <class ELFT> void ICF<ELFT>::run() {
   auto print = [&ctx = ctx]() -> ELFSyncStream {
     return {ctx, ctx.arg.printIcfSections ? DiagLevel::Msg : DiagLevel::None};
   };
+  if (safeThunksMode)
+    forEachClassRange(0, sections.size(), [&](size_t begin, size_t end) {
+      applySafeThunksToRange(begin, end);
+    });
+
   // Merge sections by the equivalence class.
   forEachClassRange(0, sections.size(), [&](size_t begin, size_t end) {
     if (end - begin == 1)
       return;
     print() << "selected section " << sections[begin];
     for (size_t i = begin + 1; i < end; ++i) {
+      if (safeThunksMode && sections[i]->keepUnique)
+        continue;
       print() << "  removing identical section " << sections[i];
       sections[begin]->replace(sections[i]);
 

lld/ELF/ICF.h

@@ -9,8 +9,10 @@
 #ifndef LLD_ELF_ICF_H
 #define LLD_ELF_ICF_H
 
+#include "Target.h"
 namespace lld::elf {
 struct Ctx;
+class TargetInfo;
 
 template <class ELFT> void doIcf(Ctx &);
 }

lld/ELF/Options.td

@@ -293,6 +293,8 @@ def icf_all: F<"icf=all">, HelpText<"Enable identical code folding">;
 
 def icf_safe: F<"icf=safe">, HelpText<"Enable safe identical code folding">;
 
+def icf_safe_thunks: F<"icf=safe_thunks">, HelpText<"Enable safe identical code folding with thunks">;
+
 def icf_none: F<"icf=none">, HelpText<"Disable identical code folding (default)">;
 
 def ignore_function_address_equality: FF<"ignore-function-address-equality">,

lld/ELF/Target.h

@@ -16,6 +16,7 @@
 #include "llvm/Object/ELF.h"
 #include "llvm/Object/ELFTypes.h"
 #include "llvm/Support/Compiler.h"
+#include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MathExtras.h"
 #include <array>
 
@@ -102,6 +103,13 @@ class TargetInfo {
   virtual void applyJumpInstrMod(uint8_t *loc, JumpModType type,
                                  JumpModType val) const {}
 
+  // Initialize the body of the safe thunk in ICF for the target.
+  virtual void initICFSafeThunkBody(InputSection *thunk, Symbol *target) const {
+    llvm_unreachable("target does not support ICF safe thunks");
+  }
+  // Returns the size of the safe thunk in ICF for the target.
+  virtual uint32_t getICFSafeThunkSize() const { return 0; }
+
   virtual ~TargetInfo();
 
   // This deletes a jump insn at the end of the section if it is a fall thru to