From 2df6a02fd75fb7a5581753133794d68c7533b7eb Mon Sep 17 00:00:00 2001
From: Simon Cook <simon.cook@embecosm.com>
Date: Thu, 9 Apr 2020 15:11:55 +0100
Subject: [PATCH] [RISCV] Implement evaluateBranch

This implements the instruction analysis required to print branch
targets as part of llvm-objdump's disassembly.

Note, this only handles those branches which can be analyzed in a single
instruction, a future patch will handle multiple-instruction patterns,
such as AUIPC/LUI+JALR instruction pairs.

Differential Revision: https://reviews.llvm.org/D77567
---
 .../RISCV/MCTargetDesc/RISCVMCTargetDesc.cpp  | 41 +++++++++++++++++++
 .../MC/Disassembler/RISCV/branch-targets.txt  | 27 ++++++++++++
 2 files changed, 68 insertions(+)
 create mode 100644 llvm/test/MC/Disassembler/RISCV/branch-targets.txt

diff --git a/llvm/lib/Target/RISCV/MCTargetDesc/RISCVMCTargetDesc.cpp b/llvm/lib/Target/RISCV/MCTargetDesc/RISCVMCTargetDesc.cpp
index d231280cc1a4..cb515520b809 100644
--- a/llvm/lib/Target/RISCV/MCTargetDesc/RISCVMCTargetDesc.cpp
+++ b/llvm/lib/Target/RISCV/MCTargetDesc/RISCVMCTargetDesc.cpp
@@ -20,6 +20,7 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/CodeGen/Register.h"
 #include "llvm/MC/MCAsmInfo.h"
+#include "llvm/MC/MCInstrAnalysis.h"
 #include "llvm/MC/MCInstrInfo.h"
 #include "llvm/MC/MCRegisterInfo.h"
 #include "llvm/MC/MCStreamer.h"
@@ -93,6 +94,45 @@ static MCTargetStreamer *createRISCVAsmTargetStreamer(MCStreamer &S,
   return new RISCVTargetAsmStreamer(S, OS);
 }
 
+namespace {
+
+class RISCVMCInstrAnalysis : public MCInstrAnalysis {
+public:
+  explicit RISCVMCInstrAnalysis(const MCInstrInfo *Info)
+      : MCInstrAnalysis(Info) {}
+
+  bool evaluateBranch(const MCInst &Inst, uint64_t Addr, uint64_t Size,
+                      uint64_t &Target) const override {
+    if (isConditionalBranch(Inst)) {
+      int64_t Imm;
+      if (Size == 2)
+        Imm = Inst.getOperand(1).getImm();
+      else
+        Imm = Inst.getOperand(2).getImm();
+      Target = Addr + Imm;
+      return true;
+    }
+
+    if (Inst.getOpcode() == RISCV::C_JAL || Inst.getOpcode() == RISCV::C_J) {
+      Target = Addr + Inst.getOperand(0).getImm();
+      return true;
+    }
+
+    if (Inst.getOpcode() == RISCV::JAL) {
+      Target = Addr + Inst.getOperand(1).getImm();
+      return true;
+    }
+
+    return false;
+  }
+};
+
+} // end anonymous namespace
+
+static MCInstrAnalysis *createRISCVInstrAnalysis(const MCInstrInfo *Info) {
+  return new RISCVMCInstrAnalysis(Info);
+}
+
 extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeRISCVTargetMC() {
   for (Target *T : {&getTheRISCV32Target(), &getTheRISCV64Target()}) {
     TargetRegistry::RegisterMCAsmInfo(*T, createRISCVMCAsmInfo);
@@ -104,6 +144,7 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeRISCVTargetMC() {
     TargetRegistry::RegisterMCSubtargetInfo(*T, createRISCVMCSubtargetInfo);
     TargetRegistry::RegisterObjectTargetStreamer(
         *T, createRISCVObjectTargetStreamer);
+    TargetRegistry::RegisterMCInstrAnalysis(*T, createRISCVInstrAnalysis);
 
     // Register the asm target streamer.
     TargetRegistry::RegisterAsmTargetStreamer(*T, createRISCVAsmTargetStreamer);
diff --git a/llvm/test/MC/Disassembler/RISCV/branch-targets.txt b/llvm/test/MC/Disassembler/RISCV/branch-targets.txt
new file mode 100644
index 000000000000..d3486d2b99f9
--- /dev/null
+++ b/llvm/test/MC/Disassembler/RISCV/branch-targets.txt
@@ -0,0 +1,27 @@
+# RUN: llvm-mc -assemble -triple riscv32 -mattr=+c -filetype=obj %s -o - 2>&1 | \
+# RUN:   llvm-objdump -d --mattr=+c -M no-aliases - | FileCheck %s
+# RUN: llvm-mc -assemble -triple riscv64 -mattr=+c -filetype=obj %s -o - 2>&1 | \
+# RUN:   llvm-objdump -d --mattr=+c -M no-aliases - | FileCheck %s
+
+label1:
+.option norvc
+  j label1
+  j label2
+  bnez a0, label1
+  bnez a0, label2
+.option rvc
+  j label1
+  j label2
+  bnez a0, label1
+  bnez a0, label2
+# CHECK-LABEL: <label1>:
+# CHECK-NEXT: jal zero, 0 <label1>
+# CHECK-NEXT: jal zero, 20 <label2>
+# CHECK-NEXT: bne a0, zero, -8 <label1>
+# CHECK-NEXT: bne a0, zero, 12 <label2>
+# CHECK-NEXT: c.j -16 <label1>
+# CHECK-NEXT: c.j 6 <label2>
+# CHECK-NEXT: c.bnez a0, -20 <label1>
+# CHECK-NEXT: c.bnez a0, 2 <label2>
+
+label2: