Fixes for multiple issues in DSC/KC Triage views.

1. Crash on kernel cache image load
2. Duplicate Load Image buttons in kernel cache view
3. Improper selection behavior in both views
4. Address column should no longer resize to be smaller than contents
5. Symbol tables should now be properly sortable

4 files changed, 194 insertions, 126 deletions

diff --git a/view/sharedcache/ui/dsctriage.cpp b/view/sharedcache/ui/dsctriage.cpp
index e2cf8147..54fe4ed0 100644
--- a/view/sharedcache/ui/dsctriage.cpp
+++ b/view/sharedcache/ui/dsctriage.cpp
@@ -1,11 +1,10 @@
+#include <QHeaderView>
+#include <QMessageBox>
+#include <utility>
 #include "dsctriage.h"
 #include "globalarea.h"
-#include "progresstask.h"
 #include "symboltable.h"
 #include "ui/fontsettings.h"
-#include <QHeaderView>
-#include <QMessageBox>
-#include <utility>
 
 using namespace BinaryNinja;
 using namespace SharedCacheAPI;
@@ -28,7 +27,6 @@ QWidget* DSCTriageViewType::create(BinaryViewRef data, ViewFrame* viewFrame)
 {
 	if (data->GetTypeName() != VIEW_NAME)
 		return nullptr;
-	// TODO: Check for dyld start. Then continue.
 	return new DSCTriageView(viewFrame, data);
 }
 
