Add instruction highlight in graph view (#1879)

src/Cutter.pro

@@ -359,6 +359,7 @@ SOURCES += \
     common/PythonManager.cpp \
     plugins/PluginManager.cpp \
     common/BasicBlockHighlighter.cpp \
+    common/BasicInstructionHighlighter.cpp \
     dialogs/LinkTypeDialog.cpp \
     widgets/ColorPicker.cpp \
     common/ColorThemeWorker.cpp \
@@ -494,6 +495,7 @@ HEADERS  += \
     common/PythonManager.h \
     plugins/PluginManager.h \
     common/BasicBlockHighlighter.h \
+    common/BasicInstructionHighlighter.h \
     common/UpdateWorker.h \
     widgets/ColorPicker.h \
     common/ColorThemeWorker.h \

src/common/BasicInstructionHighlighter.cpp

@@ -0,0 +1,75 @@
+#include "BasicInstructionHighlighter.h"
+#include <vector>
+
+/**
+ * @brief Clear the basic instruction highlighting
+ */
+void BasicInstructionHighlighter::clear(RVA address, RVA size)
+{
+    BasicInstructionIt it = biMap.lower_bound(address);
+    if (it != biMap.begin()) {
+        --it;
+    }
+
+    std::vector<RVA> addrs;
+    while (it != biMap.end() && it->first < address + size) {
+        addrs.push_back(it->first);
+        ++it;
+    }
+
+    // first and last entries may intersect, but not necessarily
+    // be contained in [address, address + size), so we need to
+    // check it and perhaps adjust their addresses.
+    std::vector<BasicInstruction> newInstructions;
+    if (!addrs.empty()) {
+        const BasicInstruction &prev = biMap[addrs.front()];
+        if (prev.address < address && prev.address + prev.size > address) {
+            newInstructions.push_back({prev.address, address - prev.address, prev.color});
+        }
+
+        const BasicInstruction &next = biMap[addrs.back()];
+        if (next.address < address + size && next.address + next.size > address + size) {
+            const RVA offset = address + size - next.address;
+            newInstructions.push_back({next.address + offset, next.size - offset, next.color});
+        }
+    }
+
+    for (RVA addr : addrs) {
+        const BasicInstruction &bi = biMap[addr];
+        if (std::max(bi.address, address) < std::min(bi.address + bi.size, address + size)) {
+            biMap.erase(addr);
+        }
+    }
+
+    for ( BasicInstruction newInstr : newInstructions) {
+        biMap[newInstr.address] = newInstr;
+    }
+}
+
+/**
+ * @brief Highlight the basic instruction at address
+ */
+void BasicInstructionHighlighter::highlight(RVA address, RVA size, QColor color)
+{
+    clear(address, size);
+    biMap[address] = {address, size, color};
+}
+
+/**
+ * @brief Return a highlighted basic instruction
+ *
+ * If there is nothing to highlight at specified address, returns nullptr
+ */
+BasicInstruction *BasicInstructionHighlighter::getBasicInstruction(RVA address)
+{
+    BasicInstructionIt it = biMap.upper_bound(address);
+    if (it == biMap.begin()) {
+        return nullptr;
+    }
+
+    BasicInstruction *bi = &(--it)->second;
+    if (bi->address <= address && address < bi->address + bi->size) {
+        return bi;
+    }
+    return nullptr;
+}

src/common/BasicInstructionHighlighter.h

@@ -0,0 +1,27 @@
+#ifndef BASICINSTRUCTIONHIGHLIGHTER_H
+#define BASICINSTRUCTIONHIGHLIGHTER_H
+
+#include "CutterCommon.h"
+#include <map>
+#include <QColor>
+
+struct BasicInstruction {
+    RVA address;
+    RVA size;
+    QColor color;
+};
+
+typedef std::map<RVA, BasicInstruction>::iterator BasicInstructionIt;
+
+class BasicInstructionHighlighter
+{
+public:
+    void clear(RVA address, RVA size);
+    void highlight(RVA address, RVA size, QColor color);
+    BasicInstruction *getBasicInstruction(RVA address);
+
+private:
+    std::map<RVA, BasicInstruction> biMap;
+};
+
+#endif // BASICINSTRUCTIONHIGHLIGHTER_H

src/core/Cutter.cpp

@@ -8,6 +8,7 @@
 #include <memory>
 
 #include "common/TempConfig.h"
+#include "common/BasicInstructionHighlighter.h"
 #include "common/Configuration.h"
 #include "common/AsyncTask.h"
 #include "common/R2Task.h"
@@ -2895,6 +2896,10 @@ BasicBlockHighlighter* CutterCore::getBBHighlighter()
     return bbHighlighter;
 }
 
