Skip to content

Commit 6920a58

Browse files
marirsmarirs
authored
Merge pull request #19 from mnaza/perf/global-features-cache
Perf: compute global features once per file; single hash lookup in Feature::evaluate
2 parents e760998 + 5b26ce4 commit 6920a58

3 files changed

Lines changed: 95 additions & 65 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,23 @@ This project follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
2626
`u16::from_be_bytes` received the chunk bytes in reverse order, turning
2727
every UTF-16BE string into non-ASCII garbage that was then dropped.
2828

29+
### Performance (closes [#18](https://github.com/marirs/capa-rs/issues/18))
30+
31+
- **Global features computed once per file instead of once per
32+
instruction** `extract_global_features` paid a full
33+
`goblin::Object::parse` for the OS feature on every call — and it was
34+
called per instruction, per basic block and per function. The result
35+
(OS/arch — constant per file) is now cached in a
36+
`once_cell::sync::OnceCell` on the smda extractor (thread-safe for the
37+
rayon per-function loop). Instruction-level extraction on
38+
`data/mimikatz.exe_` (808 KiB, ~136k instructions) drops from ~8.4 s to
39+
~0.2 s (**~40×**); `data/Demo64.dll` from ~23 ms to ~3.4 ms.
40+
- **Single hash lookup in `Feature::evaluate`** All 20 feature types used
41+
`contains_key` + indexing, doing two hash-map lookups and cloning
42+
`self` (including its `HashSet<Scope>`) twice on every hit — the
43+
hottest operation of the rule engine. All sites now do one
44+
`HashMap::get`, matching the pattern `RangeStatement` already used.
45+
2946
## [0.5.2] — xor-zero number(0), regex /i fast path, rule pre-pruning
3047

3148
### Fixed — feature extraction parity

‎src/extractor/smda.rs‎

Lines changed: 38 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,14 @@ pub struct Extractor<'a> {
110110
report: DisassemblyReport<'a>,
111111
buf: &'a [u8],
112112
path: String,
113+
/// Cache for `extract_global_features` (issue marirs/capa-rs#18):
114+
/// OS/arch are constant per file, but computing the OS feature
115+
/// costs a full `goblin::Object::parse` of the buffer — and the
116+
/// feature was recomputed on every instruction and basic block.
117+
/// `OnceCell` (not `lazy_static`) because the value is per-file;
118+
/// the `sync` flavour because `find_capabilities` shares the
119+
/// extractor across rayon worker threads.
120+
global_features_cache: once_cell::sync::OnceCell<Vec<(crate::rules::features::Feature, u64)>>,
113121
}
114122

115123
impl std::fmt::Debug for Extractor<'_> {
@@ -175,22 +183,34 @@ impl<'data> super::Extractor for Extractor<'data> {
175183
}
176184

177185
fn extract_global_features(&self) -> Result<Vec<(crate::rules::features::Feature, u64)>> {
178-
Ok(vec![
179-
(
180-
crate::rules::features::Feature::Os(crate::rules::features::OsFeature::new(
181-
&self.extract_os()?.to_string(),
182-
"",
183-
)?),
184-
0,
185-
),
186-
(
187-
crate::rules::features::Feature::Arch(crate::rules::features::ArchFeature::new(
188-
&self.extract_arch()?.to_string(),
189-
"",
190-
)?),
191-
0,
192-
),
193-
])
186+
// Issue marirs/capa-rs#18: computed once per extractor, then
187+
// cloned — the uncached version paid a full goblin parse for
188+
// the OS feature on every call (i.e. per instruction).
189+
Ok(self
190+
.global_features_cache
191+
.get_or_try_init(|| -> Result<Vec<(crate::rules::features::Feature, u64)>> {
192+
Ok(vec![
193+
(
194+
crate::rules::features::Feature::Os(
195+
crate::rules::features::OsFeature::new(
196+
&self.extract_os()?.to_string(),
197+
"",
198+
)?,
199+
),
200+
0,
201+
),
202+
(
203+
crate::rules::features::Feature::Arch(
204+
crate::rules::features::ArchFeature::new(
205+
&self.extract_arch()?.to_string(),
206+
"",
207+
)?,
208+
),
209+
0,
210+
),
211+
])
212+
})?
213+
.clone())
194214
}
195215

