From 7949c2a943c75a39922745f1e4156eea3ef42085 Mon Sep 17 00:00:00 2001
From: Maksim Panchenko <maks@fb.com>
Date: Fri, 28 Feb 2025 19:11:04 -0800
Subject: [PATCH] [BOLT][AArch64] Refactor ADR to ADRP+ADD conversion pass.
 NFCI

In preparation of using the new interface in more places, refactor the
ADR conversion pass.
---
 bolt/include/bolt/Core/MCPlusBuilder.h           | 11 +++++++----
 bolt/lib/Passes/ADRRelaxationPass.cpp            | 10 +++-------
 bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp | 14 ++++++++++++--
 3 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index fbb853656fb91..1d45a314a17b6 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -637,10 +637,6 @@ class MCPlusBuilder {
     return false;
   }
 
-  virtual void getADRReg(const MCInst &Inst, MCPhysReg &RegName) const {
-    llvm_unreachable("not implemented");
-  }
-
   virtual bool isMoveMem2Reg(const MCInst &Inst) const { return false; }
 
   virtual bool mayLoad(const MCInst &Inst) const {
@@ -1538,6 +1534,13 @@ class MCPlusBuilder {
     llvm_unreachable("not implemented");
   }
 
+  /// Undo the linker's ADRP+ADD to ADR relaxation. Take \p ADRInst and return
+  /// ADRP+ADD instruction sequence.
+  virtual InstructionListType undoAdrpAddRelaxation(const MCInst &ADRInst,
+                                                    MCContext *Ctx) const {
+    llvm_unreachable("not implemented");
+  }
+
   /// Return not 0 if the instruction CurInst, in combination with the recent
   /// history of disassembled instructions supplied by [Begin, End), is a linker
   /// generated veneer/stub that needs patching. This happens in AArch64 when
diff --git a/bolt/lib/Passes/ADRRelaxationPass.cpp b/bolt/lib/Passes/ADRRelaxationPass.cpp
index 52811edcb8273..4b37a061ac12d 100644
--- a/bolt/lib/Passes/ADRRelaxationPass.cpp
+++ b/bolt/lib/Passes/ADRRelaxationPass.cpp
@@ -71,14 +71,10 @@ void ADRRelaxationPass::runOnFunction(BinaryFunction &BF) {
             continue;
       }
 
-      MCPhysReg Reg;
-      BC.MIB->getADRReg(Inst, Reg);
-      int64_t Addend = BC.MIB->getTargetAddend(Inst);
-      InstructionListType Addr;
-
+      InstructionListType AdrpAdd;
       {
         auto L = BC.scopeLock();
-        Addr = BC.MIB->materializeAddress(Symbol, BC.Ctx.get(), Reg, Addend);
+        AdrpAdd = BC.MIB->undoAdrpAddRelaxation(Inst, BC.Ctx.get());
       }
 
       if (It != BB.begin() && BC.MIB->isNoop(*std::prev(It))) {
@@ -99,7 +95,7 @@ void ADRRelaxationPass::runOnFunction(BinaryFunction &BF) {
         PassFailed = true;
         return;
       }
-      It = BB.replaceInstruction(It, Addr);
+      It = BB.replaceInstruction(It, AdrpAdd);
     }
   }
 }
diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index f79d5a747246e..6bcea7cde4b1c 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -278,13 +278,23 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     return Inst.getOpcode() == AArch64::ADDXri;
   }
 
-  void getADRReg(const MCInst &Inst, MCPhysReg &RegName) const override {
+  MCPhysReg getADRReg(const MCInst &Inst) const {
     assert((isADR(Inst) || isADRP(Inst)) && "Not an ADR instruction");
     assert(MCPlus::getNumPrimeOperands(Inst) != 0 &&
            "No operands for ADR instruction");
     assert(Inst.getOperand(0).isReg() &&
            "Unexpected operand in ADR instruction");
-    RegName = Inst.getOperand(0).getReg();
+    return Inst.getOperand(0).getReg();
+  }
+
+  InstructionListType undoAdrpAddRelaxation(const MCInst &ADRInst,
+                                            MCContext *Ctx) const override {
+    assert(isADR(ADRInst) && "ADR instruction expected");
+
+    const MCPhysReg Reg = getADRReg(ADRInst);
+    const MCSymbol *Target = getTargetSymbol(ADRInst);
+    const uint64_t Addend = getTargetAddend(ADRInst);
+    return materializeAddress(Target, Ctx, Reg, Addend);
   }
 
   bool isTB(const MCInst &Inst) const {