@@ -134,36 +132,36 @@ void DSCTriageView::loadImagesWithAddr(const std::vector<uint64_t>& addresses, b
 
 QWidget* DSCTriageView::initImageTable()
 {
-	auto loadImageTable = new FilterableTableView(this);
+	m_imageTable = new FilterableTableView(this);
 
-	m_imageModel = new QStandardItemModel(0, 3, loadImageTable);
+	m_imageModel = new QStandardItemModel(0, 3, m_imageTable);
 	m_imageModel->setHorizontalHeaderLabels({"Address", "Loaded", "Name"});
 
 	// Apply custom column styling
-	loadImageTable->setItemDelegateForColumn(0, new AddressColorDelegate(loadImageTable));
-	loadImageTable->setItemDelegateForColumn(1, new LoadedDelegate(loadImageTable));
+	m_imageTable->setItemDelegateForColumn(0, new AddressColorDelegate(m_imageTable));
+	m_imageTable->setItemDelegateForColumn(1, new LoadedDelegate(m_imageTable));
 
 	// Context menu
-	loadImageTable->setContextMenuPolicy(Qt::CustomContextMenu);
-	connect(loadImageTable, &QWidget::customContextMenuRequested, [this, loadImageTable](const QPoint &pos) {
-		QMenu contextMenu(tr("Load Image Actions"), loadImageTable);
+	m_imageTable->setContextMenuPolicy(Qt::CustomContextMenu);
+	connect(m_imageTable, &QWidget::customContextMenuRequested, [this](const QPoint &pos) {
+		QMenu contextMenu(tr("Load Image Actions"), m_imageTable);
 
 		// Get number of selected images
-		auto selected = loadImageTable->selectionModel()->selectedRows();
+		auto selected = m_imageTable->selectionModel()->selectedRows();
 		int selectedCount = 0;
 		std::vector<uint64_t> addresses;
 		for (const auto& idx : selected)
 		{
 			// Skip rows hidden by the filter
-			if (loadImageTable->isRowHidden(idx.row()))
+			if (m_imageTable->isRowHidden(idx.row()))
 				continue;
 			addresses.push_back(idx.data().toString().toULongLong(nullptr, 16));
 			selectedCount++;
 		}
 
-		QAction noSelectionAction("No Images Selected", loadImageTable);
-		QAction loadImagesAction("", loadImageTable);
-		QAction loadImagesWithDepsAction("", loadImageTable);
+		QAction noSelectionAction("No Images Selected", m_imageTable);
+		QAction loadImagesAction("", m_imageTable);
+		QAction loadImagesWithDepsAction("", m_imageTable);
 		if (selectedCount == 0)
 		{
 			noSelectionAction.setEnabled(false);
@@ -188,20 +186,25 @@ QWidget* DSCTriageView::initImageTable()
 			contextMenu.addAction(&loadImagesWithDepsAction);
 		}
 
-		contextMenu.exec(loadImageTable->viewport()->mapToGlobal(pos));
+		contextMenu.exec(m_imageTable->viewport()->mapToGlobal(pos));
 	});
 
 	auto loadImageButton = new QPushButton();
-	connect(loadImageButton, &QPushButton::clicked, [this, loadImageTable](bool) {
-		auto selected = loadImageTable->selectionModel()->selectedRows();
+	connect(loadImageButton, &QPushButton::clicked, [this](bool) {
+		// Collect only visible selected rows
+		QModelIndexList selected;
+		for (const auto& index : m_imageTable->selectionModel()->selectedRows()) {
+			if (!m_imageTable->isRowHidden(index.row())) {
+				selected.append(index);
+			}
+		}
+
+		if (selected.empty())
+			return;
+
 		std::vector<uint64_t> addresses;
 		for (const auto& idx : selected)
-		{
-			// Skip rows hidden by the filter
-			if (loadImageTable->isRowHidden(idx.row()))
-				continue;
 			addresses.push_back(idx.data().toString().toULongLong(nullptr, 16));
-		}
 		loadImagesWithAddr(addresses);
 	});
 	loadImageButton->setText("Load Selected");
@@ -213,43 +216,43 @@ QWidget* DSCTriageView::initImageTable()
 		refreshDataButton->setText("Refresh");
 	} // refreshDataButton
 
-	auto loadImageFilterEdit = new FilterEdit(loadImageTable);
-	connect(loadImageFilterEdit, &FilterEdit::textChanged, [loadImageTable](const QString& filter) {
-		loadImageTable->setFilter(filter.toStdString());
+	auto loadImageFilterEdit = new FilterEdit(m_imageTable);
+	connect(loadImageFilterEdit, &FilterEdit::textChanged, [this](const QString& filter) {
+		m_imageTable->setFilter(filter.toStdString());
 	});
 
-	connect(loadImageTable, &FilterableTableView::activated, this, [=](const QModelIndex& index) {
+	connect(m_imageTable, &FilterableTableView::activated, this, [=](const QModelIndex& index) {
 		auto addr = m_imageModel->item(index.row(), 0)->text().toULongLong(nullptr, 16);
 		loadImagesWithAddr({addr});
 	});
 
 	auto loadImageLayout = new QVBoxLayout;
 	loadImageLayout->addWidget(loadImageFilterEdit);
-	loadImageLayout->addWidget(loadImageTable);
+	loadImageLayout->addWidget(m_imageTable);
 
-	auto buttonLayout = new QHBoxLayout;
-	buttonLayout->addWidget(loadImageButton);
-	buttonLayout->addWidget(refreshDataButton);
-	buttonLayout->setAlignment(Qt::AlignLeft);
-	loadImageLayout->addLayout(buttonLayout);
+	auto loadImageFooterLayout = new QHBoxLayout;
+	loadImageFooterLayout->addWidget(loadImageButton);
+	loadImageFooterLayout->addWidget(refreshDataButton);
+	loadImageFooterLayout->setAlignment(Qt::AlignLeft);
+	loadImageLayout->addLayout(loadImageFooterLayout);
 
 	auto loadImageWidget = new QWidget;
 	loadImageWidget->setLayout(loadImageLayout);
 
-	loadImageTable->setModel(m_imageModel);
+	m_imageTable->setModel(m_imageModel);
 
-	loadImageTable->setEditTriggers(QAbstractItemView::NoEditTriggers);
+	m_imageTable->setEditTriggers(QAbstractItemView::NoEditTriggers);
 
-	loadImageTable->horizontalHeader()->setSectionResizeMode(0, QHeaderView::Fixed);
-	loadImageTable->horizontalHeader()->setSectionResizeMode(1, QHeaderView::ResizeToContents);
-	loadImageTable->horizontalHeader()->setSectionResizeMode(2, QHeaderView::Stretch);
+	m_imageTable->horizontalHeader()->setSectionResizeMode(0, QHeaderView::Fixed);
+	m_imageTable->horizontalHeader()->setSectionResizeMode(1, QHeaderView::ResizeToContents);
+	m_imageTable->horizontalHeader()->setSectionResizeMode(2, QHeaderView::Stretch);
 
-	loadImageTable->setSelectionBehavior(QAbstractItemView::SelectRows);
-	loadImageTable->setSelectionMode(QAbstractItemView::ExtendedSelection);
+	m_imageTable->setSelectionBehavior(QAbstractItemView::SelectRows);
+	m_imageTable->setSelectionMode(QAbstractItemView::ExtendedSelection);
 
-	loadImageTable->setSortingEnabled(true);
+	m_imageTable->setSortingEnabled(true);
 
-	loadImageTable->verticalHeader()->setVisible(false);
+	m_imageTable->verticalHeader()->setVisible(false);
 
 	m_triageTabs->addTab(loadImageWidget, "Images");
 	m_triageTabs->setCanCloseTab(loadImageWidget, false);
@@ -262,51 +265,43 @@ void DSCTriageView::initSymbolTable()
 {
 	m_symbolTable = new SymbolTableView(this);
 
+	// Apply custom column styling
+	m_symbolTable->setItemDelegateForColumn(0, new AddressColorDelegate(m_symbolTable));
+
 	auto symbolFilterEdit = new FilterEdit(m_symbolTable);
 	connect(symbolFilterEdit, &FilterEdit::textChanged, [this](const QString& filter) {
 		m_symbolTable->setFilter(filter.toStdString());
 	});
 
-	// Apply custom column styling
-	m_symbolTable->setItemDelegateForColumn(0, new AddressColorDelegate(m_symbolTable));
-
 	auto loadSymbolImageButton = new QPushButton();
-	{
-		connect(loadSymbolImageButton, &QPushButton::clicked,
-				[this](bool) {
-					auto selected = m_symbolTable->selectionModel()->selectedRows();
-					std::vector<uint64_t> addresses;
-					for (const auto& row : selected)
-						addresses.push_back(row.data().toString().toULongLong(nullptr, 16));
-					loadImagesWithAddr(addresses);
-				});
-		loadSymbolImageButton->setText("Load Image");
-	} // loadImageButton
+	connect(loadSymbolImageButton, &QPushButton::clicked, [this](bool) {
+		auto selected = m_symbolTable->selectionModel()->selectedRows();
+		std::vector<uint64_t> addresses;
+		for (const auto& row : selected)
+			addresses.push_back(row.data().toString().toULongLong(nullptr, 16));
+		loadImagesWithAddr(addresses);
+	});
+	loadSymbolImageButton->setText("Load Image");
 
 	// Shows the current selected rows image name.
-	auto currentImageLabel = new QLabel(this); {
-		currentImageLabel->setText("");
-		currentImageLabel->setSizePolicy(QSizePolicy::Fixed, QSizePolicy::Fixed);
-	}
-
-	// Update the label whenever the selection changes.
-	connect(m_symbolTable->selectionModel(), &QItemSelectionModel::currentRowChanged, this,
-			[this, currentImageLabel](const QModelIndex &current, const QModelIndex &) {
-				auto symbol = m_symbolTable->getSymbolAtRow(current.row());
-				auto controller = SharedCacheController::GetController(*this->m_data);
-				if (!controller)
-					return;
-				auto image = controller->GetImageContaining(symbol.address);
-				if (image)
-					currentImageLabel->setText("Image: " + QString::fromStdString(image->name));
-				else
-					currentImageLabel->setText("");
-			});
+	auto currentImageLabel = new QLabel(this);
+	currentImageLabel->setText("");
+	currentImageLabel->setSizePolicy(QSizePolicy::Fixed, QSizePolicy::Fixed);
+	connect(m_symbolTable->selectionModel(), &QItemSelectionModel::currentRowChanged, this, [this, currentImageLabel](const QModelIndex &current, const QModelIndex &) {
+		auto symbol = m_symbolTable->getSymbolAtRow(current.row());
+		auto controller = SharedCacheController::GetController(*this->m_data);
+		if (!controller)
+			return;
+		auto image = controller->GetImageContaining(symbol.address);
+		if (image)
+			currentImageLabel->setText("Image: " + QString::fromStdString(image->name));
+		else
+			currentImageLabel->setText("");
+	});
 
 	auto symbolFooterLayout = new QHBoxLayout;
 	symbolFooterLayout->addWidget(loadSymbolImageButton);
 	symbolFooterLayout->addWidget(currentImageLabel);
-
 	symbolFooterLayout->setAlignment(Qt::AlignLeft);
 
 	auto symbolLayout = new QVBoxLayout;
diff --git a/view/sharedcache/ui/dsctriage.h b/view/sharedcache/ui/dsctriage.h
index a16457e3..e847625e 100644
--- a/view/sharedcache/ui/dsctriage.h
+++ b/view/sharedcache/ui/dsctriage.h
@@ -1,24 +1,29 @@
-#include <binaryninjaapi.h>
+#include <QHeaderView>
 #include <QItemDelegate>
+#include <QPainter>
+#include <QSortFilterProxyModel>
+#include <QStandardItemModel>
 #include <QStyledItemDelegate>
-
-#include "uitypes.h"
-#include "viewframe.h"
-#include "animation.h"
-#include "uicontext.h"
-
+#include <QTableView>
+#include <binaryninjaapi.h>
+#include <progresstask.h>
+#include <sharedcacheapi.h>
 #include "filter.h"
 #include "symboltable.h"
-
 #include "ui/fontsettings.h"
+#include "uicontext.h"
+#include "uitypes.h"
+#include "viewframe.h"
 
 #ifndef BINARYNINJA_DSCTRIAGE_H
 #define BINARYNINJA_DSCTRIAGE_H
 
+using namespace BinaryNinja;
+using namespace SharedCacheAPI;
+
 
 class AddressColorDelegate : public QStyledItemDelegate
 {
-
 public:
 	explicit AddressColorDelegate(QObject* parent = nullptr) : QStyledItemDelegate(parent) {}
 
@@ -27,6 +32,7 @@ public:
 		QStyleOptionViewItem opt = option;
 		initStyleOption(&opt, index);
 
+		opt.font = getMonospaceFont(qobject_cast<QWidget*>(parent()));
 		opt.palette.setColor(QPalette::Text, getThemeColor(BNThemeColor::AddressColor));
 		opt.displayAlignment = Qt::AlignCenter | Qt::AlignVCenter;
 
@@ -130,7 +136,10 @@ public:
 
 	void scrollToFirstItem() override {
 		if (model()->rowCount() > 0) {
-			scrollTo(model()->index(0, 0));
+			QModelIndex top = indexAt(rect().topLeft());
+			if (top.isValid()) {
+				scrollTo(top);
+			}
 		}
 	}
 
@@ -143,34 +152,39 @@ public:
 
 	void selectFirstItem() override {
 		if (model()->rowCount() > 0) {
-			QModelIndex firstIndex = model()->index(0, 0);
-			selectionModel()->select(firstIndex, QItemSelectionModel::ClearAndSelect);
+			QModelIndex top = indexAt(rect().topLeft());
+			if (top.isValid()) {
+				selectionModel()->select(top, QItemSelectionModel::ClearAndSelect);
+				setCurrentIndex(top);
+			}
 		}
 	}
 
 	void activateFirstItem() override {
 		if (model()->rowCount() > 0) {
-			QModelIndex firstIndex = model()->index(0, 0);
-			setCurrentIndex(firstIndex);
-			emit activated(firstIndex);
-		}
-	}
-
-	bool eventFilter(QObject* obj, QEvent* event) override {
-		if (event->type() == QEvent::KeyPress) {
-			auto* keyEvent = dynamic_cast<QKeyEvent*>(event);
-			if (keyEvent->key() == Qt::Key_Escape) {
-				clearSelection();
-				return true;
-			}
-			if (keyEvent->key() == Qt::Key_Enter || keyEvent->key() == Qt::Key_Return) {
-				emit activated(currentIndex());
-				return true;
+			QModelIndex topLeft = indexAt(rect().topLeft());
+			if (topLeft.isValid()) {
+				setCurrentIndex(topLeft);
+				emit activated(topLeft);
 			}
 		}
-		return QTableView::eventFilter(obj, event);
 	}
 
+//	bool eventFilter(QObject* obj, QEvent* event) override {
+//		if (event->type() == QEvent::KeyPress) {
+//			auto* keyEvent = dynamic_cast<QKeyEvent*>(event);
+//			if (keyEvent->key() == Qt::Key_Escape) {
+//				clearSelection();
+//				return true;
+//			}
+//			if (keyEvent->key() == Qt::Key_Enter || keyEvent->key() == Qt::Key_Return) {
+//				emit activated(currentIndex());
+//				return true;
+//			}
+//		}
+//		return QTableView::eventFilter(obj, event);
+//	}
+
 signals:
 	void filterTextChanged(const QString& text);
 };
@@ -184,6 +198,7 @@ class DSCTriageView : public QWidget, public View, public UIContextNotification
 	SplitTabWidget* m_triageTabs;
 	DockableTabCollection* m_triageCollection;
 
+	FilterableTableView* m_imageTable;
 	QStandardItemModel* m_imageModel;
 
 	SymbolTableView* m_symbolTable;
diff --git a/view/sharedcache/ui/symboltable.cpp b/view/sharedcache/ui/symboltable.cpp
index 1df7d529..b0247f2b 100644
--- a/view/sharedcache/ui/symboltable.cpp
+++ b/view/sharedcache/ui/symboltable.cpp
@@ -10,23 +10,27 @@
 using namespace BinaryNinja;
 using namespace SharedCacheAPI;
 
+
 SymbolTableModel::SymbolTableModel(SymbolTableView* parent)
 	: QAbstractTableModel(parent), m_parent(parent) {
 	// TODO: Need to implement updating this font if it is changed by the user
 	m_font = getMonospaceFont(parent);
 }
 
+
 int SymbolTableModel::rowCount(const QModelIndex& parent) const {
 	Q_UNUSED(parent);
 	return static_cast<int>(m_modelSymbols.size());
 }
 
+
 int SymbolTableModel::columnCount(const QModelIndex& parent) const {
 	Q_UNUSED(parent);
 	// We have 3 columns: Address, Type, Name
 	return 3;
 }
 
+
 QVariant SymbolTableModel::data(const QModelIndex& index, int role) const {
 	if (!index.isValid() || (role != Qt::DisplayRole && role != Qt::FontRole)) {
 		return QVariant();
@@ -58,6 +62,7 @@ QVariant SymbolTableModel::data(const QModelIndex& index, int role) const {
 	}
 }
 
+
 QVariant SymbolTableModel::headerData(int section, Qt::Orientation orientation, int role) const {
 	if (role != Qt::DisplayRole || orientation != Qt::Horizontal) {
 		return QVariant();
@@ -75,12 +80,58 @@ QVariant SymbolTableModel::headerData(int section, Qt::Orientation orientation,
 	}
 }
 
+
+void SymbolTableModel::sort(int column, Qt::SortOrder order)
+{
+	beginResetModel();
+
+	std::function<bool(const CacheSymbol&, const CacheSymbol&)> comparator;
+
+	switch (column)
+	{
+	case 0: // Address column
+		comparator = [](const CacheSymbol& a, const CacheSymbol& b) {
+			return a.address < b.address;
+		};
+		break;
+	case 1: // Type column
+		comparator = [](const CacheSymbol& a, const CacheSymbol& b) {
+			return GetSymbolTypeAsString(a.type) < GetSymbolTypeAsString(b.type);
+		};
+		break;
+	case 2: // Name column
+		comparator = [](const CacheSymbol& a, const CacheSymbol& b) {
+			return a.name < b.name;
+		};
+		break;
+	default:
+		endResetModel();
+		return;
+	}
+
+	if (order == Qt::DescendingOrder)
+	{
+		std::sort(m_modelSymbols.begin(), m_modelSymbols.end(),
+				  [&comparator](const CacheSymbol& a, const CacheSymbol& b) {
+					  return comparator(b, a);
+				  });
+	}
+	else
+	{
+		std::sort(m_modelSymbols.begin(), m_modelSymbols.end(), comparator);
+	}
+
+	endResetModel();
+}
+
+
 void SymbolTableModel::updateSymbols(std::vector<CacheSymbol>&& symbols)
 {
 	m_preparedSymbols = symbols;
 	setFilter(m_filter);
 }
 
+
 const CacheSymbol& SymbolTableModel::symbolAt(int row) const
 {
 	return m_modelSymbols.at(row);
diff --git a/view/sharedcache/ui/symboltable.h b/view/sharedcache/ui/symboltable.h
index 2fc5c892..51ba88f0 100644
--- a/view/sharedcache/ui/symboltable.h
+++ b/view/sharedcache/ui/symboltable.h
@@ -2,12 +2,13 @@
 
 #include <sharedcacheapi.h>
 #include "viewframe.h"
-#include "animation.h"
 
 #include <QTableView>
 #include <QStandardItemModel>
 #include "filter.h"
 
+#ifndef BINARYNINJA_DSCSYMBOLTABLE_H
+#define BINARYNINJA_DSCSYMBOLTABLE_H
 
 class SymbolTableView;
 
@@ -29,9 +30,11 @@ public:
 	int columnCount(const QModelIndex& parent) const override;
 	QVariant data(const QModelIndex& index, int role) const override;
 	QVariant headerData(int section, Qt::Orientation orientation, int role) const override;
+	void sort(int column, Qt::SortOrder order) override;
 	void updateSymbols(std::vector<SharedCacheAPI::CacheSymbol>&& symbols);
 	void setFilter(std::string text);
 	const SharedCacheAPI::CacheSymbol& symbolAt(int row) const;
+
 };
 
 
@@ -44,45 +47,46 @@ Q_OBJECT
 
 public:
 	explicit SymbolTableView(QWidget* parent);
-
 	~SymbolTableView() override;
 
+	// Call this to populate the symbols from the given view.
+	void populateSymbols(BinaryNinja::BinaryView& view);
+
 	void scrollToFirstItem() override
 	{
-		if (model()->rowCount() > 0)
-		{
-			scrollTo(model()->index(0, 0));
+		if (model()->rowCount() > 0) {
+			QModelIndex top = indexAt(rect().topLeft());
+			if (top.isValid())
+				scrollTo(top);
 		}
 	}
 
-	// Call this to populate the symbols from the given view.
-	void populateSymbols(BinaryNinja::BinaryView& view);
-
 	void scrollToCurrentItem() override
 	{
 		QModelIndex currentIndex = selectionModel()->currentIndex();
 		if (currentIndex.isValid())
-		{
 			scrollTo(currentIndex);
-		}
 	}
 
 	void selectFirstItem() override
 	{
-		if (model()->rowCount() > 0)
-		{
-			QModelIndex firstIndex = model()->index(0, 0);
-			selectionModel()->select(firstIndex, QItemSelectionModel::ClearAndSelect);
+		if (model()->rowCount() > 0) {
+			QModelIndex top = indexAt(rect().topLeft());
+			if (top.isValid()) {
+				selectionModel()->select(top, QItemSelectionModel::ClearAndSelect);
+				setCurrentIndex(top);
+			}
 		}
 	}
 
 	void activateFirstItem() override
 	{
-		if (model()->rowCount() > 0)
-		{
-			QModelIndex firstIndex = model()->index(0, 0);
-			setCurrentIndex(firstIndex);
-			emit activated(firstIndex);
+		if (model()->rowCount() > 0) {
+			QModelIndex topLeft = indexAt(rect().topLeft());
+			if (topLeft.isValid()) {
+				setCurrentIndex(topLeft);
+				emit activated(topLeft);
+			}
 		}
 	}
 
@@ -92,4 +96,7 @@ public:
 	}
 
 	void setFilter(const std::string& filter) override;
-};
\ No newline at end of file
+};
+
+
+#endif  // BINARYNINJA_DSCSYMBOLTABLE_H