196216
fn extract_file_features(&self) -> Result<Vec<(crate::rules::features::Feature, u64)>> {
@@ -455,6 +475,7 @@ impl<'data> Extractor<'data> {
455475
report,
456476
buf: data,
457477
path: path.to_string(),
478+
global_features_cache: once_cell::sync::OnceCell::new(),
458479
})
459480
}
460481

@@ -483,6 +504,7 @@ impl<'data> Extractor<'data> {
483504
// Synthetic path — no real file backs a buffer-mode
484505
// extractor. Kept stable so `Debug` output is uniform.
485506
path: "<buffer>".to_string(),
507+
global_features_cache: once_cell::sync::OnceCell::new(),
486508
})
487509
}
488510

‎src/rules/features.rs‎

Lines changed: 40 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -413,8 +413,8 @@ impl FunctionNameFeature {
413413
&self,
414414
features: &std::collections::HashMap<Feature, Vec<u64>>,
415415
) -> Result<(bool, Vec<u64>)> {
416-
if features.contains_key(&Feature::FunctionName(self.clone())) {
417-
return Ok((true, features[&Feature::FunctionName(self.clone())].clone()));
416+
if let Some(locations) = features.get(&Feature::FunctionName(self.clone())) {
417+
return Ok((true, locations.clone()));
418418
}
419419
Ok((false, vec![]))
420420
}
@@ -458,8 +458,8 @@ impl SectionFeature {
458458
&self,
459459
features: &std::collections::HashMap<Feature, Vec<u64>>,
460460
) -> Result<(bool, Vec<u64>)> {
461-
if features.contains_key(&Feature::Section(self.clone())) {
462-
return Ok((true, features[&Feature::Section(self.clone())].clone()));
461+
if let Some(locations) = features.get(&Feature::Section(self.clone())) {
462+
return Ok((true, locations.clone()));
463463
}
464464
Ok((false, vec![]))
465465
}
@@ -503,8 +503,8 @@ impl ImportFeature {
503503
&self,
504504
features: &std::collections::HashMap<Feature, Vec<u64>>,
505505
) -> Result<(bool, Vec<u64>)> {
506-
if features.contains_key(&Feature::Import(self.clone())) {
507-
return Ok((true, features[&Feature::Import(self.clone())].clone()));
506+
if let Some(locations) = features.get(&Feature::Import(self.clone())) {
507+
return Ok((true, locations.clone()));
508508
}
509509
Ok((false, vec![]))
510510
}
@@ -548,8 +548,8 @@ impl ExportFeature {
548548
&self,
549549
features: &std::collections::HashMap<Feature, Vec<u64>>,
550550
) -> Result<(bool, Vec<u64>)> {
551-
if features.contains_key(&Feature::Export(self.clone())) {
552-
return Ok((true, features[&Feature::Export(self.clone())].clone()));
551+
if let Some(locations) = features.get(&Feature::Export(self.clone())) {
552+
return Ok((true, locations.clone()));
553553
}
554554
Ok((false, vec![]))
555555
}
@@ -588,8 +588,8 @@ impl BasicBlockFeature {
588588
&self,
589589
features: &std::collections::HashMap<Feature, Vec<u64>>,
590590
) -> Result<(bool, Vec<u64>)> {
591-
if features.contains_key(&Feature::BasicBlock(self.clone())) {
592-
return Ok((true, features[&Feature::BasicBlock(self.clone())].clone()));
591+
if let Some(locations) = features.get(&Feature::BasicBlock(self.clone())) {
592+
return Ok((true, locations.clone()));
593593
}
594594
Ok((false, vec![]))
595595
}
@@ -614,8 +614,8 @@ impl MnemonicFeature {
614614
&self,
615615
features: &std::collections::HashMap<Feature, Vec<u64>>,
616616
) -> Result<(bool, Vec<u64>)> {
617-
if features.contains_key(&Feature::Mnemonic(self.clone())) {
618-
return Ok((true, features[&Feature::Mnemonic(self.clone())].clone()));
617+
if let Some(locations) = features.get(&Feature::Mnemonic(self.clone())) {
618+
return Ok((true, locations.clone()));
619619
}
620620
Ok((false, vec![]))
621621
}
@@ -661,8 +661,8 @@ impl OffsetFeature {
661661
&self,
662662
features: &std::collections::HashMap<Feature, Vec<u64>>,
663663
) -> Result<(bool, Vec<u64>)> {
664-
if features.contains_key(&Feature::Offset(self.clone())) {
665-
return Ok((true, features[&Feature::Offset(self.clone())].clone()));
664+
if let Some(locations) = features.get(&Feature::Offset(self.clone())) {
665+
return Ok((true, locations.clone()));
666666
}
667667
Ok((false, vec![]))
668668
}
@@ -708,11 +708,8 @@ impl OperandOffsetFeature {
708708
&self,
709709
features: &std::collections::HashMap<Feature, Vec<u64>>,
710710
) -> Result<(bool, Vec<u64>)> {
711-
if features.contains_key(&Feature::OperandOffset(self.clone())) {
712-
return Ok((
713-
true,
714-
features[&Feature::OperandOffset(self.clone())].clone(),
715-
));
711+
if let Some(locations) = features.get(&Feature::OperandOffset(self.clone())) {
712+
return Ok((true, locations.clone()));
716713
}
717714
Ok((false, vec![]))
718715
}
@@ -766,8 +763,8 @@ impl NumberFeature {
766763
&self,
767764
features: &std::collections::HashMap<Feature, Vec<u64>>,
768765
) -> Result<(bool, Vec<u64>)> {
769-
if features.contains_key(&Feature::Number(self.clone())) {
770-
return Ok((true, features[&Feature::Number(self.clone())].clone()));
766+
if let Some(locations) = features.get(&Feature::Number(self.clone())) {
767+
return Ok((true, locations.clone()));
771768
}
772769
Ok((false, vec![]))
773770
}
@@ -813,11 +810,8 @@ impl OperandNumberFeature {
813810
&self,
814811
features: &std::collections::HashMap<Feature, Vec<u64>>,
815812
) -> Result<(bool, Vec<u64>)> {
816-
if features.contains_key(&Feature::OperandNumber(self.clone())) {
817-
return Ok((
818-
true,
819-
features[&Feature::OperandNumber(self.clone())].clone(),
820-
));
813+
if let Some(locations) = features.get(&Feature::OperandNumber(self.clone())) {
814+
return Ok((true, locations.clone()));
821815
}
822816
Ok((false, vec![]))
823817
}
@@ -888,8 +882,8 @@ impl ApiFeature {
888882
&self,
889883
features: &std::collections::HashMap<Feature, Vec<u64>>,
890884
) -> Result<(bool, Vec<u64>)> {
891-
if features.contains_key(&Feature::Api(self.clone())) {
892-
return Ok((true, features[&Feature::Api(self.clone())].clone()));
885+
if let Some(locations) = features.get(&Feature::Api(self.clone())) {
886+
return Ok((true, locations.clone()));
893887
}
894888
Ok((false, vec![]))
895889
}
@@ -936,8 +930,8 @@ impl PropertyFeature {
936930
&self,
937931
features: &std::collections::HashMap<Feature, Vec<u64>>,
938932
) -> Result<(bool, Vec<u64>)> {
939-
if features.contains_key(&Feature::Property(self.clone())) {
940-
return Ok((true, features[&Feature::Property(self.clone())].clone()));
933+
if let Some(locations) = features.get(&Feature::Property(self.clone())) {
934+
return Ok((true, locations.clone()));
941935
}
942936
Ok((false, vec![]))
943937
}
@@ -989,8 +983,8 @@ impl MatchedRuleFeature {
989983
&self,
990984
features: &std::collections::HashMap<Feature, Vec<u64>>,
991985
) -> Result<(bool, Vec<u64>)> {
992-
if features.contains_key(&Feature::MatchedRule(self.clone())) {
993-
return Ok((true, features[&Feature::MatchedRule(self.clone())].clone()));
986+
if let Some(locations) = features.get(&Feature::MatchedRule(self.clone())) {
987+
return Ok((true, locations.clone()));
994988
}
995989
Ok((false, vec![]))
996990
}
@@ -1068,11 +1062,8 @@ impl CharacteristicFeature {
10681062
&self,
10691063
features: &std::collections::HashMap<Feature, Vec<u64>>,
10701064
) -> Result<(bool, Vec<u64>)> {
1071-
if features.contains_key(&Feature::Characteristic(self.clone())) {
1072-
return Ok((
1073-
true,
1074-
features[&Feature::Characteristic(self.clone())].clone(),
1075-
));
1065+
if let Some(locations) = features.get(&Feature::Characteristic(self.clone())) {
1066+
return Ok((true, locations.clone()));
10761067
}
10771068
Ok((false, vec![]))
10781069
}
@@ -1124,8 +1115,8 @@ impl StringFeature {
11241115
&self,
11251116
features: &std::collections::HashMap<Feature, Vec<u64>>,
11261117
) -> Result<(bool, Vec<u64>)> {
1127-
if features.contains_key(&Feature::String(self.clone())) {
1128-
return Ok((true, features[&Feature::String(self.clone())].clone()));
1118+
if let Some(locations) = features.get(&Feature::String(self.clone())) {
1119+
return Ok((true, locations.clone()));
11291120
}
11301121
Ok((false, vec![]))
11311122
}
@@ -1536,8 +1527,8 @@ impl ArchFeature {
15361527
&self,
15371528
features: &std::collections::HashMap<Feature, Vec<u64>>,
15381529
) -> Result<(bool, Vec<u64>)> {
1539-
if features.contains_key(&Feature::Arch(self.clone())) {
1540-
return Ok((true, features[&Feature::Arch(self.clone())].clone()));
1530+
if let Some(locations) = features.get(&Feature::Arch(self.clone())) {
1531+
return Ok((true, locations.clone()));
15411532
}
15421533
Ok((false, vec![]))
15431534
}
@@ -1594,8 +1585,8 @@ impl NamespaceFeature {
15941585
&self,
15951586
features: &std::collections::HashMap<Feature, Vec<u64>>,
15961587
) -> Result<(bool, Vec<u64>)> {
1597-
if features.contains_key(&Feature::Namespace(self.clone())) {
1598-
return Ok((true, features[&Feature::Namespace(self.clone())].clone()));
1588+
if let Some(locations) = features.get(&Feature::Namespace(self.clone())) {
1589+
return Ok((true, locations.clone()));
15991590
}
16001591
Ok((false, vec![]))
16011592
}
@@ -1645,8 +1636,8 @@ impl ClassFeature {
16451636
&self,
16461637
features: &std::collections::HashMap<Feature, Vec<u64>>,
16471638
) -> Result<(bool, Vec<u64>)> {
1648-
if features.contains_key(&Feature::Class(self.clone())) {
1649-
return Ok((true, features[&Feature::Class(self.clone())].clone()));
1639+
if let Some(locations) = features.get(&Feature::Class(self.clone())) {
1640+
return Ok((true, locations.clone()));
16501641
}
16511642
Ok((false, vec![]))
16521643
}
@@ -1700,8 +1691,8 @@ impl OsFeature {
17001691
&self,
17011692
features: &std::collections::HashMap<Feature, Vec<u64>>,
17021693
) -> Result<(bool, Vec<u64>)> {
1703-
if features.contains_key(&Feature::Os(self.clone())) {
1704-
return Ok((true, features[&Feature::Os(self.clone())].clone()));
1694+
if let Some(locations) = features.get(&Feature::Os(self.clone())) {
1695+
return Ok((true, locations.clone()));
17051696
}
17061697
Ok((false, vec![]))
17071698
}
@@ -1762,8 +1753,8 @@ impl FormatFeature {
17621753
&self,
17631754
features: &std::collections::HashMap<Feature, Vec<u64>>,
17641755
) -> Result<(bool, Vec<u64>)> {
1765-
if features.contains_key(&Feature::Format(self.clone())) {
1766-
return Ok((true, features[&Feature::Format(self.clone())].clone()));
1756+
if let Some(locations) = features.get(&Feature::Format(self.clone())) {
1757+
return Ok((true, locations.clone()));
17671758
}
17681759
Ok((false, vec![]))
17691760
}

0 commit comments

Comments
 (0)