[MC] Support returning marked up ranges in the disassembly

[seiya]

Summary

This patch introduces MCInstPrinter::setMarkupSpans, which allows the caller to
use the output of rich disassembly without parsing its output string.

An upcoming patch replaces markup() with startMarkup(), endMarkup(), and withMarkup().

This feature will be used for a disassembly highlighting feature in llvm-objdump.

[jhenderson]

llvm/include/llvm/MC/MCInstPrinter.h

44

Doxygen style comments please, to match the rest of the file.

64

Why is this a unique_ptr? Can't you just have a std::vector here directly?

80

This isn't a stack, it's a vector. Use the correct terminology for the type.

Could it be a vector of ArrayRefs?

83

Why not just Spans?

85–86

Don't use names with trailing underscores. Just do NewEnabled and NewSpans or whatever.

202

I'm not sure this comment makes sense. Also, the comments here should be in doxygen style, to match existing members.

Perhaps "Set the vector to write marked up ranges to" would be better phrasing.

[jhenderson]

llvm/include/llvm/MC/MCInstPrinter.h

205–206

This feels like bad code design to me, since it would be incredibly easy for a use to forget to do this. How about creating printInst and printInstImpl functions. The former calls resetMarkup and then printInstImpl. The latter is the pure virtual method that sub-classes implement.

llvm/lib/MC/MCInstPrinter.cpp

156–157

I'd probably just delete the second sentence here.

162

Use C++ style comments, and start the comment with an upper-case letter. It would be better to write this in the passive too: "Length and InnerLength will be set later."

[MaskRay]

llvm/include/llvm/MC/MCInstPrinter.h

74

MarkupState can be nested in MCInstprinter. It isn't expected to be accessed by other files.

[seiya]

• Addressed some review comments.

[seiya]

llvm/include/llvm/MC/MCInstPrinter.h

80

This isn't a stack, it's a vector. Use the correct terminology for the type.

Here I use std::vector as stack. I'd like to use std::stack instead but it does not provide clear method. I've updated the comment to stress my intention.

Could it be a vector of ArrayRefs?

No. We need a pointer to a vector to append elements into the vector.

205–206

I'm totally agree with you and adding printInstImpl sounds good idea. Thanks!

[seiya]

llvm/include/llvm/MC/MCInstPrinter.h

64

I think it should be a unique_ptr to prevent dangling pointer reference here:

auto *InnerSpans =
    const_cast<std::vector<MarkupSpan> *>(Span.InnerSpans.get());
M.State.Spans.back()->push_back(std::move(Span));
// Span is now moved to the "M.State.Spans.back()" but InnerSpans is still valid
// since Span.InnerSpans is an unique_ptr.
M.State.Spans.push_back(InnerSpans);

[seiya]

I'll consider the way to eliminate MarkupState.reset so that we don't have to introduce printInstImpl to existing all InstPrinters.

[seiya]

llvm/include/llvm/MC/MCInstPrinter.h

64

Ah, I realized we don't need unique_ptr here. We can rewrite the code above as:

M.State.Spans.back()->push_back(std::move(Span));
M.State.Spans.push_back(&M.State.Spans.back()->back().InnerSpans);

I'll update the patch later.

[seiya]

• Removed MarkupState so that we don't have to update all MCInstPrinters.
• Addressed remaining review comments.

[seiya]

• Fixed a comment.

[MaskRay]

llvm/lib/MC/MCInstPrinter.cpp

145

I suggest deleting this TODO. After rL166639 (2012), the document has never been touched. Actually, markup() has been adopted in every target so explicitly mentioning this may not be very necessary.

[seiya]

• Addressed a review comment.

[jhenderson]

llvm/include/llvm/MC/MCInstPrinter.h

167–168

You could just use while(!MarupSpans.empty()) MarkupSpans.pop(); instead of clear, if you want to use a std::stack.

[seiya]

• Addressed a review comment.

llvm/include/llvm/MC/MCInstPrinter.h

167–168

Fixed. First I thought iterating all elements (O(n)) is not good but it won't be issue because n is small (typically n < 3) here...

[jhenderson]

LGTM.

llvm/include/llvm/MC/MCInstPrinter.h

167–168

I think vector::clear is likely O(n) time as well - it has to destroy all its elements.

[MaskRay]

llvm/lib/MC/MCInstPrinter.cpp

147

M.Spans.top()->emplace_back(M.Type, OS.tell(), 0, OS.tell() + Length, 0);

[seiya]

• Addressed a review comment.

[seiya]

llvm/include/llvm/MC/MCInstPrinter.h

167–168

There's no guarantee but I implicitly assumed that there must be an O(1) optimization for
trivially destructible types like pointers.

[rupprecht]

llvm/include/llvm/MC/MCInstPrinter.h

72

Can you add brief class-level comments for MarkupStart/MarkupEnd too?

