From e3876ddbd47e14494e34f9cc9c42eef2f97e9eb2 Mon Sep 17 00:00:00 2001
From: Rui Ueyama <ruiu@cs.stanford.edu>
Date: Sun, 21 Jul 2024 21:58:29 +0900
Subject: [PATCH] Refactor

---
 elf/mold.h           |  5 ++---
 elf/output-chunks.cc | 11 ++++++++---
 elf/tls.cc           | 47 ++++++++++++--------------------------------
 3 files changed, 23 insertions(+), 40 deletions(-)

diff --git a/elf/mold.h b/elf/mold.h
index 7bb5f99e8f..fe6295fbcf 100644
--- a/elf/mold.h
+++ b/elf/mold.h
@@ -335,9 +335,8 @@ class __attribute__((aligned(4))) InputSection {
 // tls.cc
 //
 
-template <typename E> u64 get_tls_begin(Context<E> &);
-template <typename E> u64 get_tp_addr(Context<E> &);
-template <typename E> u64 get_dtp_addr(Context<E> &);
+template <typename E> u64 get_tp_addr(const ElfPhdr<E> &);
+template <typename E> u64 get_dtp_addr(const ElfPhdr<E> &);
 
 //
 // output-chunks.cc
diff --git a/elf/output-chunks.cc b/elf/output-chunks.cc
index 1e997a6e0d..74c054b8f4 100644
--- a/elf/output-chunks.cc
+++ b/elf/output-chunks.cc
@@ -361,9 +361,14 @@ void OutputPhdr<E>::update_shdr(Context<E> &ctx) {
   phdrs = create_phdr(ctx);
   this->shdr.sh_size = phdrs.size() * sizeof(ElfPhdr<E>);
 
-  ctx.tls_begin = get_tls_begin(ctx);
-  ctx.tp_addr = get_tp_addr(ctx);
-  ctx.dtp_addr = get_dtp_addr(ctx);
+  for (ElfPhdr<E> &phdr : phdrs) {
+    if (phdr.p_type == PT_TLS) {
+      ctx.tls_begin = phdr.p_vaddr;
+      ctx.tp_addr = get_tp_addr(phdr);
+      ctx.dtp_addr = get_dtp_addr(phdr);
+      break;
+    }
+  }
 }
 
 template <typename E>
diff --git a/elf/tls.cc b/elf/tls.cc
index 8d391ace96..fc4f955fad 100644
--- a/elf/tls.cc
+++ b/elf/tls.cc
@@ -124,42 +124,24 @@
 
 namespace mold::elf {
 
-template <typename E>
-static ElfPhdr<E> *get_tls_segment(Context<E> &ctx) {
-  if (ctx.phdr)
-    for (ElfPhdr<E> &phdr : ctx.phdr->phdrs)
-      if (phdr.p_type == PT_TLS)
-        return &phdr;
-  return nullptr;
-}
-
-template <typename E>
-u64 get_tls_begin(Context<E> &ctx) {
-  if (ElfPhdr<E> *phdr = get_tls_segment(ctx))
-    return phdr->p_vaddr;
-  return 0;
-}
-
 // Returns the TP address which can be used for efficient TLV accesses in
 // the main executable. TP at runtime refers to a per-process TLS block
 // whose address is not known at link-time. So the address returned from
 // this function is the TP if the TLS template image were a TLS block.
 template <typename E>
-u64 get_tp_addr(Context<E> &ctx) {
-  ElfPhdr<E> *phdr = get_tls_segment(ctx);
-  if (!phdr)
-    return 0;
+u64 get_tp_addr(const ElfPhdr<E> &phdr) {
+  assert(phdr.p_type == PT_TLS);
 
   if constexpr (is_x86<E> || is_sparc<E> || is_s390x<E>) {
     // On x86, SPARC and s390x, TP (%gs on i386, %fs on x86-64, %g7 on SPARC
     // and %a0/%a1 on s390x) refers to past the end of the TLS block for
     // historical reasons. TLVs are accessed with negative offsets from TP.
-    return align_to(phdr->p_vaddr + phdr->p_memsz, phdr->p_align);
+    return align_to(phdr.p_vaddr + phdr.p_memsz, phdr.p_align);
   } else if constexpr (is_arm<E> || is_sh4<E> || is_alpha<E>) {
     // On ARM, SH4 and Alpha, the runtime appends two words at the beginning
     // of TLV template image when copying TLVs to the TLS block, so we need
     // to offset it.
-    return align_down(phdr->p_vaddr - sizeof(Word<E>) * 2, phdr->p_align);
+    return align_down(phdr.p_vaddr - sizeof(Word<E>) * 2, phdr.p_align);
   } else if constexpr (is_ppc<E> || is_m68k<E>) {
     // On PowerPC and m68k, TP is 0x7000 (28 KiB) past the beginning
     // of the TLV block to maximize the addressable range of load/store
@@ -167,24 +149,22 @@ u64 get_tp_addr(Context<E> &ctx) {
     // (32 KiB) off because there's a small implementation-defined piece of
     // data before the initial TLV block, and the runtime wants to access
     // them efficiently too.
-    return phdr->p_vaddr + 0x7000;
+    return phdr.p_vaddr + 0x7000;
   } else {
     // RISC-V and LoongArch just uses the beginning of the main executable's
     // TLV block as TP. Their load/store instructions usually take 12-bits
     // signed immediates, so the beginning of the TLS block ± 2 KiB is
     // accessible with a single load/store instruction.
     static_assert(is_riscv<E> || is_loongarch<E>);
-    return phdr->p_vaddr;
+    return phdr.p_vaddr;
   }
 }
 
 // Returns the address __tls_get_addr() would return if it's called
 // with offset 0.
 template <typename E>
-u64 get_dtp_addr(Context<E> &ctx) {
-  ElfPhdr<E> *phdr = get_tls_segment(ctx);
-  if (!phdr)
-    return 0;
+u64 get_dtp_addr(const ElfPhdr<E> &phdr) {
+  assert(phdr.p_type == PT_TLS);
 
   if constexpr (is_ppc<E> || is_m68k<E>) {
     // On PowerPC and m68k, R_DTPOFF is resolved to the address 0x8000
@@ -193,21 +173,20 @@ u64 get_dtp_addr(Context<E> &ctx) {
     // immediates. That is, if the offset were right at the beginning of the
     // start of the TLS block, the half of addressible space (negative
     // immediates) would have been wasted.
-    return phdr->p_vaddr + 0x8000;
+    return phdr.p_vaddr + 0x8000;
   } else if constexpr (is_riscv<E>) {
     // On RISC-V, the bias is 0x800 as the load/store instructions in the
     // ISA usually have a 12-bit immediate.
-    return phdr->p_vaddr + 0x800;
+    return phdr.p_vaddr + 0x800;
   } else {
     // On other targets, DTP simply refers to the beginning of the TLS block.
-    return phdr->p_vaddr;
+    return phdr.p_vaddr;
   }
 }
 
 using E = MOLD_TARGET;
 
-template u64 get_tls_begin<E>(Context<E> &);
-template u64 get_tp_addr<E>(Context<E> &);
-template u64 get_dtp_addr<E>(Context<E> &);
+template u64 get_tp_addr<E>(const ElfPhdr<E> &);
+template u64 get_dtp_addr<E>(const ElfPhdr<E> &);
 
 } // namespace mold::elf