logicaffeine_compile/optimize/egraph/
extract.rs

1//! Bottom-up cost extraction (cycle-tolerant).
2//!
3//! LogosCost is a static per-op estimate tuned to the VM/JIT tier: shifts,
4//! adds and compares are unit ops; multiply is a few; divide/modulo are an
5//! order worse. Costs propagate to a per-class fixpoint — cyclic classes
6//! (x ≡ x + 0) simply never beat their finite leaf, so extraction always
7//! terminates with a finite tree.
8
9use std::collections::HashMap;
10
11use super::{CompilerEGraph, CompilerENode, NodeId};
12
13#[derive(Debug)]
14pub struct ExtractTree {
15    pub node: CompilerENode,
16    pub children: Vec<ExtractTree>,
17}
18
19fn op_cost(node: &CompilerENode) -> u64 {
20    use CompilerENode::*;
21    match node {
22        Int(_) | Bool(_) | Float(_) | Var(..) => 1,
23        // Opaque subtrees re-emit the original expression; bias against
24        // duplicating them when a modeled form exists.
25        Opaque(_) => 4,
26        Add(..) | Sub(..) | Shl(..) | Shr(..) | BitXor(..) | BitAnd(..) | BitOr(..)
27        | And(..) | Or(..) | Not(..)
28        | Eq(..) | Ne(..) | Lt(..) | Le(..) | Gt(..) | Ge(..) => 1,
29        Len(..) => 2,
30        Index(..) => 3,
31        Mul(..) => 4,
32        Contains(..) => 6,
33        Slice(..) => 8,
34        // O(n) materialization — the fusion algebra exists to delete these.
35        Copy(..) => 10,
36        Concat(..) => 8,
37        Div(..) | Mod(..) => 12,
38    }
39}
40
41/// Per-class best (cost, representative node id), to fixpoint. Nodes the
42/// `admissible` filter rejects never become representatives — extraction
43/// for one statement must not reuse a STALE variable version or another
44/// statement's opaque (that would duplicate its effect).
45fn best_costs(
46    eg: &mut CompilerEGraph,
47    admissible: &dyn Fn(&CompilerENode) -> bool,
48) -> HashMap<NodeId, (u64, NodeId)> {
49    let mut best: HashMap<NodeId, (u64, NodeId)> = HashMap::new();
50    loop {
51        let mut changed = false;
52        for id in 0..eg.node_count() {
53            let node = eg.canonical_node(id);
54            if !admissible(&node) {
55                continue;
56            }
57            let mut total = op_cost(&node);
58            let mut ready = true;
59            for c in node.children() {
60                match best.get(&c) {
61                    Some(&(cc, _)) => total = total.saturating_add(cc),
62                    None => {
63                        ready = false;
64                        break;
65                    }
66                }
67            }
68            if !ready {
69                continue;
70            }
71            let root = eg.find(id);
72            match best.get(&root) {
73                Some(&(old, _)) if old <= total => {}
74                _ => {
75                    best.insert(root, (total, id));
76                    changed = true;
77                }
78            }
79        }
80        if !changed {
81            break;
82        }
83    }
84    best
85}
86
87fn build_tree(
88    eg: &mut CompilerEGraph,
89    best: &HashMap<NodeId, (u64, NodeId)>,
90    class: NodeId,
91) -> Option<ExtractTree> {
92    let root = eg.find(class);
93    let (_, rep) = *best.get(&root)?;
94    let node = eg.canonical_node(rep);
95    let children = node
96        .children()
97        .into_iter()
98        .map(|c| build_tree(eg, best, c))
99        .collect::<Option<Vec<_>>>()?;
100    Some(ExtractTree { node, children })
101}
102
103/// The cheapest finite tree representing `root`'s class.
104pub fn best_tree(eg: &mut CompilerEGraph, root: NodeId) -> ExtractTree {
105    let best = best_costs(eg, &|_| true);
106    build_tree(eg, &best, root)
107        .unwrap_or_else(|| panic!("class has no finite-cost representative"))
108}
109
110/// The cheapest finite tree whose every node passes `admissible`, or None
111/// when no such tree exists (the caller keeps the original expression).
112pub fn best_tree_filtered(
113    eg: &mut CompilerEGraph,
114    root: NodeId,
115    admissible: &dyn Fn(&CompilerENode) -> bool,
116) -> Option<ExtractTree> {
117    let best = best_costs(eg, admissible);
118    build_tree(eg, &best, root)
119}
logicaffeine_compile/optimize/egraph/extract.rs

logicaffeine_compile/optimize/egraph/
extract.rs
