[SharedCache] Fix some bugs

4 files changed, 81 insertions, 24 deletions

diff --git a/view/sharedcache/ui/dscpicker.cpp b/view/sharedcache/ui/dscpicker.cpp
index c6448ff4..340b8aa8 100644
--- a/view/sharedcache/ui/dscpicker.cpp
+++ b/view/sharedcache/ui/dscpicker.cpp
@@ -11,7 +11,7 @@ using namespace SharedCacheAPI;
 
 void DisplayDSCPicker(UIContext* ctx, Ref<BinaryView> dscView)
 {
-	static auto getImageNames = [dscView](QVariant var) {
+	auto getImageNames = [dscView](QVariant var) {
 		auto controller = SharedCacheController::GetController(*dscView);
 		if (!controller)
 			return QStringList();
@@ -22,7 +22,7 @@ void DisplayDSCPicker(UIContext* ctx, Ref<BinaryView> dscView)
 		return entries;
 	};
 
-	static auto getChosenImage = [ctx](QVariant var) {
+	auto getChosenImage = [ctx](QVariant var) {
 		QStringList entries = var.toStringList();
 
 		auto choiceDialog = new MetadataChoiceDialog(ctx ? ctx->mainWindow() : nullptr, "Pick Image", "Select", entries);
@@ -36,7 +36,7 @@ void DisplayDSCPicker(UIContext* ctx, Ref<BinaryView> dscView)
 		return QVariant(QString::fromStdString(entries.at((qsizetype)choiceDialog->GetChosenEntry().value().idx).toStdString()));
 	};
 
-	static auto loadSelectedImage = [dscView](QVariant var) {
+	auto loadSelectedImage = [dscView](QVariant var) {
 		auto selectedImageName = var.toString().toStdString();
 		if (selectedImageName.empty())
 			return;
@@ -54,8 +54,8 @@ void DisplayDSCPicker(UIContext* ctx, Ref<BinaryView> dscView)
 
 
 	BackgroundThread::create(ctx ? ctx->mainWindow() : nullptr)
-		->thenBackground([](QVariant var){ return getImageNames(var); })
-		->thenMainThread([](QVariant var){ return getChosenImage(var); })
-		->thenBackground([](QVariant var){ return loadSelectedImage(var); })
+		->thenBackground([getImageNames](QVariant var){ return getImageNames(var); })
+		->thenMainThread([getChosenImage](QVariant var){ return getChosenImage(var); })
+		->thenBackground([loadSelectedImage](QVariant var){ return loadSelectedImage(var); })
 		->start();
 }
diff --git a/view/sharedcache/ui/dsctriage.cpp b/view/sharedcache/ui/dsctriage.cpp
index aab71aac..297612d5 100644
--- a/view/sharedcache/ui/dsctriage.cpp
+++ b/view/sharedcache/ui/dsctriage.cpp
@@ -54,7 +54,7 @@ DSCTriageView::DSCTriageView(QWidget* parent, BinaryViewRef data) : QWidget(pare
 
 	QWidget* defaultWidget = nullptr;
 
-	static auto loadImagesWithAddr = [this](const std::vector<uint64_t>& addresses) {
+	auto loadImagesWithAddr = [this](const std::vector<uint64_t>& addresses) {
 		auto controller = SharedCacheController::GetController(*this->m_data);
 		if (!controller)
 			return;
@@ -102,10 +102,13 @@ DSCTriageView::DSCTriageView(QWidget* parent, BinaryViewRef data) : QWidget(pare
 			m_imageModel->setHorizontalHeaderLabels({"Address", "Loaded", "Name"});
 		} // loadImageModel
 
+		// Show the icon in the loaded column.
+		loadImageTable->setItemDelegateForColumn(1, new LoadedDelegate());
+
 		auto loadImageButton = new CustomStyleFlatPushButton();
 		{
 			connect(loadImageButton, &QPushButton::clicked,
-				[loadImageTable](bool) {
+				[loadImageTable, loadImagesWithAddr](bool) {
 					auto selected = loadImageTable->selectionModel()->selectedRows();
 					std::vector<uint64_t> addresses;
 					for (const auto& row : selected)
@@ -188,7 +191,7 @@ DSCTriageView::DSCTriageView(QWidget* parent, BinaryViewRef data) : QWidget(pare
 		auto loadSymbolImageButton = new CustomStyleFlatPushButton();
 		{
 			connect(loadSymbolImageButton, &QPushButton::clicked,
-				[this](bool) {
+				[this, loadImagesWithAddr](bool) {
 					auto selected = m_symbolTable->selectionModel()->selectedRows();
 					std::vector<uint64_t> addresses;
 					for (const auto& row : selected)
@@ -236,12 +239,6 @@ DSCTriageView::DSCTriageView(QWidget* parent, BinaryViewRef data) : QWidget(pare
 		auto symbolWidget = new QWidget;
 		symbolWidget->setLayout(symbolLayout);
 
-		m_symbolTable->horizontalHeader()->setSectionResizeMode(0, QHeaderView::ResizeToContents); // Address
-		m_symbolTable->horizontalHeader()->setSectionResizeMode(1, QHeaderView::Stretch);          // Name
-
-		m_symbolTable->setSelectionBehavior(QAbstractItemView::SelectRows);
-		m_symbolTable->setSelectionMode(QAbstractItemView::SingleSelection);
-
 		connect(m_symbolTable, &SymbolTableView::activated, this, [=](const QModelIndex& index){
 				auto symbol = m_symbolTable->getSymbolAtRow(index.row());
 				auto dialog = new QMessageBox(this);
@@ -444,9 +441,6 @@ void DSCTriageView::RefreshData()
 
 void DSCTriageView::setImageLoaded(const uint64_t imageHeaderAddr)
 {
-	// TODO: Better icon... probably something like ✅
-	static QIcon loadedIcon(":/icons/images/plus.png");
-
 	// Go through the m_loadImageModel and find the image associated with the address
 	// then set the image as loaded.
 	for (int i = 0; i < m_imageModel->rowCount(); i++)
@@ -456,8 +450,9 @@ void DSCTriageView::setImageLoaded(const uint64_t imageHeaderAddr)
 		if (addr == imageHeaderAddr)
 		{
 			auto statusCol = m_imageModel->index(i, 1);
-			m_imageModel->setData(statusCol, QString("•"), Qt::DisplayRole);
+			// See the `LoadedDelegate` class, we set 1 to indicate that this image is loaded.
+			m_imageModel->setData(statusCol, "1", Qt::DisplayRole);
 			break;
 		}
 	}
-}
+}
\ No newline at end of file
diff --git a/view/sharedcache/ui/dsctriage.h b/view/sharedcache/ui/dsctriage.h
index 8c4ceb3f..5aa36646 100644
--- a/view/sharedcache/ui/dsctriage.h
+++ b/view/sharedcache/ui/dsctriage.h
@@ -3,6 +3,7 @@
 //
 
 #include <binaryninjaapi.h>
+#include <QItemDelegate>
 #include <QStyledItemDelegate>
 
 #include "uitypes.h"
@@ -137,5 +138,58 @@ public:
 	static void Register();
 };
 
+class LoadedDelegate : public QItemDelegate
+{
+	Q_OBJECT
+public:
+	void paint(QPainter *painter, const QStyleOptionViewItem &option,
+			   const QModelIndex &index) const override
+	{
+		if (!index.isValid())
+			return;
+
+		painter->save();
+
+		// Highlight if the item is selected
+		if (option.state & QStyle::State_Selected)
+			painter->fillRect(option.rect, option.palette.highlight());
+
+		// "1" is the indicator that its loaded.
+		if (index.data(Qt::DisplayRole).toString() == "1")
+		{
+			QPixmap loadedIcon;
+			pixmapForBWMaskIcon(":/icons/images/download.png", &loadedIcon, SidebarHeaderTextColor);
+			if (!loadedIcon.isNull())
+			{
+				QSize pixmapSize(20, 20);
+				QPixmap scaledPixmap = loadedIcon.scaled(pixmapSize, Qt::KeepAspectRatio, Qt::SmoothTransformation);
+
+				// Calculate the rectangle for centering the pixmap
+				int x = option.rect.x() + (option.rect.width() - scaledPixmap.width()) / 2; // Center horizontally
+				int y = option.rect.y() + (option.rect.height() - scaledPixmap.height()) / 2; // Center vertically
+				QRect iconRect(x, y, scaledPixmap.width(), scaledPixmap.height());
+
+				// Draw the pixmap
+				painter->drawPixmap(iconRect, scaledPixmap);
+			}
+		}
+
+		painter->restore();
+	}
+
+	QSize sizeHint(const QStyleOptionViewItem &option,
+				   const QModelIndex &index) const override
+	{
+		Q_UNUSED(option);
+		Q_UNUSED(index);
+		return QSize(50, 24);
+	}
+
+	void setEditorData(QWidget *editor, const QModelIndex &index) const override
+	{
+		Q_UNUSED(editor);
+		Q_UNUSED(index);
+	}
+};
 
 #endif	// BINARYNINJA_DSCTRIAGE_H
diff --git a/view/sharedcache/ui/symboltable.cpp b/view/sharedcache/ui/symboltable.cpp
index c5e2d7bf..c7393d59 100644
--- a/view/sharedcache/ui/symboltable.cpp
+++ b/view/sharedcache/ui/symboltable.cpp
@@ -19,8 +19,8 @@ int SymbolTableModel::rowCount(const QModelIndex& parent) const {
 
 int SymbolTableModel::columnCount(const QModelIndex& parent) const {
 	Q_UNUSED(parent);
-	// We have 2 columns: Address, Name
-	return 2;
+	// We have 3 columns: Address, Type, Name
+	return 3;
 }
 
 QVariant SymbolTableModel::data(const QModelIndex& index, int role) const {
@@ -29,11 +29,14 @@ QVariant SymbolTableModel::data(const QModelIndex& index, int role) const {
 	}
 
 	const CacheSymbol& symbol = m_symbols.at(index.row());
+	auto symbolType = GetSymbolTypeAsString(symbol.type);
 
 	switch (index.column()) {
 	case 0: // Address column
 		return QString("0x%1").arg(symbol.address, 0, 16); // Display address as hexadecimal
-	case 1: // Name column
+	case 1: // Type column
+		return QString::fromUtf8(symbolType.c_str(), symbolType.size());
+	case 2: // Name column
 		return QString::fromUtf8(symbol.name.c_str(), symbol.name.size());
 	default:
 		return QVariant();
@@ -49,6 +52,8 @@ QVariant SymbolTableModel::headerData(int section, Qt::Orientation orientation,
 	case 0:
 		return QString("Address");
 	case 1:
+		return QString("Type");
+	case 2:
 		return QString("Name");
 	default:
 		return QVariant();
@@ -100,7 +105,9 @@ SymbolTableView::SymbolTableView(QWidget* parent)
 	setModel(m_model);
 
 	// Configure view settings
-	horizontalHeader()->setSectionResizeMode(QHeaderView::Stretch);
+	horizontalHeader()->setSectionResizeMode(0, QHeaderView::ResizeToContents);
+	horizontalHeader()->setSectionResizeMode(1, QHeaderView::ResizeToContents);
+	horizontalHeader()->setSectionResizeMode(2, QHeaderView::Stretch);
 	setSelectionBehavior(QAbstractItemView::SelectRows);
 	setSelectionMode(QAbstractItemView::SingleSelection);
 
@@ -116,6 +123,7 @@ void SymbolTableView::populateSymbols(BinaryView &view)
 	if (auto controller = SharedCacheController::GetController(view)) {
 		BackgroundThread::create(this)
 			->thenBackground([this, controller]() {
+				// TODO: There is a crash related to `this` being destructed. https://github.com/Vector35/binaryninja-api/issues/6300
 				std::vector<CacheSymbol> newSymbols = controller->GetSymbols();
 				m_symbols.swap(newSymbols);
 			})