From 6384d1785700ab27734302caf4083456c20740ed Mon Sep 17 00:00:00 2001 From: Kavi Gupta Date: Tue, 17 Jun 2025 14:00:26 -0400 Subject: [PATCH] optimize bottom up utility correction --- src/compression.rs | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/src/compression.rs b/src/compression.rs index db598e7a..d46fea29 100644 --- a/src/compression.rs +++ b/src/compression.rs @@ -1690,6 +1690,12 @@ fn get_utility_of_loc_once(pattern: &Pattern, shared: &SharedData) -> Vec { fn bottom_up_utility_correction(pattern: &Pattern, shared:&SharedData, utility_of_loc_once: &[i32]) -> (Vec,FxHashMap) { let mut cumulative_utility_of_node: Vec = vec![0; shared.corpus_span.len()]; let mut corrected_utils: FxHashMap = Default::default(); + let mut indices = vec![-1; shared.corpus_span.len()]; + for (idx, node) in pattern.match_locations.iter().enumerate() { + // we want to keep track of the index of the match location in the pattern + // so that we can use it later to determine whether to rewrite or not + indices[*node] = idx as i32; + } for node in shared.corpus_span.clone() { @@ -1702,8 +1708,11 @@ fn bottom_up_utility_correction(pattern: &Pattern, shared:&SharedData, utility_o assert!(utility_without_rewrite >= 0); - if let Ok(idx) = pattern.match_locations.binary_search(&node) { + let idxi = indices[node]; + + if idxi >= 0 { // this node is a potential rewrite location + let idx = idxi as usize; let utility_of_args: i32 = pattern.first_zid_of_ivar.iter() .map(|zid| cumulative_utility_of_node[shared.arg_of_zid_node[*zid][&node].unshifted_id])