[WARP] Improve UX surrounding removal of matched functions

- Adds Python API to remove matched function
- Adds command + UI actions to remove matched function
- Adds command + UI actions to ignore function in subsequent matches

4 files changed, 104 insertions, 1 deletions

diff --git a/plugins/warp/ui/shared/function.cpp b/plugins/warp/ui/shared/function.cpp
index e177c285..06d4cbcf 100644
--- a/plugins/warp/ui/shared/function.cpp
+++ b/plugins/warp/ui/shared/function.cpp
@@ -70,6 +70,15 @@ void WarpFunctionItemModel::InsertFunction(uint64_t address, WarpFunctionItem* i
 	m_insertableFunctionRows[address] = rowCount() - 1;
 }
 
+void WarpFunctionItemModel::RemoveFunction(uint64_t address)
+{
+	const auto iter = m_insertableFunctionRows.find(address);
+	if (iter == m_insertableFunctionRows.end())
+		return;
+	removeRow(iter->second);
+	m_insertableFunctionRows.erase(iter);
+}
+
 WarpFunctionItem* WarpFunctionItemModel::GetItem(const QModelIndex& index) const
 {
 	if (!index.isValid())
@@ -173,7 +182,7 @@ WarpFunctionTableWidget::WarpFunctionTableWidget(QWidget* parent) : QWidget(pare
 	m_table->setSelectionBehavior(QAbstractItemView::SelectRows);
 	m_table->setSelectionMode(QAbstractItemView::SingleSelection);
 	m_table->setEditTriggers(QAbstractItemView::NoEditTriggers);
-	m_table->setFocusPolicy(Qt::NoFocus);
+	m_table->setFocusPolicy(Qt::FocusPolicy::StrongFocus);
 	m_table->setShowGrid(false);
 	m_table->setAlternatingRowColors(false);
 	m_table->setSortingEnabled(true);
@@ -219,6 +228,15 @@ WarpFunctionTableWidget::WarpFunctionTableWidget(QWidget* parent) : QWidget(pare
 		if (!item || !item->GetFunction())
 			return;
 
+		for (QAction* action : m_contextMenu->actions())
+		{
+			bool enabled = true;
+			auto iter = m_contextMenuIsValid.find(action->text());
+			if (iter != m_contextMenuIsValid.end())
+				enabled = iter->second(item, m_model->GetAddress(sourceIndex));
+			action->setEnabled(enabled);
+		}
+
 		// Execute the menu and get the selected action
 		const QAction* selectedAction = m_contextMenu->exec(m_table->viewport()->mapToGlobal(pos));
 		if (!selectedAction)
@@ -238,6 +256,16 @@ void WarpFunctionTableWidget::RegisterContextMenuAction(
 	m_contextMenuActions[name] = callback;
 }
 
+void WarpFunctionTableWidget::RegisterContextMenuAction(
+	const QString& name,
+	const std::function<void(WarpFunctionItem*, std::optional<uint64_t>)>& callback,
+	const std::function<bool(WarpFunctionItem*, std::optional<uint64_t>)>& isValid)
+{
+	// Reuse existing registration then add optional validator
+	RegisterContextMenuAction(name, callback);
+	m_contextMenuIsValid[name] = isValid;
+}
+
 void WarpFunctionTableWidget::SetFunctions(QVector<WarpFunctionItem*> functions)
 {
 	// Clear matches as they are no longer valid.
@@ -266,6 +294,11 @@ void WarpFunctionTableWidget::InsertFunction(uint64_t address, WarpFunctionItem*
 	m_model->InsertFunction(address, function);
 }
 
+void WarpFunctionTableWidget::RemoveFunction(uint64_t address)
+{
+	m_model->RemoveFunction(address);
+}
+
 void WarpFunctionTableWidget::setFilter(const std::string& filter)
 {
 	m_proxyModel->setFilterFixedString(QString::fromStdString(filter));
diff --git a/plugins/warp/ui/shared/function.h b/plugins/warp/ui/shared/function.h
index 5cfaff75..6281944d 100644
--- a/plugins/warp/ui/shared/function.h
+++ b/plugins/warp/ui/shared/function.h
@@ -49,6 +49,8 @@ public:
 
 	void InsertFunction(uint64_t address, WarpFunctionItem* item);
 
+	void RemoveFunction(uint64_t address);
+
 	WarpFunctionItem* GetItem(const QModelIndex& index) const;
 
 	std::optional<uint64_t> GetAddress(const QModelIndex& index) const;
@@ -68,6 +70,8 @@ public:
 			emit dataChanged(topLeft, bottomRight);
 		}
 	}
+
+	Warp::Ref<Warp::Function> GetMatchedFunction() const { return m_matchedFunction; }
 };
 
 class WarpFunctionFilterModel : public QSortFilterProxyModel
@@ -95,6 +99,7 @@ class WarpFunctionTableWidget : public QWidget, public FilterTarget
 	FilteredView* m_filterView;
 	QMenu* m_contextMenu;
 	std::map<QString, std::function<void(WarpFunctionItem*, std::optional<uint64_t>)>> m_contextMenuActions;
+	std::map<QString, std::function<bool(WarpFunctionItem*, std::optional<uint64_t>)>> m_contextMenuIsValid;
 
 public:
 	explicit WarpFunctionTableWidget(QWidget* parent = nullptr);
@@ -107,10 +112,17 @@ public:
 	void RegisterContextMenuAction(
 		const QString& name, const std::function<void(WarpFunctionItem*, std::optional<uint64_t>)>& callback);
 
+	void RegisterContextMenuAction(
+		const QString &name,
+		const std::function<void(WarpFunctionItem *, std::optional<uint64_t>)> &callback,
+		const std::function<bool(WarpFunctionItem *, std::optional<uint64_t>)> &isValid);
+
 	void SetFunctions(QVector<WarpFunctionItem*> functions);
 
 	void InsertFunction(uint64_t address, WarpFunctionItem* function);
 
+	void RemoveFunction(uint64_t address);
+
 	void setFilter(const std::string&) override;
 
 	void scrollToFirstItem() override {}
diff --git a/plugins/warp/ui/shared/misc.cpp b/plugins/warp/ui/shared/misc.cpp
index f84237c9..6de314b1 100644
--- a/plugins/warp/ui/shared/misc.cpp
+++ b/plugins/warp/ui/shared/misc.cpp
@@ -1,5 +1,6 @@
 #include "misc.h"
 
+#include <QDialogButtonBox>
 #include <QGridLayout>
 #include <QHeaderView>
 
@@ -140,3 +141,44 @@ ParsedQuery::ParsedQuery(const QString& rawQuery)
 	// Normalize whitespace
 	query = query.simplified();
 }
+
+WarpRemoveMatchDialog::WarpRemoveMatchDialog(QWidget *parent, FunctionRef func) : QDialog(parent), m_func(func)
+{
+	setWindowTitle("Remove Matching Function");
+	setModal(true);
+
+	auto* vbox = new QVBoxLayout(this);
+	auto* text = new QLabel("Remove the match for this function? You can also mark it as ignored to prevent future automatic matches.");
+	text->setWordWrap(true);
+	vbox->addWidget(text);
+
+	m_ignoreCheck = new QCheckBox("Tag function as ignored");
+	m_ignoreCheck->setChecked(true);
+	vbox->addWidget(m_ignoreCheck);
+
+	auto* buttons = new QDialogButtonBox(QDialogButtonBox::Ok | QDialogButtonBox::Cancel, this);
+	connect(buttons, &QDialogButtonBox::accepted, this, &QDialog::accept);
+	connect(buttons, &QDialogButtonBox::rejected, this, &QDialog::reject);
+	vbox->addWidget(buttons);
+}
+
+bool WarpRemoveMatchDialog::execute()
+{
+	if (!m_func)
+		return false;
+	if (exec() != QDialog::Accepted)
+		return false;
+	Warp::Function::RemoveMatch(*m_func);
+	if (m_ignoreCheck->isChecked())
+	{
+		// TODO: For now we just assume the tag type to exist (the matcher activity will create it)
+		const TagTypeRef tagType = m_func->GetView()->GetTagTypeByName("WARP: Ignored Function");
+		if (!tagType)
+			return false;
+		const TagRef tag = new BinaryNinja::Tag(tagType, "");
+		if (tagType)
+			m_func->AddUserFunctionTag(tag);
+	}
+	m_func->Reanalyze();
+	return true;
+}
diff --git a/plugins/warp/ui/shared/misc.h b/plugins/warp/ui/shared/misc.h
index 9d187594..97adcab0 100644
--- a/plugins/warp/ui/shared/misc.h
+++ b/plugins/warp/ui/shared/misc.h
@@ -1,4 +1,6 @@
 #pragma once
+#include <QCheckBox>
+#include <QDialog>
 #include <qmetatype.h>
 #include <QSortFilterProxyModel>
 #include <qstandarditemmodel.h>
@@ -103,6 +105,20 @@ struct ParsedQuery
 	}
 };
 
+// TODO: Consolidate with `WARP\\Remove Matched Function` plugin command?
+class WarpRemoveMatchDialog : public QDialog
+{
+	Q_OBJECT
+public:
+	explicit WarpRemoveMatchDialog(QWidget* parent, FunctionRef func);
+
+	bool execute();
+
+private:
+	FunctionRef m_func;
+	QCheckBox* m_ignoreCheck{nullptr};
+};
+
 constexpr const char* ALLOWED_TAGS_SETTING = "warp.fetcher.allowedSourceTags";
 constexpr const char* BATCH_SIZE_SETTING = "warp.fetcher.fetchBatchSize";