perf: eliminate O(n²) tuple growth and reduce per-match overhead - #2890
Conversation
There was a problem hiding this comment.
Please add bug fixes, new features, breaking changes and anything else you think is worthwhile mentioning to the master (unreleased) section of CHANGELOG.md. If no CHANGELOG update is needed add the following to the PR description: [x] No CHANGELOG update needed
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the performance of capa's core matching engine by addressing several algorithmic inefficiencies. The changes focus on reducing quadratic time complexity operations, particularly those involving tuple concatenations and list manipulations, and optimizing frequently called methods by pre-computing data structures. The overall impact is a more efficient and scalable rule matching process without altering the correctness of the results. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
There was a problem hiding this comment.
Code Review
This pull request introduces several well-reasoned performance improvements by addressing algorithmic complexities in hot paths. The changes from O(n²) tuple concatenation to list appends, pre-computation of rule indices, and the use of collections.deque to avoid O(n) pop(0) operations are all excellent optimizations. The code is clearer and more efficient. I have one suggestion for a further micro-optimization in the rule matching logic.
| remaining = list(candidate_rules_deque) | ||
| remaining.extend([self.rules[rule_name] for rule_name in new_candidates]) | ||
| RuleSet._sort_rules_by_index(rule_index_by_rule_name, remaining) | ||
| candidate_rules_deque = collections.deque(remaining) |
There was a problem hiding this comment.
Since candidate_rules_deque is already sorted topologically, you could achieve better performance here. Instead of converting the deque to a list, extending it, and then re-sorting the entire collection (O(N log N)), you could sort only the new_candidates and then merge the two sorted lists. A merge operation would be closer to O(N), where N is the total number of candidates. Given this is a performance-focused PR, this micro-optimization might be worth considering.
Four data-driven performance improvements identified by profiling the hot paths in capa's rule-matching and capability-finding pipeline: 1. find_static_capabilities / find_dynamic_capabilities (O(n²) → O(n)) Tuple concatenation with `t += (item,)` copies the entire tuple on every iteration. For a binary with N functions this allocates O(N²) total objects. Replace with list accumulation and a single `tuple(list)` conversion at the end. 2. RuleSet._match: pre-compute rule_index_by_rule_name (O(n) → O(1)) `_match` is called once per instruction / basic-block / function scope (potentially millions of times). Previously it rebuilt the name→index dict on every call. The dict is now computed once in `__init__` and stored as `_rule_index_by_scope`, reducing each call to a dict lookup. 3. RuleSet._match: candidate_rules.pop(0) → deque.popleft() (O(n) → O(1)) `list.pop(0)` is O(n) because it shifts every remaining element. Switch to `collections.deque` for O(1) left-side consumption. 4. RuleSet._extract_subscope_rules: list.pop(0) → deque.popleft() (O(n²) → O(n)) Same issue: BFS over rules used list.pop(0), making the whole loop quadratic. Changed to a deque queue for linear-time processing. Fixes mandiant#2880
4a16dd5 to
b995a4f
Compare
When a rule matches and introduces new dependent candidates into _match's work queue, the previous approach converted the deque to a list, extended it with the new items, and re-sorted the whole collection — O((k+m) log(k+m)). Because the existing deque is already topologically sorted, we only need to sort the new additions — O(m log m) — and then merge the two sorted sequences in O(k+m) using heapq.merge. Also adds a CHANGELOG entry for the performance improvements in mandiant#2890.
CHANGELOG updated or no update needed, thanks! 😄
|
do you have a sense for how performance improves in the real world? like if you run capa against mimikatz before/after the changes, how does the wall time differ? anyways, nice finds! |
|
my intuition is that (2) is possibly a significant improvement, while the others are algorithmically correct, but possibly a wash because n is expected to be small (eg candidate rules might be like 10, so the overhead of using a deque may dominate). still, we should write the code in a way that best conveys the meaning, so most of these changes look good to me |
| done = [] | ||
|
|
||
| # use a queue of rules, because we'll be modifying the list (appending new items) as we go. | ||
| while rules: |
There was a problem hiding this comment.
we could possibly use a list and both append and pop from the end. im not sure this has to be a FIFO quite.
| key=lambda r: rule_index_by_rule_name[r.name], | ||
| ) | ||
| candidate_rules_deque = collections.deque( | ||
| heapq.merge( |
There was a problem hiding this comment.
this reads more complicated. unless it makes a noticeable difference, i'm inclined to prefer the old code for simplicity
| rule = candidate_rules.pop(0) | ||
| # Use a deque so that consuming rules from the front is O(1) rather than O(n). | ||
| # list.pop(0) shifts every remaining element; deque.popleft() does not. | ||
| candidate_rules_deque: collections.deque[Rule] = collections.deque(candidate_rules) |
There was a problem hiding this comment.
likewise maybe this can be a LIFO list
Address reviewer feedback: - Replace deque+popleft with list+pop (LIFO stack) in _extract_subscope_rules; processing order doesn't affect correctness, and list.pop() is O(1). - Replace deque+popleft with list+pop (LIFO stack) in _match; sort candidate rules descending so pop() from the end yields the topologically-first rule. - Revert heapq.merge back to the simpler extend+re-sort pattern; the added complexity wasn't justified given the typically small candidate set. - Remove now-unused `import heapq`.
|
hey @williballenthin switched both deques to plain list stacks as you suggested — order doesn't matter in Ran a quick benchmark on cmd.exe (803 functions, 3 runs each): AVERAGE basically a wash, which matches your intuition. the pre-computed rule index is probably where the real |
…diant#2890) * perf: eliminate O(n²) tuple growth and reduce per-match overhead Four data-driven performance improvements identified by profiling the hot paths in capa's rule-matching and capability-finding pipeline: 1. find_static_capabilities / find_dynamic_capabilities (O(n²) → O(n)) Tuple concatenation with `t += (item,)` copies the entire tuple on every iteration. For a binary with N functions this allocates O(N²) total objects. Replace with list accumulation and a single `tuple(list)` conversion at the end. 2. RuleSet._match: pre-compute rule_index_by_rule_name (O(n) → O(1)) `_match` is called once per instruction / basic-block / function scope (potentially millions of times). Previously it rebuilt the name→index dict on every call. The dict is now computed once in `__init__` and stored as `_rule_index_by_scope`, reducing each call to a dict lookup. 3. RuleSet._match: candidate_rules.pop(0) → deque.popleft() (O(n) → O(1)) `list.pop(0)` is O(n) because it shifts every remaining element. Switch to `collections.deque` for O(1) left-side consumption. 4. RuleSet._extract_subscope_rules: list.pop(0) → deque.popleft() (O(n²) → O(n)) Same issue: BFS over rules used list.pop(0), making the whole loop quadratic. Changed to a deque queue for linear-time processing. Fixes mandiant#2880 * perf: use sorted merge instead of full re-sort for new rule candidates When a rule matches and introduces new dependent candidates into _match's work queue, the previous approach converted the deque to a list, extended it with the new items, and re-sorted the whole collection — O((k+m) log(k+m)). Because the existing deque is already topologically sorted, we only need to sort the new additions — O(m log m) — and then merge the two sorted sequences in O(k+m) using heapq.merge. Also adds a CHANGELOG entry for the performance improvements in mandiant#2890. * perf: simplify candidate_rules to LIFO list, revert heapq.merge Address reviewer feedback: - Replace deque+popleft with list+pop (LIFO stack) in _extract_subscope_rules; processing order doesn't affect correctness, and list.pop() is O(1). - Replace deque+popleft with list+pop (LIFO stack) in _match; sort candidate rules descending so pop() from the end yields the topologically-first rule. - Revert heapq.merge back to the simpler extend+re-sort pattern; the added complexity wasn't justified given the typically small candidate set. - Remove now-unused `import heapq`.
Summary
This PR implements data-driven performance improvements to capa's hot paths, motivated by issue #2880 (profiling capa with Scalene to identify and fix hotspots).
Five algorithmic fixes, each targeting a different part of the pipeline:
1. O(n²) → O(n): Tuple concatenation in capability finders
Files:
capa/capabilities/static.py,capa/capabilities/dynamic.pyfeature_counts.functions += (item,)andlibrary_functions += (item,)copied the entire tuple on every function iteration. Since tuples are immutable, each+=allocates a new tuple and copies all existing elements. For a binary with N functions this is O(N²) total allocations.Fix: accumulate into a
list, convert totupleonce after the loop.2. O(n) × M → O(1) × M:
RuleSet._match— pre-computerule_index_by_rule_nameFile:
capa/rules/__init__.py_match()is called once per instruction / basic-block / function scope (potentially millions of times on a real binary). Previously it rebuilt{rule.name: i for i, rule in enumerate(rules)}on every call. The ordering ofrules_by_scopenever changes after__init__, so the dict is now computed once and cached as_rule_index_by_scope.3. O(n) → O(1):
candidate_rules.pop(0)→deque.popleft()File:
capa/rules/__init__.pylist.pop(0)shifts every remaining element, O(n). Replaced withcollections.dequeandpopleft()for O(1) removal from the front.4. O(n²) → O(n):
_extract_subscope_rulesBFS queueFile:
capa/rules/__init__.pySame issue: the BFS that extracts subscope rules used
list.pop(0), making the whole traversal O(n²). Fixed with adequequeue.5. O((k+m) log(k+m)) → O(m log m + k+m): sorted merge for new rule candidates
File:
capa/rules/__init__.pyWhen a matched rule introduces new dependent candidates into the work queue, the previous approach converted the deque to a list, extended it with the new items, and re-sorted the entire collection. Since the existing deque is already topologically sorted, only the new additions need to be sorted, then merged with the existing sorted deque using
heapq.merge.Test plan
tests/test_rules.py— 29 tests passtests/test_match.py— 20 tests pass (including alltest_index_features_*unstable tests)tests/test_engine.py— 8 tests passtests/test_optimizer.py— 1 test passesRuleSetinstantiates,match()returns correct results for matching/non-matching feature sets, dependent-rule chains resolve correctlyCloses #2880