llvm/lib/MC/MCInstPrinter.cpp

159

Another design consideration would be to put the functionality of endMarkup() into the destructor MarkupStart::~MarkupStart()

e.g.

OS << startMarkup(MarkupType::Reg) << '%' << getRegisterName(RegNo)
     << endMarkup();

turns into:

OS << startMarkup(MarkupType::Reg) << '%' << getRegisterName(RegNo); /* ~MarkupStart() ends the markup */

Though it would require changing all the call sites so the register/operand/whatever can be printed as part of the same OS << ... expression, so I wouldn't recommend trying that until later. But it may be worth looking into to avoid programmer error of forgetting either of markupStart()/markupEnd()

[seiya]

I'll try out the destructor approach suggested by @rupprecht.

llvm/lib/MC/MCInstPrinter.cpp

159

The destructor approach sounds good idea! I'll try it.

[seiya]

• Added WithMarkup.
• Added class-level comments for MarkupStart and MarkupEnd.

[seiya]

llvm/lib/MC/MCInstPrinter.cpp

159

I've added WithMarkup, which does the functionality of endMarkup in the destructor (just like WithColor).

[jhenderson]

You should update your patch title and description too.

llvm/include/llvm/MC/MCInstPrinter.h

74

This comment doesn't make much sense (it refers to startupMarkup() too many times, I think). Please fix it.

[rupprecht]

llvm/include/llvm/MC/MCInstPrinter.h

64

Can you include an example where you might have two InnerSpans?

80

When do you need to have multiple Spans on the stack?
Can you add test coverage for that case?

[seiya]

• Updated comments.

[seiya]

In D65189#1625303, @jhenderson wrote:

You should update your patch title and description too.

What do you think what the patch title should describe?

How about "[MC] Support returning marked up ranges in the disassembly" or "[MCInstPrinter] Add MarkupStart/MarkupEnd/WithMarkup?

llvm/include/llvm/MC/MCInstPrinter.h

80

When do you need to have multiple Spans on the stack?

The stack has multiple Spans if the marked up ranges are nested:

movq <reg:%rax>, <mem:256(<reg:%rdi>)>
     ^^^^^^^^^^  ^^^^^^^^^----------^^
      1st span    2nd span (*)
                          ^^^^^^^^^^
                           1st span in the (*)'s InnerSpans

In this example, when we start disassembling <reg:%rdi>, the pointer to InnerSpans of the <mem:...> span will be appended into the stack.

Can you add test coverage for that case?

test/MC/Disassembler/X86/marked-up.txt includes such a case. We don't add a test in this patch so this should be committed along with D65190, btw.

[jhenderson]

In D65189#1626565, @seiya wrote:

In D65189#1625303, @jhenderson wrote:

You should update your patch title and description too.

What do you think what the patch title should describe?

How about "[MC] Support returning marked up ranges in the disassembly" or "[MCInstPrinter] Add MarkupStart/MarkupEnd/WithMarkup?

I think the first of your suggestions is good.

Nit: "This feature will be used for the highlighting feature in llvm-objdump." should be "This feature will be used for a disassembly highlighting feature in llvm-objdump."

[seiya]

Updated the title and description. Thanks @jhenderson!

[seiya]

