Skip to content

Commit c35fd93

Browse files
committed
[WARP] Deduplicate matched functions within the matcher implicitly
Avoids having to coalesce function objects from multiple sources when enumerating for matched functions.
1 parent f83b7bd commit c35fd93

File tree

1 file changed

+22
-5
lines changed

1 file changed

+22
-5
lines changed

plugins/warp/src/plugin/workflow.rs

Lines changed: 22 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,11 @@ use binaryninja::workflow::{activity, Activity, AnalysisContext, Workflow, Workf
1818
use itertools::Itertools;
1919
use rayon::iter::IntoParallelIterator;
2020
use rayon::iter::ParallelIterator;
21+
use std::cmp::Ordering;
2122
use std::collections::HashMap;
22-
use std::sync::atomic::{AtomicUsize, Ordering};
23-
use std::time::Instant;
23+
use std::sync::atomic::AtomicUsize;
24+
use std::sync::{Arc, Mutex};
25+
use std::time::{Duration, Instant};
2426
use warp::r#type::class::function::{Location, RegisterLocation, StackLocation};
2527
use warp::signature::function::{Function, FunctionGUID};
2628
use warp::target::Target;
@@ -149,7 +151,7 @@ pub fn run_matcher(view: &BinaryView) {
149151
function_guid_with_sources
150152
.into_par_iter()
151153
.for_each(|(guid, sources)| {
152-
let matched_functions: Vec<Function> = sources
154+
let mut matched_functions: Vec<Function> = sources
153155
.iter()
154156
.flat_map(|source| {
155157
container
@@ -158,6 +160,21 @@ pub fn run_matcher(view: &BinaryView) {
158160
})
159161
.collect();
160162

163+
// We sort primarily by symbol, then by type, so we can deduplicate in-place.
164+
matched_functions.sort_unstable_by(|a, b| match a.symbol.cmp(&b.symbol) {
165+
Ordering::Equal => match (&a.ty, &b.ty) {
166+
(None, None) => Ordering::Equal,
167+
(None, Some(_)) => Ordering::Less,
168+
(Some(_), None) => Ordering::Greater,
169+
// TODO: We still need to order the types, probably cant do this in place.
170+
// TODO: Once Type can be ordered, we can remove this entire explicit match stmt.
171+
(Some(_), Some(_)) => Ordering::Equal,
172+
},
173+
other => other,
174+
});
175+
// This removes consecutive duplicates efficiently
176+
matched_functions.dedup_by(|a, b| a.symbol == b.symbol && a.ty == b.ty);
177+
161178
// NOTE: See the comment in `match_function_from_constraints` about this fast fail.
162179
if matcher
163180
.settings
@@ -182,7 +199,7 @@ pub fn run_matcher(view: &BinaryView) {
182199
if let Some(matched_function) =
183200
matcher.match_function_from_constraints(function, &matched_functions)
184201
{
185-
matched_count.fetch_add(1, Ordering::Relaxed);
202+
matched_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
186203
// We were able to find a match, add it to the match cache and then mark the function
187204
// as requiring updates; this is so that we know about it in the applier activity.
188205
insert_cached_function_match(function, Some(matched_function.clone()));
@@ -199,7 +216,7 @@ pub fn run_matcher(view: &BinaryView) {
199216
log::info!(
200217
"Function matching took {:.3} seconds and matched {} functions",
201218
start.elapsed().as_secs_f64(),
202-
matched_count.load(Ordering::Relaxed)
219+
matched_count.load(std::sync::atomic::Ordering::Relaxed)
203220
);
204221
background_task.finish();
205222

0 commit comments

Comments
 (0)