From 81572a6da3a9e806eccde5b91447bb5265280a91 Mon Sep 17 00:00:00 2001 From: Yihong Zhang Date: Sun, 19 Jul 2026 02:02:23 +0000 Subject: [PATCH 1/4] Materialize rows and use dense sort ids in the extractor's Bellman-Ford loop The extractor re-scanned every reachable function table on each Bellman-Ford pass through the bridge's for_each machinery, and looked up per-sort cost/topo maps by string name for every child of every row. Materialize each function's non-subsumed rows once per extractor build, resolve each child column's handling (eq/container/base) up front, and index the per-sort maps by a dense id. Relaxation order is unchanged, so extraction results are identical. Co-Authored-By: Claude Fable 5 --- src/extract.rs | 304 ++++++++++++++++++++++++++++++------------------- 1 file changed, 186 insertions(+), 118 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index 3c7aa40ba..509de9b75 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -128,10 +128,36 @@ pub struct Extractor { rootsorts: Vec, funcs: Vec, cost_model: Box>, - costs: HashMap>, + /// Dense id assigned to each eq sort that some extractable function outputs; + /// indexes into `costs`, `topo_rnk`, and `parent_edge`. + sort_ids: HashMap, + costs: Vec>, topo_rnk_cnt: usize, - topo_rnk: HashMap>, - parent_edge: HashMap)>>, + topo_rnk: Vec>, + parent_edge: Vec)>>, +} + +/// How extraction treats one child column of a function, resolved once so the +/// per-row cost loops avoid repeated sort dispatch and name-keyed lookups. +enum ChildKind { + /// An eq-sort column; `Some` holds the sort's dense id, `None` means no + /// extractable function outputs this sort (its terms have no cost). + EqSort(Option), + Container(ArcSort), + Base(ArcSort), +} + +/// Per-function data for [`Extractor::bellman_ford`]: resolved schema facts +/// plus a flat materialized copy of the table's non-subsumed rows, so each +/// relaxation pass iterates memory instead of re-scanning the table. +struct FuncData { + name: String, + /// Row width in `rows`; 0 iff the table had no (non-subsumed) rows. + arity: usize, + output_idx: usize, + output_sort_id: usize, + child_kinds: Vec, + rows: Vec, } impl Extractor { @@ -224,29 +250,26 @@ impl Extractor { } // Initialize the tables to have the reachable entries - let mut costs: HashMap> = Default::default(); - let mut topo_rnk: HashMap> = Default::default(); - let mut parent_edge: HashMap)>> = - Default::default(); - + let mut sort_ids: HashMap = Default::default(); for func_name in funcs.iter() { let func = egraph.functions.get(func_name).unwrap(); let output_sort_name = func.extraction_output_sort().name(); - if !costs.contains_key(output_sort_name) { - costs.insert(output_sort_name.to_owned(), Default::default()); - topo_rnk.insert(output_sort_name.to_owned(), Default::default()); - parent_edge.insert(output_sort_name.to_owned(), Default::default()); - } + let next_id = sort_ids.len(); + sort_ids + .entry(output_sort_name.to_owned()) + .or_insert(next_id); } + let n_sorts = sort_ids.len(); let mut extractor = Extractor { rootsorts, funcs, cost_model: Box::new(cost_model), - costs, + sort_ids, + costs: (0..n_sorts).map(|_| Default::default()).collect(), topo_rnk_cnt: 0, - topo_rnk, - parent_edge, + topo_rnk: (0..n_sorts).map(|_| Default::default()).collect(), + parent_edge: (0..n_sorts).map(|_| Default::default()).collect(), }; extractor.bellman_ford(egraph); @@ -269,7 +292,9 @@ impl Extractor { .container_cost(egraph, sort, value, &ch_costs), ) } else if sort.is_eq_sort() { - self.costs.get(sort.name())?.get(&value).cloned() + self.costs[*self.sort_ids.get(sort.name())?] + .get(&value) + .cloned() } else { // Primitive Some(self.cost_model.base_value_cost(egraph, sort, value)) @@ -311,8 +336,8 @@ impl Extractor { usize::max(ret, self.compute_topo_rnk_node(egraph, *value, sort)) }) } else if sort.is_eq_sort() { - if let Some(t) = self.topo_rnk.get(sort.name()) { - *t.get(&value).unwrap_or(&usize::MAX) + if let Some(id) = self.sort_ids.get(sort.name()) { + *self.topo_rnk[*id].get(&value).unwrap_or(&usize::MAX) } else { usize::MAX } @@ -321,23 +346,6 @@ impl Extractor { } } - fn compute_topo_rnk_hyperedge( - &self, - egraph: &EGraph, - row: &egglog_bridge::ScanEntry, - func: &Function, - ) -> usize { - let sorts = &func.schema.input; - let num_children = func.extraction_num_children(); - row.vals - .iter() - .take(num_children) - .zip(sorts.iter()) - .fold(0, |ret, (value, sort)| { - usize::max(ret, self.compute_topo_rnk_node(egraph, *value, sort)) - }) - } - /// We use Bellman-Ford to compute the costs of the relevant eq sorts' terms /// [Bellman-Ford](https://en.wikipedia.org/wiki/Bellman%E2%80%93Ford_algorithm) is a shortest path algorithm. /// The version implemented here computes the shortest path from any node in a set of sources to all the reachable nodes. @@ -348,98 +356,161 @@ impl Extractor { /// Additionally, to avoid cycles in the extraction even when the cost model can assign an equal cost to a term and its subterm. /// It computes a topological rank for each eclass /// and only allows each eclass to have children of classes of strictly smaller ranks in the extraction. - fn bellman_ford(&mut self, egraph: &EGraph) { - let mut ensure_fixpoint = false; + /// Compute the child costs of one materialized row into `ch_costs`, and fold them + /// into the row's total cost. Returns `None` if any child is uncomputed so far. + fn row_cost( + &self, + egraph: &EGraph, + func: &Function, + f: &FuncData, + row: &[Value], + ch_costs: &mut Vec, + ) -> Option { + ch_costs.clear(); + for (kind, value) in f.child_kinds.iter().zip(row.iter()) { + let cost = match kind { + ChildKind::EqSort(Some(id)) => self.costs[*id].get(value)?.clone(), + ChildKind::EqSort(None) => return None, + ChildKind::Container(sort) => self.compute_cost_node(egraph, *value, sort)?, + ChildKind::Base(sort) => self.cost_model.base_value_cost(egraph, sort, *value), + }; + ch_costs.push(cost); + } + let enode = Enode { + children: &row[..f.output_idx], + eclass: row[f.output_idx], + subsumed: false, + }; + Some(self.cost_model.fold( + func.extraction_term_name(), + ch_costs, + self.cost_model.enode_cost(egraph, func, &enode), + )) + } - let funcs = self.funcs.clone(); + fn bellman_ford(&mut self, egraph: &EGraph) { + // Materialize each function's non-subsumed rows once, so the relaxation + // passes below iterate plain memory instead of re-scanning every table. + let func_data: Vec = self + .funcs + .iter() + .map(|func_name| { + let func = egraph.functions.get(func_name).unwrap(); + let num_children = func.extraction_num_children(); + let child_kinds = func.schema.input[..num_children] + .iter() + .map(|sort| { + if sort.is_container_sort() { + ChildKind::Container(sort.clone()) + } else if sort.is_eq_sort() { + ChildKind::EqSort(self.sort_ids.get(sort.name()).copied()) + } else { + ChildKind::Base(sort.clone()) + } + }) + .collect(); + let mut arity = 0; + let mut rows = Vec::new(); + egraph.backend.for_each(func.backend_id, |row| { + if !row.subsumed { + arity = row.vals.len(); + rows.extend_from_slice(row.vals); + } + }); + FuncData { + name: func_name.clone(), + arity, + output_idx: func.extraction_output_index(), + output_sort_id: self.sort_ids[func.extraction_output_sort().name()], + child_kinds, + rows, + } + }) + .collect(); + let mut ch_costs: Vec = Vec::new(); + let mut ensure_fixpoint = false; while !ensure_fixpoint { ensure_fixpoint = true; - for func_name in funcs.iter() { - let func = egraph.functions.get(func_name).unwrap(); - let target_sort = func.extraction_output_sort(); - - let output_idx = func.extraction_output_index(); - let relax_hyperedge = |row: egglog_bridge::ScanEntry| { - if !row.subsumed { - let target = &row.vals[output_idx]; - let mut updated = false; - if let Some(new_cost) = self.compute_cost_hyperedge(egraph, &row, func) { - match self - .costs - .get_mut(target_sort.name()) - .unwrap() - .entry(*target) - { - HEntry::Vacant(e) => { - updated = true; - e.insert(new_cost); - } - HEntry::Occupied(mut e) => { - if new_cost < *(e.get()) { - updated = true; - e.insert(new_cost); - } - } - } + for f in &func_data { + if f.rows.is_empty() { + continue; + } + let func = egraph.functions.get(&f.name).unwrap(); + for row in f.rows.chunks_exact(f.arity) { + let Some(new_cost) = self.row_cost(egraph, func, f, row, &mut ch_costs) else { + continue; + }; + let target = row[f.output_idx]; + let updated = match self.costs[f.output_sort_id].entry(target) { + HEntry::Vacant(e) => { + e.insert(new_cost); + true } - // record the chronological order of the updates - // which serves as a topological order that avoids cycles - // even when a term has a cost equal to its subterms - if updated { - ensure_fixpoint = false; - self.topo_rnk_cnt += 1; - self.topo_rnk - .get_mut(target_sort.name()) - .unwrap() - .insert(*target, self.topo_rnk_cnt); + HEntry::Occupied(mut e) => { + if new_cost < *(e.get()) { + e.insert(new_cost); + true + } else { + false + } } + }; + // record the chronological order of the updates + // which serves as a topological order that avoids cycles + // even when a term has a cost equal to its subterms + if updated { + ensure_fixpoint = false; + self.topo_rnk_cnt += 1; + self.topo_rnk[f.output_sort_id].insert(target, self.topo_rnk_cnt); } - }; - - egraph.backend.for_each(func.backend_id, relax_hyperedge); + } } } // Save the edges for reconstruction - for func_name in funcs.iter() { - let func = egraph.functions.get(func_name).unwrap(); - let target_sort = func.extraction_output_sort(); - let output_idx = func.extraction_output_index(); - - let save_best_parent_edge = |row: egglog_bridge::ScanEntry| { - if !row.subsumed { - let target = &row.vals[output_idx]; - if let Some(best_cost) = self.costs.get(target_sort.name()).unwrap().get(target) - && Some(best_cost.clone()) - == self.compute_cost_hyperedge(egraph, &row, func) - { - // one of the possible best parent edges - let target_topo_rnk = *self - .topo_rnk - .get(target_sort.name()) - .unwrap() - .get(target) - .unwrap(); - if target_topo_rnk > self.compute_topo_rnk_hyperedge(egraph, &row, func) { - // one of the parent edges that avoids cycles - if let HEntry::Vacant(e) = self - .parent_edge - .get_mut(target_sort.name()) - .unwrap() - .entry(*target) - { - e.insert((func.decl.name.clone(), row.vals.to_vec())); - } - } + for f in &func_data { + if f.rows.is_empty() { + continue; + } + let func = egraph.functions.get(&f.name).unwrap(); + for row in f.rows.chunks_exact(f.arity) { + let target = row[f.output_idx]; + let Some(best_cost) = self.costs[f.output_sort_id].get(&target) else { + continue; + }; + if Some(best_cost.clone()) != self.row_cost(egraph, func, f, row, &mut ch_costs) { + continue; + } + // one of the possible best parent edges + let target_topo_rnk = *self.topo_rnk[f.output_sort_id].get(&target).unwrap(); + let edge_topo_rnk = + f.child_kinds + .iter() + .zip(row.iter()) + .fold(0, |ret, (kind, value)| { + usize::max( + ret, + match kind { + ChildKind::EqSort(Some(id)) => { + *self.topo_rnk[*id].get(value).unwrap_or(&usize::MAX) + } + ChildKind::EqSort(None) => usize::MAX, + ChildKind::Container(sort) => { + self.compute_topo_rnk_node(egraph, *value, sort) + } + ChildKind::Base(_) => 0, + }, + ) + }); + if target_topo_rnk > edge_topo_rnk { + // one of the parent edges that avoids cycles + if let HEntry::Vacant(e) = self.parent_edge[f.output_sort_id].entry(target) { + e.insert((func.decl.name.clone(), row.to_vec())); } } - }; - - egraph - .backend - .for_each(func.backend_id, save_best_parent_edge); + } } } @@ -482,10 +553,7 @@ impl Extractor { ch_terms, ) } else if sort.is_eq_sort() { - let (func_name, hyperedge) = self - .parent_edge - .get(sort.name()) - .unwrap() + let (func_name, hyperedge) = self.parent_edge[self.sort_ids[sort.name()]] .get(&value) .unwrap(); let func = egraph.functions.get(func_name).unwrap(); From 28a1122319bf4660a2d61718d9c5ee8f7c28a6e2 Mon Sep 17 00:00:00 2001 From: Yihong Zhang Date: Sun, 19 Jul 2026 02:43:19 +0000 Subject: [PATCH 2/4] Skip converged functions in the extractor's Bellman-Ford passes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Track a version counter per eq sort, bumped on every cost update, and snapshot each function's dependency-sort versions just before its rows are relaxed. A function whose dependency sorts (eq children, plus eq sorts reachable inside container children) are unchanged since its last visit recomputes identical costs for every row, so it is skipped without altering the relaxation trace — topo ranks and extracted terms are byte-identical. Co-Authored-By: Claude Fable 5 --- src/extract.rs | 72 +++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 71 insertions(+), 1 deletion(-) diff --git a/src/extract.rs b/src/extract.rs index 509de9b75..41fdeecb0 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -388,6 +388,30 @@ impl Extractor { )) } + /// Collect the ids of every eq sort reachable inside a container sort, + /// i.e. the sorts a container child's cost can depend on. `visited` guards + /// against container sorts that reach themselves (e.g. an `UnstableFn` + /// closing over values of its own sort). + fn collect_container_dep_sorts( + &self, + sort: &ArcSort, + visited: &mut HashSet, + out: &mut Vec, + ) { + if !visited.insert(sort.name().to_owned()) { + return; + } + if sort.is_container_sort() { + for inner in sort.inner_sorts() { + self.collect_container_dep_sorts(&inner, visited, out); + } + } else if sort.is_eq_sort() + && let Some(id) = self.sort_ids.get(sort.name()) + { + out.push(*id); + } + } + fn bellman_ford(&mut self, egraph: &EGraph) { // Materialize each function's non-subsumed rows once, so the relaxation // passes below iterate plain memory instead of re-scanning every table. @@ -428,15 +452,60 @@ impl Extractor { }) .collect(); + // The eq sorts each function's row costs depend on: eq children directly, + // plus every eq sort reachable inside a container child. + let func_deps: Vec> = func_data + .iter() + .map(|f| { + let mut deps: Vec = Vec::new(); + let mut visited: HashSet = Default::default(); + for kind in &f.child_kinds { + match kind { + ChildKind::EqSort(Some(id)) => deps.push(*id), + ChildKind::EqSort(None) | ChildKind::Base(_) => {} + ChildKind::Container(sort) => { + self.collect_container_dep_sorts(sort, &mut visited, &mut deps) + } + } + } + deps.sort_unstable(); + deps.dedup(); + deps + }) + .collect(); + + // A function's pass recomputes exactly the same cost for every row unless + // one of its dependency sorts gained a cost update since the snapshot + // taken right before the function was last processed. Skipping it in that + // case drops no updates, so the relaxation trace (and hence topo ranks + // and extracted terms) is unchanged. + let mut sort_versions: Vec = vec![0; self.costs.len()]; + let mut last_seen: Vec>> = vec![None; func_data.len()]; + let mut ch_costs: Vec = Vec::new(); let mut ensure_fixpoint = false; while !ensure_fixpoint { ensure_fixpoint = true; - for f in &func_data { + for (fi, f) in func_data.iter().enumerate() { if f.rows.is_empty() { continue; } + match &mut last_seen[fi] { + Some(seen) + if func_deps[fi] + .iter() + .zip(seen.iter()) + .all(|(d, s)| sort_versions[*d] == *s) => + { + continue; + } + Some(seen) => { + seen.clear(); + seen.extend(func_deps[fi].iter().map(|d| sort_versions[*d])); + } + slot => *slot = Some(func_deps[fi].iter().map(|d| sort_versions[*d]).collect()), + } let func = egraph.functions.get(&f.name).unwrap(); for row in f.rows.chunks_exact(f.arity) { let Some(new_cost) = self.row_cost(egraph, func, f, row, &mut ch_costs) else { @@ -462,6 +531,7 @@ impl Extractor { // even when a term has a cost equal to its subterms if updated { ensure_fixpoint = false; + sort_versions[f.output_sort_id] += 1; self.topo_rnk_cnt += 1; self.topo_rnk[f.output_sort_id].insert(target, self.topo_rnk_cnt); } From b2f727a7541a326e4295cb8bee0728f1a64a4476 Mon Sep 17 00:00:00 2001 From: Yihong Zhang Date: Sun, 19 Jul 2026 02:56:34 +0000 Subject: [PATCH 3/4] Track dirty rows semi-naively in the extractor's Bellman-Ford loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the function-level dependency gate with row-level tracking: a reverse index maps each (sort, value) an extractable row reads — eq children plus eq values nested inside container children — to the rows reading it, and each cost update marks only those rows dirty. Sweeps visit dirty rows in the same (function, row) order as the naive passes, so the update trace, topo ranks, and extracted terms are unchanged. Co-Authored-By: Claude Fable 5 --- src/extract.rs | 135 +++++++++++++++++++++++++++++-------------------- 1 file changed, 79 insertions(+), 56 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index 41fdeecb0..6a654ccf7 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -388,27 +388,26 @@ impl Extractor { )) } - /// Collect the ids of every eq sort reachable inside a container sort, - /// i.e. the sorts a container child's cost can depend on. `visited` guards - /// against container sorts that reach themselves (e.g. an `UnstableFn` - /// closing over values of its own sort). - fn collect_container_dep_sorts( + /// Report every (sort id, value) pair a container value's cost depends on: + /// the eq values stored anywhere inside it, found by recursing through + /// nested container values. + fn register_container_deps( &self, + egraph: &EGraph, sort: &ArcSort, - visited: &mut HashSet, - out: &mut Vec, + value: Value, + register: &mut impl FnMut(usize, Value), ) { - if !visited.insert(sort.name().to_owned()) { - return; - } if sort.is_container_sort() { - for inner in sort.inner_sorts() { - self.collect_container_dep_sorts(&inner, visited, out); + for (inner_sort, inner_value) in + sort.inner_values(egraph.backend.container_values(), value) + { + self.register_container_deps(egraph, &inner_sort, inner_value, register); } } else if sort.is_eq_sort() && let Some(id) = self.sort_ids.get(sort.name()) { - out.push(*id); + register(*id, value); } } @@ -452,62 +451,76 @@ impl Extractor { }) .collect(); - // The eq sorts each function's row costs depend on: eq children directly, - // plus every eq sort reachable inside a container child. - let func_deps: Vec> = func_data - .iter() - .map(|f| { - let mut deps: Vec = Vec::new(); - let mut visited: HashSet = Default::default(); - for kind in &f.child_kinds { + // Reverse dependency index: (sort id, value) -> rows reading that value + // as an eq child, directly or inside a container child. + let mut child_index: HashMap<(usize, Value), Vec<(u32, u32)>> = Default::default(); + for (fi, f) in func_data.iter().enumerate() { + if f.rows.is_empty() { + continue; + } + for (ri, row) in f.rows.chunks_exact(f.arity).enumerate() { + for (kind, value) in f.child_kinds.iter().zip(row.iter()) { + let mut register = |sid: usize, v: Value| { + child_index + .entry((sid, v)) + .or_default() + .push((fi as u32, ri as u32)); + }; match kind { - ChildKind::EqSort(Some(id)) => deps.push(*id), + ChildKind::EqSort(Some(id)) => register(*id, *value), ChildKind::EqSort(None) | ChildKind::Base(_) => {} ChildKind::Container(sort) => { - self.collect_container_dep_sorts(sort, &mut visited, &mut deps) + self.register_container_deps(egraph, sort, *value, &mut register) } } } - deps.sort_unstable(); - deps.dedup(); - deps + } + } + + // Semi-naive relaxation: a row recomputes exactly the same cost against a + // never-increasing target unless one of its child (sort, value) costs + // changed since the row was last evaluated, so only such "dirty" rows are + // (re)visited. Rows are swept in the same (function, row) order as the + // naive pass loop, so the update trace — and hence topo ranks and + // extracted terms — is unchanged. + let mut dirty: Vec> = func_data + .iter() + .map(|f| { + vec![ + true; + if f.arity == 0 { + 0 + } else { + f.rows.len() / f.arity + } + ] }) .collect(); - - // A function's pass recomputes exactly the same cost for every row unless - // one of its dependency sorts gained a cost update since the snapshot - // taken right before the function was last processed. Skipping it in that - // case drops no updates, so the relaxation trace (and hence topo ranks - // and extracted terms) is unchanged. - let mut sort_versions: Vec = vec![0; self.costs.len()]; - let mut last_seen: Vec>> = vec![None; func_data.len()]; + let mut dirty_count: Vec = dirty.iter().map(|d| d.len()).collect(); let mut ch_costs: Vec = Vec::new(); - let mut ensure_fixpoint = false; - while !ensure_fixpoint { - ensure_fixpoint = true; - + loop { + let mut any = false; for (fi, f) in func_data.iter().enumerate() { - if f.rows.is_empty() { + if dirty_count[fi] == 0 { continue; } - match &mut last_seen[fi] { - Some(seen) - if func_deps[fi] - .iter() - .zip(seen.iter()) - .all(|(d, s)| sort_versions[*d] == *s) => - { + any = true; + let func = egraph.functions.get(&f.name).unwrap(); + // Marks set at or behind the sweep cursor (including a row + // re-marking itself) stay for the next sweep; marks ahead of it + // are picked up in this one, matching the naive pass exactly. + let n_rows = dirty[fi].len(); + let mut ri = 0; + while ri < n_rows { + if !dirty[fi][ri] { + ri += 1; continue; } - Some(seen) => { - seen.clear(); - seen.extend(func_deps[fi].iter().map(|d| sort_versions[*d])); - } - slot => *slot = Some(func_deps[fi].iter().map(|d| sort_versions[*d]).collect()), - } - let func = egraph.functions.get(&f.name).unwrap(); - for row in f.rows.chunks_exact(f.arity) { + dirty[fi][ri] = false; + dirty_count[fi] -= 1; + let row = &f.rows[ri * f.arity..(ri + 1) * f.arity]; + ri += 1; let Some(new_cost) = self.row_cost(egraph, func, f, row, &mut ch_costs) else { continue; }; @@ -530,13 +543,23 @@ impl Extractor { // which serves as a topological order that avoids cycles // even when a term has a cost equal to its subterms if updated { - ensure_fixpoint = false; - sort_versions[f.output_sort_id] += 1; self.topo_rnk_cnt += 1; self.topo_rnk[f.output_sort_id].insert(target, self.topo_rnk_cnt); + if let Some(readers) = child_index.get(&(f.output_sort_id, target)) { + for &(dfi, dri) in readers { + let (dfi, dri) = (dfi as usize, dri as usize); + if !dirty[dfi][dri] { + dirty[dfi][dri] = true; + dirty_count[dfi] += 1; + } + } + } } } } + if !any { + break; + } } // Save the edges for reconstruction From 371f45a1917124edfd7a365235b2bb5586af606c Mon Sep 17 00:00:00 2001 From: Yihong Zhang Date: Sun, 19 Jul 2026 03:16:24 +0000 Subject: [PATCH 4/4] Cache the extractor across consecutive extract commands MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An extract command built a fresh Extractor — recomputing every cost via Bellman-Ford — even when nothing changed since the previous extract, which dominates programs that extract many roots after a run. Cache the extractor's data on the EGraph, keyed by root sort and a fingerprint of all backend table versions (any merged mutation, including union-find changes, alters some table's version). The cache deliberately clones empty so push/pop don't carry cost tables. Co-Authored-By: Claude Fable 5 --- egglog-bridge/src/lib.rs | 14 +++++++ src/extract.rs | 79 ++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 23 +++++++++--- 3 files changed, 110 insertions(+), 6 deletions(-) diff --git a/egglog-bridge/src/lib.rs b/egglog-bridge/src/lib.rs index bc890c007..6d01dc585 100644 --- a/egglog-bridge/src/lib.rs +++ b/egglog-bridge/src/lib.rs @@ -476,6 +476,20 @@ impl EGraph { self.db.get_table(self.funcs[table].table).len() } + /// The version of every table in the database, in table-id order. + /// + /// Any merged mutation that reaches a table (inserts, rebuilds, union-find + /// changes, clears, or newly declared tables) changes the result, so equal + /// vectors mean the readable state of the database is unchanged. + pub fn table_versions(&self) -> Vec<(u64, u64)> { + (0..self.db.next_table_id().index()) + .map(|i| { + let v = self.db.get_table(TableId::from_usize(i)).version(); + (v.major.index() as u64, v.minor.index() as u64) + }) + .collect() + } + /// Remove every row from the given function's backing table. /// /// This is the bulk counterpart to staging a `remove` for every key in the diff --git a/src/extract.rs b/src/extract.rs index 6a654ccf7..f2241b1ee 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -137,6 +137,85 @@ pub struct Extractor { parent_edge: Vec)>>, } +/// The reusable state of a default-cost-model [`Extractor`], detached from its +/// (unsendable) boxed cost model so it can be stored on the [`EGraph`]. See +/// [`ExtractorCache`]. +pub(crate) struct ExtractorData { + rootsorts: Vec, + funcs: Vec, + sort_ids: HashMap, + costs: Vec>, + topo_rnk_cnt: usize, + topo_rnk: Vec>, + parent_edge: Vec)>>, +} + +impl Extractor { + pub(crate) fn from_data(data: ExtractorData) -> Self { + Extractor { + rootsorts: data.rootsorts, + funcs: data.funcs, + cost_model: Box::new(TreeAdditiveCostModel::default()), + sort_ids: data.sort_ids, + costs: data.costs, + topo_rnk_cnt: data.topo_rnk_cnt, + topo_rnk: data.topo_rnk, + parent_edge: data.parent_edge, + } + } + + pub(crate) fn into_data(self) -> ExtractorData { + ExtractorData { + rootsorts: self.rootsorts, + funcs: self.funcs, + sort_ids: self.sort_ids, + costs: self.costs, + topo_rnk_cnt: self.topo_rnk_cnt, + topo_rnk: self.topo_rnk, + parent_edge: self.parent_edge, + } + } +} + +/// Cache of the extractor built by the most recent `extract` command, keyed by +/// its root sort and the backend's table-version fingerprint, so runs of +/// consecutive extracts against an unchanged e-graph build the extractor once. +/// +/// Deliberately clones empty: a pushed e-graph rebuilds its extractor on first +/// use rather than carrying a copy of the cost tables through push/pop. +#[derive(Default)] +pub(crate) struct ExtractorCache(Option<(String, Vec<(u64, u64)>, ExtractorData)>); + +impl Clone for ExtractorCache { + fn clone(&self) -> Self { + ExtractorCache(None) + } +} + +impl ExtractorCache { + /// Take the cached data if it was built for `rootsort` against a database + /// whose table versions still equal `versions`. + pub(crate) fn take_if_valid( + &mut self, + rootsort: &str, + versions: &[(u64, u64)], + ) -> Option { + match self.0.take() { + Some((s, v, data)) if s == rootsort && v == versions => Some(data), + _ => None, + } + } + + pub(crate) fn store( + &mut self, + rootsort: String, + versions: Vec<(u64, u64)>, + data: ExtractorData, + ) { + self.0 = Some((rootsort, versions, data)); + } +} + /// How extraction treats one child column of a function, resolved once so the /// per-row cost loops avoid repeated sort dispatch and name-keyed lookups. enum ChildKind { diff --git a/src/lib.rs b/src/lib.rs index aee7bd30b..a99c20cc1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -307,6 +307,7 @@ pub struct EGraph { proof_state: EncodingState, /// In proof mode, this is the program before proof instrumentation and the version we use for proof checking. proof_check_program: Vec, + extractor_cache: extract::ExtractorCache, } /// A user-defined command allows users to inject custom command that can be called @@ -416,6 +417,7 @@ impl Default for EGraph { command_macros: Default::default(), proof_state, proof_check_program: vec![], + extractor_cache: Default::default(), }; add_base_sort(&mut eg, UnitSort, span!()).unwrap(); add_base_sort(&mut eg, StringSort, span!()).unwrap(); @@ -1677,12 +1679,18 @@ impl EGraph { let mut termdag = TermDag::default(); - let extractor = Extractor::compute_costs_from_rootsorts( - Some(vec![sort]), - self, - TreeAdditiveCostModel::default(), - ); - return if n == 0 { + // Consecutive extracts against an unchanged e-graph reuse the + // previous extractor instead of recomputing all costs. + let versions = self.backend.table_versions(); + let extractor = match self.extractor_cache.take_if_valid(sort.name(), &versions) { + Some(data) => Extractor::from_data(data), + None => Extractor::compute_costs_from_rootsorts( + Some(vec![sort.clone()]), + self, + TreeAdditiveCostModel::default(), + ), + }; + let output = if n == 0 { if let Some((cost, term)) = extractor.extract_best(self, &mut termdag, x) { // dont turn termdag into a string if we have messages disabled for performance reasons if log_enabled!(Level::Info) { @@ -1710,6 +1718,9 @@ impl EGraph { } Ok(vec![CommandOutput::ExtractVariants(termdag, terms)]) }; + self.extractor_cache + .store(sort.name().to_owned(), versions, extractor.into_data()); + return output; } ResolvedNCommand::Push(n) => { (0..n).for_each(|_| self.push());