[WARP] Server-side constraint matching

Reduce networked functions by constraining on the returned set of functions on the server

4 files changed, 35 insertions, 34 deletions

diff --git a/plugins/warp/ui/shared/fetchdialog.cpp b/plugins/warp/ui/shared/fetchdialog.cpp
index 328f00c2..44d7c237 100644
--- a/plugins/warp/ui/shared/fetchdialog.cpp
+++ b/plugins/warp/ui/shared/fetchdialog.cpp
@@ -67,12 +67,6 @@ WarpFetchDialog::WarpFetchDialog(BinaryViewRef bv, std::shared_ptr<WarpFetcher>
 	for (const auto& t : GetAllowedTagsFromView(m_bv))
 		AddListItem(m_tagsList, QString::fromStdString(t));
 
-	// Batch size and matcher checkbox
-	m_batchSize = new QSpinBox(this);
-	m_batchSize->setRange(10, 1000);
-	m_batchSize->setValue(GetBatchSizeFromView(m_bv));
-	m_batchSize->setToolTip("Number of functions to fetch in each batch");
-
 	m_rerunMatcher = new QCheckBox("Re-run matcher after fetch", this);
 	m_rerunMatcher->setChecked(true);
 
@@ -83,7 +77,6 @@ WarpFetchDialog::WarpFetchDialog(BinaryViewRef bv, std::shared_ptr<WarpFetcher>
 
 	form->addRow(new QLabel("Container: "), m_containerCombo);
 	form->addRow(new QLabel("Allowed Tags: "), tagWrapper);
-	form->addRow(new QLabel("Batch Size: "), m_batchSize);
 	form->addRow(m_rerunMatcher);
 	form->addRow(m_clearProcessed);
 
@@ -144,7 +137,6 @@ void WarpFetchDialog::onAccept()
 	if (idx > 0)  // 0 == All Containers
 		containerIndex = static_cast<size_t>(idx - 1);
 
-	const auto batch = static_cast<size_t>(m_batchSize->value());
 	const bool rerun = m_rerunMatcher->isChecked();
 
 	const auto tags = collectTags();
@@ -155,7 +147,7 @@ void WarpFetchDialog::onAccept()
 		m_fetchProcessor->ClearProcessed();
 
 	// Execute the network fetch in batches
-	runBatchedFetch(containerIndex, tags, batch, rerun);
+	runBatchedFetch(containerIndex, tags, rerun);
 
 	accept();
 }
@@ -169,7 +161,7 @@ void WarpFetchDialog::onReject()
 }
 
 void WarpFetchDialog::runBatchedFetch(const std::optional<size_t>& containerIndex,
-	const std::vector<Warp::SourceTag>& allowedTags, size_t batchSize, bool rerunMatcher)
+	const std::vector<Warp::SourceTag>& allowedTags, bool rerunMatcher)
 {
 	if (!m_bv)
 		return;
@@ -177,42 +169,34 @@ void WarpFetchDialog::runBatchedFetch(const std::optional<size_t>& containerInde
 	std::vector<Ref<Function>> funcs = m_bv->GetAnalysisFunctionList();
 	if (funcs.empty())
 		return;
-	const size_t totalFuncs = funcs.size();
-	const size_t totalBatches = (totalFuncs + batchSize - 1) / batchSize;
 
 	// Create a background task to show progress in the UI
 	Ref<BackgroundTask> task =
-		new BackgroundTask("Fetching WARP functions (0 / " + std::to_string(totalBatches) + ")", false);
+		new BackgroundTask("Fetching WARP functions (0 / " + std::to_string(funcs.size()) + ")", true);
 
 	auto fetcher = m_fetchProcessor;
 	auto bv = m_bv;
 
 	// TODO: Too many captures in this thing lol.
 	WorkerInteractiveEnqueue(
-		[fetcher, bv, funcs = std::move(funcs), batchSize, rerunMatcher, task, allowedTags]() mutable {
+		[fetcher, bv, funcs = std::move(funcs), rerunMatcher, task, allowedTags]() mutable {
+			const auto batchSize = GetBatchSizeFromView(bv);
 			size_t processed = 0;
-			size_t batchIndex = 0;
-
 			while (processed < funcs.size())
 			{
+				if (task->IsCancelled())
+					break;
 				const size_t remaining = funcs.size() - processed;
 				const size_t thisBatchCount = std::min(batchSize, remaining);
-
 				for (size_t i = 0; i < thisBatchCount; ++i)
 					fetcher->AddPendingFunction(funcs[processed + i]);
-
 				fetcher->FetchPendingFunctions(allowedTags);
-
-				++batchIndex;
 				processed += thisBatchCount;
-
-				task->SetProgressText("Fetching WARP functions (" + std::to_string(batchIndex) + " / "
-					+ std::to_string((funcs.size() + batchSize - 1) / batchSize) + ")");
+				task->SetProgressText("Fetching WARP functions (" + std::to_string(processed) + " / " + std::to_string(funcs.size()) + ")");
 			}
 
 			task->Finish();
-			// TODO: Print how long it took?
-			Logger("WARP Fetcher").LogInfo("Finished fetching WARP functions...");
+			Logger("WARP Fetcher").LogInfo("Finished fetching WARP functions in %d seconds...", task->GetRuntimeSeconds());
 
 			if (rerunMatcher && bv)
 				Warp::RunMatcher(*bv);
diff --git a/plugins/warp/ui/shared/fetchdialog.h b/plugins/warp/ui/shared/fetchdialog.h
index c642d31e..72b8c57e 100644
--- a/plugins/warp/ui/shared/fetchdialog.h
+++ b/plugins/warp/ui/shared/fetchdialog.h
@@ -22,7 +22,6 @@ class WarpFetchDialog : public QDialog
 	QPushButton* m_removeTagBtn;
 	QPushButton* m_resetTagBtn;
 
-	QSpinBox* m_batchSize;
 	QCheckBox* m_rerunMatcher;
 	QCheckBox* m_clearProcessed;
 
@@ -51,7 +50,7 @@ private:
 	std::vector<Warp::SourceTag> collectTags() const;
 
 	void runBatchedFetch(const std::optional<size_t>& containerIndex, const std::vector<Warp::SourceTag>& allowedTags,
-		size_t batchSize, bool rerunMatcher);
+		bool rerunMatcher);
 };
 
 void RegisterWarpFetchFunctionsCommand();
diff --git a/plugins/warp/ui/shared/fetcher.cpp b/plugins/warp/ui/shared/fetcher.cpp
index 51ab018c..2eab910a 100644
--- a/plugins/warp/ui/shared/fetcher.cpp
+++ b/plugins/warp/ui/shared/fetcher.cpp
@@ -63,22 +63,40 @@ void WarpFetcher::FetchPendingFunctions(const std::vector<Warp::SourceTag>& allo
 
 	// Because we must fetch for a single target we map the function guids to the associated platform to perform fetches
 	// for each.
-	std::map<PlatformRef, std::vector<Warp::FunctionGUID>> platformMappedGuids;
+	std::map<PlatformRef, std::unordered_set<Warp::FunctionGUID>> platformMappedGuidSet;
+	std::map<PlatformRef, std::unordered_set<Warp::ConstraintGUID>> platformMappedConstraintSet;
 	for (const auto& func : requests)
 	{
-		const auto guid = Warp::GetAnalysisFunctionGUID(*func);
-		if (!guid.has_value())
+		const auto warpFunc = Warp::Function::Get(*func);
+		if (!warpFunc)
 			continue;
 		auto platform = func->GetPlatform();
-		platformMappedGuids[platform].push_back(guid.value());
+		platformMappedGuidSet[platform].insert(warpFunc->GetGUID());
+
+		// We want to keep track of the guids so we can constrain the server response to only return functions with any of them.
+		const auto constraints = warpFunc->GetConstraints();
+		std::vector<Warp::ConstraintGUID> constraintGuids;
+		constraintGuids.reserve(constraints.size());
+		for (const auto& constraint : constraints)
+			constraintGuids.push_back(constraint.guid);
+		platformMappedConstraintSet[platform].insert(constraintGuids.begin(), constraintGuids.end());
 	}
 
+	std::map<PlatformRef, std::vector<Warp::FunctionGUID>> platformMappedGuids;
+	for (const auto& [platform, guids] : platformMappedGuidSet)
+		platformMappedGuids[platform] = std::vector(guids.begin(), guids.end());
+
+	// We keep them in the set above so we don't duplicate a bunch for functions with the same set of constraint guids.
+	std::map<PlatformRef, std::vector<Warp::ConstraintGUID>> platformMappedConstraints;
+	for (const auto& [platform, guids] : platformMappedConstraintSet)
+		platformMappedConstraints[platform] = std::vector(guids.begin(), guids.end());
+
 	for (const auto& [platform, guids] : platformMappedGuids)
 	{
 		m_logger->LogDebugF("Fetching {} functions for platform {}", guids.size(), platform->GetName());
 		auto target = Warp::Target::FromPlatform(*platform);
 		for (const auto& container : Warp::Container::All())
-			container->FetchFunctions(*target, guids, allowedTags);
+			container->FetchFunctions(*target, guids, allowedTags, platformMappedConstraints[platform]);
 
 		std::lock_guard<std::mutex> lock(m_requestMutex);
 		for (const auto& guid : guids)
diff --git a/plugins/warp/ui/shared/misc.h b/plugins/warp/ui/shared/misc.h
index 97adcab0..92e27fd8 100644
--- a/plugins/warp/ui/shared/misc.h
+++ b/plugins/warp/ui/shared/misc.h
@@ -138,10 +138,10 @@ inline void SetTagsToView(const BinaryViewRef& view, const std::vector<Warp::Sou
 	settings->Set(ALLOWED_TAGS_SETTING, tags, view);
 }
 
-inline int GetBatchSizeFromView(const BinaryViewRef& view)
+inline size_t GetBatchSizeFromView(const BinaryViewRef& view)
 {
 	auto settings = BinaryNinja::Settings::Instance();
 	if (!settings->Contains(BATCH_SIZE_SETTING))
-		return 100;
+		return 10000;
 	return settings->Get<uint64_t>(BATCH_SIZE_SETTING, view);
 }
\ No newline at end of file