Refactor

rui314 · Jul 28, 2024 · 3f7236d · 3f7236d
1 parent 766ce80
commit 3f7236d
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 58 deletions.
diff --git a/elf/arch-x86-64.cc b/elf/arch-x86-64.cc
@@ -28,6 +28,8 @@
 
 #include "mold.h"
 
+#include <tbb/parallel_for_each.h>
+
 namespace mold::elf {
 
 using E = X86_64;
@@ -815,4 +817,58 @@ void InputSection<E>::scan_relocations(Context<E> &ctx) {
   }
 }
 
+// Rewrite the leading endbr64 instruction with a nop if a function
+// symbol's address was not taken.
+void rewrite_endbr(Context<E> &ctx) {
+  Timer t(ctx, "rewrite_endbr");
+
+  auto mark = [&](Symbol<E> *sym) {
+    if (sym) {
+      std::scoped_lock lock(sym->mu);
+      sym->address_taken = true;
+    }
+  };
+
+  // Compute address-taken bit for each symbol
+  tbb::parallel_for_each(ctx.objs, [&](ObjectFile<E> *file) {
+    for (std::unique_ptr<InputSection<E>> &isec : file->sections)
+      if (isec && isec->is_alive && (isec->shdr().sh_flags & SHF_ALLOC))
+        for (const ElfRel<E> &rel : isec->get_rels(ctx))
+          if (!is_func_call_rel(rel))
+            if (Symbol<E> *sym = file->symbols[rel.r_sym];
+                sym->esym().st_type == STT_FUNC)
+              mark(sym);
+  });
+
+  // Exported symbols are conservatively assumed to be address-taken.
+  if (ctx.dynsym)
+    for (Symbol<E> *sym : ctx.dynsym->symbols)
+      if (sym && sym->is_exported)
+        mark(sym);
+
+  // Some symbols are implicitly address-taken
+  mark(ctx.arg.entry);
+  mark(ctx.arg.init);
+  mark(ctx.arg.fini);
+
+  constexpr u8 endbr64[] = {0xf3, 0x0f, 0x1e, 0xfa};
+  constexpr u8 nop[] = {0x0f, 0x1f, 0x40, 0x00};
+
+  // Rewrite endbr64 with nop
+  tbb::parallel_for_each(ctx.objs, [&](ObjectFile<E> *file) {
+    for (Symbol<E> *sym : file->symbols) {
+      if (sym->file == file && sym->esym().st_type == STT_FUNC &&
+          !sym->address_taken) {
+        if (InputSection<E> *isec = sym->get_input_section()) {
+          if (OutputSection<E> *osec = isec->output_section) {
+            u8 *buf = ctx.buf + osec->shdr.sh_offset + isec->offset + sym->value;
+            if (memcmp(buf, endbr64, 4) == 0)
+              memcpy(buf, nop, 4);
+          }
+        }
+      }
+    }
+  });
+}
+
 } // namespace mold::elf
diff --git a/elf/main.cc b/elf/main.cc
@@ -648,8 +648,9 @@ int elf_main(int argc, char **argv) {
   // Copy input sections to the output file and apply relocations.
   copy_chunks(ctx);
 
-  if (ctx.arg.z_rewrite_endbr)
-    rewrite_endbr(ctx);
+  if constexpr (is_x86_64<E>)
+    if (ctx.arg.z_rewrite_endbr)
+      rewrite_endbr(ctx);
 
   // Dynamic linker works better with sorted .rela.dyn section,
   // so we sort them.

diff --git a/elf/mold.h b/elf/mold.h
@@ -1505,7 +1505,6 @@ template <typename E> void create_output_symtab(Context<E> &);
 template <typename E> void report_undef_errors(Context<E> &);
 template <typename E> void create_reloc_sections(Context<E> &);
 template <typename E> void copy_chunks(Context<E> &);
-template <typename E> void rewrite_endbr(Context<E> &);
 template <typename E> void apply_version_script(Context<E> &);
 template <typename E> void parse_symbol_version(Context<E> &);
 template <typename E> void compute_import_export(Context<E> &);
@@ -1521,6 +1520,12 @@ template <typename E> void write_separate_debug_file(Context<E> &ctx);
 template <typename E> void write_dependency_file(Context<E> &);
 template <typename E> void show_stats(Context<E> &);
 
+//
+// arch-x86-64.cc
+//
+
+void rewrite_endbr(Context<X86_64> &ctx);
+
 //
 // arch-arm32.cc
 //

diff --git a/elf/passes.cc b/elf/passes.cc
@@ -1691,60 +1691,6 @@ void copy_chunks(Context<E> &ctx) {
   zero(chunks.back(), ctx.output_file->filesize);
 }
 
-// Rewrite the leading endbr64 instruction with a nop if a function
-// symbol's address was not taken.
-template <typename E>
-void rewrite_endbr(Context<E> &ctx) {
-  Timer t(ctx, "rewrite_endbr");
-  assert(is_x86_64<E>);
-
-  // Compute address-taken bit for each symbol
-  tbb::parallel_for_each(ctx.objs, [&](ObjectFile<E> *file) {
-    for (std::unique_ptr<InputSection<E>> &isec : file->sections) {
-      if (isec && isec->is_alive && (isec->shdr().sh_flags & SHF_ALLOC)) {
-        for (const ElfRel<E> &rel : isec->get_rels(ctx)) {
-          Symbol<E> &sym = *file->symbols[rel.r_sym];
-          if (!is_func_call_rel(rel) && sym.esym().st_type == STT_FUNC) {
-            std::scoped_lock lock(sym.mu);
-            sym.address_taken = true;
-          }
-        }
-      }
-    }
-  });
-
-  // Exported symbols are conservatively assumed to be address-taken.
-  if (ctx.dynsym)
-    for (Symbol<E> *sym : ctx.dynsym->symbols)
-      if (sym && sym->is_exported)
-        sym->address_taken = true;
-
-  // Some symbols are implicitly address-taken
-  ctx.arg.entry->address_taken = true;
-  ctx.arg.init->address_taken = true;
-  ctx.arg.fini->address_taken = true;
-
-  // Rewrite endbr64 with nop
-  u8 endbr64[] = {0xf3, 0x0f, 0x1e, 0xfa};
-  u8 nop[] = {0x0f, 0x1f, 0x40, 0x00};
-
-  tbb::parallel_for_each(ctx.objs, [&](ObjectFile<E> *file) {
-    for (Symbol<E> *sym : file->symbols) {
-      if (sym->file == file && sym->esym().st_type == STT_FUNC &&
-          !sym->address_taken) {
-        if (InputSection<E> *isec = sym->get_input_section()) {
-          if (OutputSection<E> *osec = isec->output_section) {
-            u8 *buf = ctx.buf + osec->shdr.sh_offset + isec->offset +
-                      sym->value;
-            if (memcmp(buf, endbr64, 4) == 0)
-              memcpy(buf, nop, 4);
-          }
-        }
-      }
-    }
-  });
-}
-
 template <typename E>
 void construct_relr(Context<E> &ctx) {
   Timer t(ctx, "construct_relr");
@@ -3328,7 +3274,6 @@ template void scan_relocations(Context<E> &);
 template void report_undef_errors(Context<E> &);
 template void create_reloc_sections(Context<E> &);
 template void copy_chunks(Context<E> &);
-template void rewrite_endbr(Context<E> &);
 template void construct_relr(Context<E> &);
 template void sort_dynsyms(Context<E> &);
 template void create_output_symtab(Context<E> &);