• Make MarkupSpans mutable (suggested in https://reviews.llvm.org/D65190#inline-589645)

This is the last change to this patch if it's accepted.

[seiya]

• Improved comments for WithMarkup a bit.

[jhenderson]

LGTM. No need to get further review from my point of view after you've fixed my remaining nit.

llvm/include/llvm/MC/MCInstPrinter.h

211–212

Nit: looks like you've got two spaces before the const in these two lines.

[seiya]

• Removed redundant whitespaces.

[MaskRay]

Does this need rebase now?

llvm/include/llvm/MC/MCInstPrinter.h

 //===- MCInstPrinter.h - MCInst to target assembly syntax -------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_MC_MCINSTPRINTER_H
 #define LLVM_MC_MCINSTPRINTER_H
 
 #include "llvm/Support/Format.h"
 #include <cstdint>
+#include <stack>
 
 namespace llvm {
 
 class MCAsmInfo;
 class MCInst;
 class MCInstrInfo;
 class MCRegisterInfo;
 class MCSubtargetInfo;
 class raw_ostream;
 class StringRef;
 
 /// Convert `Bytes' to a hex string and output to `OS'
 void dumpBytes(ArrayRef<uint8_t> Bytes, raw_ostream &OS);
 
 namespace HexStyle {
 
 enum Style {
   C,  ///< 0xff
   Asm ///< 0ffh
 };
 
 } // end namespace HexStyle
 
+enum class MarkupType {
+  Reg,
+  Imm,
+  Mem,
+};
+
+/// MarkupSpan represents a marked up range in the disassembly. For example:
+///
+///     Pos   InnerPos
+///      v    v
+///  ... <mem:256(<reg:%rip>)> ...
+///           ~~~~~~~~~~~~~~~ InnerLenth
+///      ~~~~~~~~~~~~~~~~~~~~~ Length
+///
+struct MarkupSpan {
+  MarkupType Type;
+  /// The offset of the beginning of the marked up range in the stream.
+  size_t Pos;
+  /// The length of the marked up range.
+  size_t Length;
+  /// The offset of the beginning of the inner text in the stream.
+  size_t InnerPos;
+  /// The length of the inner text.
+  size_t InnerLength;
+  /// Marked up ranges in the inner text. For example:
+  ///
+  ///   movq <reg:%rax>, <mem:256(<reg:%rdi>)>
+  ///        ^^^^^^^^^^  ^^^^^^^^^----------^^
+  ///         1st span    2nd span (*)
+  ///                             ^^^^^^^^^^
+  ///                              1st span in the (*)'s InnerSpans
+  ///
+  /// In this example, the top-level InnerSpans contains two elements that
+  /// represent <reg:%rax> and <mem:...> respectively. InnerSpans in the
+  /// latter one contains a span which represents <reg:%rdi>.
+  std::vector<MarkupSpan> InnerSpans;
+
+  MarkupSpan(MarkupType Type, size_t Pos, size_t Length, size_t InnerPos,
+             size_t InnerLength)
+      : Type(Type), Pos(Pos), Length(Length), InnerPos(InnerPos),
+        InnerLength(InnerLength) {}
+};
+
+/// An object which creates a new markup range. When it is streamed to a
+/// raw_ostream, along with printing the beginning part of markup string
+/// (e.g. "<reg:"), a new incomplete MarkupSpan is appended to the given Spans.
+///
+/// The incomplete MarkupSpan will be updated by the corresponding MarkupEnd.
+///
+/// MarkupStart should be instantiated through MCInstPrinter::startMarkup():
+///
+///   OS << startMarkup(MarkupType::Reg) << "%rax" << endMarkup();
+///
+class MarkupStart {
+  bool Enabled;
+  std::stack<std::vector<MarkupSpan> *> &Spans;
+  MarkupType Type;
+
+public:
+  MarkupStart(bool Enabled, std::stack<std::vector<MarkupSpan> *> &Spans,
+              MarkupType Type)
+      : Enabled(Enabled), Spans(Spans), Type(Type) {}
+  friend raw_ostream &operator<<(raw_ostream &OS, const MarkupStart &M);
+};
+
+/// An object which closes the current markup range created by MarkupStart.
+///
+/// When it is streamed to a raw_ostream, along with printing the ending part
+/// of markup string (">"), the last MarkupSpan in Spans will be updated with
+/// correct values.
+///
+/// MarkupEnd should be instantiated through MCInstPrinter::endMarkup():
+///
+///   OS << startMarkup(MarkupType::Reg) << "%rax" << endMarkup();
+///
+class MarkupEnd {
+  bool Enabled;
+  std::stack<std::vector<MarkupSpan> *> &Spans;
+
+public:
+  MarkupEnd(bool Enabled, std::stack<std::vector<MarkupSpan> *> &Spans)
+      : Enabled(Enabled), Spans(Spans) {}
+  friend raw_ostream &operator<<(raw_ostream &OS, const MarkupEnd &M);
+};
+
+/// A RAII object which automatically closes the markup range to be used like:
+///
+///   WithMarkup(OS, MarkupType::Reg) << "%rax";
+///
+/// This is equivalent to:
+///   OS << startMarkup(MarkupType::Reg) << "%rax" << endMarkup();
+class WithMarkup {
+  raw_ostream &OS;
+  bool Enabled;
+  std::stack<std::vector<MarkupSpan> *> &Spans;
+
+public:
+  WithMarkup(raw_ostream &OS, bool Enabled,
+             std::stack<std::vector<MarkupSpan> *> &Spans, MarkupType Type);
+  ~WithMarkup();
+
+  template <typename T> WithMarkup &operator<<(const T &O) {
+    OS << O;
+    return *this;
+  }
+};
+
 /// This is an instance of a target assembly language printer that
 /// converts an MCInst to valid target assembly syntax.
 class MCInstPrinter {
 protected:
   /// A stream that comments can be emitted to if desired.  Each comment
   /// must end with a newline.  This will be null if verbose assembly emission
   /// is disabled.
   raw_ostream *CommentStream = nullptr;
   const MCAsmInfo &MAI;
   const MCInstrInfo &MII;
   const MCRegisterInfo &MRI;
 
   /// True if we are printing marked up assembly.
   bool UseMarkup = false;
 
   /// True if we are printing immediates as hex.
   bool PrintImmHex = false;
 
   /// Which style to use for printing hexadecimal values.
   HexStyle::Style PrintHexStyle = HexStyle::C;
 
+  /// A stack of pointers which points to &MS set by setMarkupSpans or
+  /// InnerSpans of unclosed MarkupSpans. This is mutable because it will be
+  /// updated in a const virtual method "printRegName" and making it non-const
+  /// requires relatively large changes to the existing code base.
+  mutable std::stack<std::vector<MarkupSpan> *> MarkupSpans;
+
   /// Utility function for printing annotations.
   void printAnnotation(raw_ostream &OS, StringRef Annot);
 
 public:
   MCInstPrinter(const MCAsmInfo &mai, const MCInstrInfo &mii,
                 const MCRegisterInfo &mri) : MAI(mai), MII(mii), MRI(mri) {}
 
   virtual ~MCInstPrinter();
public:
   StringRef getOpcodeName(unsigned Opcode) const;
 
   /// Print the assembler register name.
   virtual void printRegName(raw_ostream &OS, unsigned RegNo) const;
 
   bool getUseMarkup() const { return UseMarkup; }
   void setUseMarkup(bool Value) { UseMarkup = Value; }
 
+  /// Set the vector to write marked up ranges to.
+  void setMarkupSpans(std::vector<MarkupSpan> &MS) {
+    while (!MarkupSpans.empty())
+      MarkupSpans.pop();
+    MarkupSpans.push(&MS);
+  }
+
   /// Utility functions to make adding mark ups simpler.
   StringRef markup(StringRef s) const;
+  MarkupStart startMarkup(MarkupType Type)  const;
+  MarkupEnd endMarkup()  const;
+  WithMarkup withMarkup(raw_ostream &OS, MarkupType Type) const;
 
   bool getPrintImmHex() const { return PrintImmHex; }
   void setPrintImmHex(bool Value) { PrintImmHex = Value; }
 
   void setPrintHexStyle(HexStyle::Style Value) { PrintHexStyle = Value; }
 
   /// Utility function to print immediates in decimal or hex.
   format_object<int64_t> formatImm(int64_t Value) const {

llvm/lib/MC/MCInstPrinter.cpp

format_object<uint64_t> MCInstPrinter::formatHex(uint64_t Value) const {
   case HexStyle::Asm:
     if (needsLeadingZero(Value))
       return format("0%" PRIx64 "h", Value);
     else
       return format("%" PRIx64 "h", Value);
   }
   llvm_unreachable("unsupported print style");
 }
+
+MarkupStart MCInstPrinter::startMarkup(MarkupType Type) const {
+  return MarkupStart(getUseMarkup(), MarkupSpans, Type);
+}
+
+MarkupEnd MCInstPrinter::endMarkup() const {
+  return MarkupEnd(getUseMarkup(), MarkupSpans);
+}
+
+WithMarkup MCInstPrinter::withMarkup(raw_ostream &OS, MarkupType Type) const {
+  return WithMarkup(OS, getUseMarkup(), MarkupSpans, Type);
+}
+
+raw_ostream &llvm::operator<<(raw_ostream &OS, const MarkupStart &M) {
+  if (M.Enabled) {
+    StringRef TypeStr;
+    switch (M.Type) {
+    case MarkupType::Imm:
+      TypeStr = "imm";
+      break;
+    case MarkupType::Reg:
+      TypeStr = "reg";
+      break;
+    case MarkupType::Mem:
+      TypeStr = "mem";
+      break;
+    }
+
+    size_t Length = 2 + TypeStr.size();
+    if (!M.Spans.empty()) {
+      // Length and InnerLength will be set later.
+      M.Spans.top()->emplace_back(M.Type, OS.tell(), 0, OS.tell() + Length, 0);
+      M.Spans.push(&M.Spans.top()->back().InnerSpans);
+    }
+    OS << "<" << TypeStr << ":";
+  }
+
+  return OS;
+}
+
+raw_ostream &llvm::operator<<(raw_ostream &OS, const MarkupEnd &M) {
+  if (M.Enabled) {
+    if (!M.Spans.empty()) {
+      assert(M.Spans.size() > 1 && "Missing the corresponding markupStart().");
+      M.Spans.pop();
+      MarkupSpan &Span = M.Spans.top()->back();
+      Span.Length = OS.tell() - Span.Pos + 1;
+      Span.InnerLength = OS.tell() - Span.InnerPos;
+    }
+    OS << ">";
+  }
+
+  return OS;
+}
+
+WithMarkup::WithMarkup(raw_ostream &OS, bool Enabled,
+                       std::stack<std::vector<MarkupSpan> *> &Spans,
+                       MarkupType Type)
+    : OS(OS), Enabled(Enabled), Spans(Spans) {
+  OS << MarkupStart(Enabled, Spans, Type);
+}
+
+WithMarkup::~WithMarkup() { OS << MarkupEnd(Enabled, Spans); }