[llvm] [RISCV][Disassemble] Ensure the comment stream of the disassembler is set. (PR #125962)

[Francesco Petrogalli via llvm-commits]

https://github.com/fpetrogalli created https://github.com/llvm/llvm-project/pull/125962

None

>From 4e1f72a8acc714157ae5b894fd97e6895e136087 Mon Sep 17 00:00:00 2001
From: Francesco Petrogalli <francesco.petrogalli at apple.com>
Date: Wed, 5 Feb 2025 15:45:09 -0800
Subject: [PATCH] [RISCV][Disassemble] Ensure the comment stream of the
 disassembler is set.

---
 llvm/include/llvm/MC/MCDisassembler/MCDisassembler.h     | 5 +++++
 llvm/lib/MC/MCDisassembler/MCDisassembler.cpp            | 8 ++++++--
 llvm/lib/Target/RISCV/Disassembler/RISCVDisassembler.cpp | 1 +
 3 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/llvm/include/llvm/MC/MCDisassembler/MCDisassembler.h b/llvm/include/llvm/MC/MCDisassembler/MCDisassembler.h
index 901bfcf5fa54f4e..cb773d0f9589c46 100644
--- a/llvm/include/llvm/MC/MCDisassembler/MCDisassembler.h
+++ b/llvm/include/llvm/MC/MCDisassembler/MCDisassembler.h
@@ -131,6 +131,11 @@ class MCDisassembler {
   ///                   MCDisassembler::SoftFail if the instruction was
   ///                                            disassemblable but invalid,
   ///                   MCDisassembler::Fail if the instruction was invalid.
+  ///
+  /// Note: to avoid potential issues caused by the field
+  /// `MCDisassembler::CommentStream` being set nullptr (its default
+  /// value), an implementation of this method should make sure to set
+  /// `CommentStream = &CStream;`.
   virtual DecodeStatus getInstruction(MCInst &Instr, uint64_t &Size,
                                       ArrayRef<uint8_t> Bytes, uint64_t Address,
                                       raw_ostream &CStream) const = 0;
diff --git a/llvm/lib/MC/MCDisassembler/MCDisassembler.cpp b/llvm/lib/MC/MCDisassembler/MCDisassembler.cpp
index 6aa4b0e4fcb99da..0dfa17eb1346fdb 100644
--- a/llvm/lib/MC/MCDisassembler/MCDisassembler.cpp
+++ b/llvm/lib/MC/MCDisassembler/MCDisassembler.cpp
@@ -29,17 +29,21 @@ bool MCDisassembler::tryAddingSymbolicOperand(MCInst &Inst, int64_t Value,
                                               uint64_t Address, bool IsBranch,
                                               uint64_t Offset, uint64_t OpSize,
                                               uint64_t InstSize) const {
-  if (Symbolizer)
+  if (Symbolizer) {
+    assert(CommentStream && "CommentStream is not set.");
     return Symbolizer->tryAddingSymbolicOperand(Inst, *CommentStream, Value,
                                                 Address, IsBranch, Offset,
                                                 OpSize, InstSize);
+  }
   return false;
 }
 
 void MCDisassembler::tryAddingPcLoadReferenceComment(int64_t Value,
                                                      uint64_t Address) const {
-  if (Symbolizer)
+  if (Symbolizer) {
+    assert(CommentStream && "CommentStream is not set.");
     Symbolizer->tryAddingPcLoadReferenceComment(*CommentStream, Value, Address);
+  }
 }
 
 void MCDisassembler::setSymbolizer(std::unique_ptr<MCSymbolizer> Symzer) {
diff --git a/llvm/lib/Target/RISCV/Disassembler/RISCVDisassembler.cpp b/llvm/lib/Target/RISCV/Disassembler/RISCVDisassembler.cpp
index 3ec465810b1d114..01648ec0cbe9ea9 100644
--- a/llvm/lib/Target/RISCV/Disassembler/RISCVDisassembler.cpp
+++ b/llvm/lib/Target/RISCV/Disassembler/RISCVDisassembler.cpp
@@ -789,6 +789,7 @@ DecodeStatus RISCVDisassembler::getInstruction(MCInst &MI, uint64_t &Size,
                                                ArrayRef<uint8_t> Bytes,
                                                uint64_t Address,
                                                raw_ostream &CS) const {
+  CommentStream = &CS;
   // It's a 16 bit instruction if bit 0 and 1 are not 0b11.
   if ((Bytes[0] & 0b11) != 0b11)
     return getInstruction16(MI, Size, Bytes, Address, CS);