+BasicInstructionHighlighter* CutterCore::getBIHighlighter()
+{
+    return &biHighlighter;
+}
 
 /**
  * @brief get a compact disassembly preview for tooltips

src/core/Cutter.h

@@ -3,6 +3,7 @@
 
 #include "core/CutterCommon.h"
 #include "core/CutterDescriptions.h"
+#include "common/BasicInstructionHighlighter.h"
 
 #include <QMap>
 #include <QDebug>
@@ -14,6 +15,7 @@
 #include <QMutex>
 
 class AsyncTaskManager;
+class BasicInstructionHighlighter;
 class CutterCore;
 class Decompiler;
 
@@ -405,6 +407,7 @@ public:
 
     static QString ansiEscapeToHtml(const QString &text);
     BasicBlockHighlighter *getBBHighlighter();
+    BasicInstructionHighlighter *getBIHighlighter();
 
 signals:
     void refreshAll();
@@ -471,6 +474,7 @@ private:
 
     bool emptyGraph = false;
     BasicBlockHighlighter *bbHighlighter;
+    BasicInstructionHighlighter biHighlighter;
 };
 
 class RCoreLocked

src/widgets/DisassemblerGraphView.cpp

@@ -9,6 +9,7 @@
 #include "common/TempConfig.h"
 #include "common/SyntaxHighlighter.h"
 #include "common/BasicBlockHighlighter.h"
+#include "common/BasicInstructionHighlighter.h"
 #include "dialogs/MultitypeFileSaveDialog.h"
 #include "common/Helpers.h"
 
@@ -45,7 +46,8 @@ DisassemblerGraphView::DisassemblerGraphView(QWidget *parent, CutterSeekable *se
       contextMenu(new QMenu(this)),
       seekable(seekable),
       actionExportGraph(this),
-      actionUnhighlight(this)
+      actionUnhighlight(this),
+      actionUnhighlightInstruction(this)
 {
     highlight_token = nullptr;
     auto *layout = new QVBoxLayout(this);
@@ -155,8 +157,21 @@ DisassemblerGraphView::DisassemblerGraphView(QWidget *parent, CutterSeekable *se
         Config()->colorsUpdated();
     });
 
+    QAction *highlightBI = new QAction(this);
+    actionUnhighlightInstruction.setVisible(false);
+
+    highlightBI->setText(tr("Highlight instruction"));
+    connect(highlightBI, &QAction::triggered, this,
+            &DisassemblerGraphView::onActionHighlightBITriggered);
+
+    actionUnhighlightInstruction.setText(tr("Unhighlight instruction"));
+    connect(&actionUnhighlightInstruction, &QAction::triggered, this,
+            &DisassemblerGraphView::onActionUnhighlightBITriggered);
+
     blockMenu->addAction(highlightBB);
     blockMenu->addAction(&actionUnhighlight);
+    blockMenu->addAction(highlightBI);
+    blockMenu->addAction(&actionUnhighlightInstruction);
 
 
     // Include all actions from generic context menu in block specific menu
@@ -426,7 +441,6 @@ void DisassemblerGraphView::drawBlock(QPainter &p, GraphView::GraphBlock &block,
     // Render node
     DisassemblyBlock &db = disassembly_blocks[block.entry];
     bool block_selected = false;
-    bool PCInBlock = false;
     RVA selected_instruction = RVA_INVALID;
 
     // Figure out if the current block is selected
@@ -437,9 +451,6 @@ void DisassemblerGraphView::drawBlock(QPainter &p, GraphView::GraphBlock &block,
             block_selected = true;
             selected_instruction = instr.addr;
         }
-        if (instr.contains(PCAddr)) {
-            PCInBlock = true;
-        }
 
         // TODO: L219
     }
@@ -480,23 +491,6 @@ void DisassemblerGraphView::drawBlock(QPainter &p, GraphView::GraphBlock &block,
         return;
     }
 
-    // Draw different background for selected instruction
-    if (selected_instruction != RVA_INVALID) {
-        int y = firstInstructionY;
-        for (const Instr &instr : db.instrs) {
-            if (instr.addr > selected_instruction) {
-                break;
-            }
-            auto selected = instr.addr == selected_instruction;
-            if (selected) {
-                p.fillRect(QRect(static_cast<int>(block.x + charWidth), y,
-                                 static_cast<int>(block.width - (10 + padding)),
-                                 int(instr.text.lines.size()) * charHeight), disassemblySelectionColor);
-            }
-            y += int(instr.text.lines.size()) * charHeight;
-        }
-    }
-
     // Highlight selected tokens
     if (highlight_token != nullptr) {
         int y = firstInstructionY;
@@ -533,23 +527,6 @@ void DisassemblerGraphView::drawBlock(QPainter &p, GraphView::GraphBlock &block,
         }
     }
 
-    // Highlight program counter
-    if (PCInBlock) {
-        int y = firstInstructionY;
-        for (const Instr &instr : db.instrs) {
-            if (instr.addr > PCAddr) {
-                break;
-            }
-            auto PC = instr.addr == PCAddr;
-            if (PC) {
-                p.fillRect(QRect(static_cast<int>(block.x + charWidth), y,
-                                 static_cast<int>(block.width - (10 + padding)),
-                                 int(instr.text.lines.size()) * charHeight), PCSelectionColor);
-            }
-            y += int(instr.text.lines.size()) * charHeight;
-        }
-    }
-
     // Render node text
     auto x = block.x + padding;
     int y = block.y + getTextOffset(0).y();
@@ -559,17 +536,29 @@ void DisassemblerGraphView::drawBlock(QPainter &p, GraphView::GraphBlock &block,
         y += charHeight;
     }
 
+    auto bih = Core()->getBIHighlighter();
     for (const Instr &instr : db.instrs) {
+        const QRect instrRect = QRect(static_cast<int>(block.x + charWidth), y,
+                                      static_cast<int>(block.width - (10 + padding)),
+                                      int(instr.text.lines.size()) * charHeight);
+
+        QColor instrColor;
         if (Core()->isBreakpoint(breakpoints, instr.addr)) {
-            p.fillRect(QRect(static_cast<int>(block.x + charWidth), y,
-                             static_cast<int>(block.width - (10 + padding)),
-                             int(instr.text.lines.size()) * charHeight), ConfigColor("gui.breakpoint_background"));
-            if (instr.addr == selected_instruction) {
-                p.fillRect(QRect(static_cast<int>(block.x + charWidth), y,
-                                 static_cast<int>(block.width - (10 + padding)),
-                                 int(instr.text.lines.size()) * charHeight), disassemblySelectionColor);
+            instrColor = ConfigColor("gui.breakpoint_background");
+        } else if (instr.addr == PCAddr) {
+            instrColor = PCSelectionColor;
+        } else if (auto background = bih->getBasicInstruction(instr.addr)) {
+            instrColor = background->color;
         }
+
+        if (instrColor.isValid()) {
+            p.fillRect(instrRect, instrColor);
         }
+
+        if (selected_instruction != RVA_INVALID && selected_instruction == instr.addr) {
+            p.fillRect(instrRect, disassemblySelectionColor);
+        }
+
         for (auto &line : instr.text.lines) {
             int rectSize = qRound(charWidth);
             if (rectSize % 2) {
@@ -755,6 +744,21 @@ DisassemblerGraphView::DisassemblyBlock *DisassemblerGraphView::blockForAddress(
     return nullptr;
 }
 
+const DisassemblerGraphView::Instr *DisassemblerGraphView::instrForAddress(RVA addr)
+{
+    DisassemblyBlock *block = blockForAddress(addr);
+    for (const Instr &i : block->instrs) {
+        if (i.addr == RVA_INVALID || i.size == RVA_INVALID) {
+            continue;
+        }
+
+        if (i.contains(addr)) {
+            return &i;
+        }
+    }
+    return nullptr;
+}
+
 void DisassemblerGraphView::onSeekChanged(RVA addr)
 {
     blockMenu->setOffset(addr);
@@ -971,7 +975,9 @@ void DisassemblerGraphView::blockClicked(GraphView::GraphBlock &block, QMouseEve
 void DisassemblerGraphView::blockContextMenuRequested(GraphView::GraphBlock &block,
                                                       QContextMenuEvent *event, QPoint pos)
 {
+    const RVA offset = this->seekable->getOffset();
     actionUnhighlight.setVisible(Core()->getBBHighlighter()->getBasicBlock(block.entry));
+    actionUnhighlightInstruction.setVisible(Core()->getBIHighlighter()->getBasicInstruction(offset));
     event->accept();
     blockMenu->exec(event->globalPos());
 }
@@ -1123,6 +1129,43 @@ void DisassemblerGraphView::on_actionExportGraph_triggered()
 
 }
 
+void DisassemblerGraphView::onActionHighlightBITriggered()
+{
+    const RVA offset = this->seekable->getOffset();
+    const Instr *instr = instrForAddress(offset);
+
+    if (!instr) {
+        return;
+    }
+
+    auto bih = Core()->getBIHighlighter();
+    QColor background = ConfigColor("linehl");
+    if (auto currentColor = bih->getBasicInstruction(offset)) {
+        background = currentColor->color;
+    }
+
+    QColor c = QColorDialog::getColor(background, this, QString(),
+                                      QColorDialog::DontUseNativeDialog);
+    if (c.isValid()) {
+        bih->highlight(instr->addr, instr->size, c);
+    }
+    Config()->colorsUpdated();
+}
+
+void DisassemblerGraphView::onActionUnhighlightBITriggered()
+{
+    const RVA offset = this->seekable->getOffset();
+    const Instr *instr = instrForAddress(offset);
+
+    if (!instr) {
+        return;
+    }
+
+    auto bih = Core()->getBIHighlighter();
+    bih->clear(instr->addr, instr->size);
+    Config()->colorsUpdated();
+}
+
 void DisassemblerGraphView::exportGraph(QString filePath, GraphExportType type)
 {
     switch (type) {

src/widgets/DisassemblerGraphView.h

@@ -157,6 +157,8 @@ protected:
 
 private slots:
     void on_actionExportGraph_triggered();
+    void onActionHighlightBITriggered();
+    void onActionUnhighlightBITriggered();
 
 private:
     bool transition_dont_seek = false;
@@ -193,6 +195,7 @@ private:
      */
     QRectF getInstrRect(GraphView::GraphBlock &block, RVA addr) const;
     void showInstruction(GraphView::GraphBlock &block, RVA addr);
+    const Instr *instrForAddress(RVA addr);
     DisassemblyBlock *blockForAddress(RVA addr);
     void seekLocal(RVA addr, bool update_viewport = true);
     void seekInstruction(bool previous_instr);
@@ -224,6 +227,7 @@ private:
 
     QAction actionExportGraph;
     QAction actionUnhighlight;
+    QAction actionUnhighlightInstruction;
 
     QLabel *emptyText = nullptr;