tidepool_optimize/
dce.rs

1use crate::occ::{get_occ, occ_analysis, Occ};
2use tidepool_eval::{Changed, Pass};
3use tidepool_repr::{get_children, replace_subtree, CoreExpr, CoreFrame};
4
5/// Dead Code Elimination pass.
6/// Removes `LetNonRec` bindings where the binder is unused.
7/// Removes `LetRec` groups where all binders are unused.
8pub struct Dce;
9
10impl Pass for Dce {
11    fn run(&self, expr: &mut CoreExpr) -> Changed {
12        if expr.nodes.is_empty() {
13            return false;
14        }
15        let occ_map = occ_analysis(expr);
16        match try_dce(expr, &occ_map) {
17            Some(new_expr) => {
18                *expr = new_expr;
19                true
20            }
21            None => false,
22        }
23    }
24
25    fn name(&self) -> &str {
26        "Dce"
27    }
28}
29
30fn try_dce(expr: &CoreExpr, occ_map: &crate::occ::OccMap) -> Option<CoreExpr> {
31    try_dce_at(expr, expr.nodes.len() - 1, occ_map)
32}
33
34fn try_dce_at(expr: &CoreExpr, idx: usize, occ_map: &crate::occ::OccMap) -> Option<CoreExpr> {
35    match &expr.nodes[idx] {
36        CoreFrame::LetNonRec { binder, body, .. } => {
37            if get_occ(occ_map, *binder) == Occ::Dead {
38                // Drop the binding, keep just body
39                let body_tree = expr.extract_subtree(*body);
40                Some(replace_subtree(expr, idx, &body_tree))
41            } else {
42                try_children(expr, idx, occ_map)
43            }
44        }
45        CoreFrame::LetRec { bindings, body } => {
46            let all_dead = bindings
47                .iter()
48                .all(|(binder, _)| get_occ(occ_map, *binder) == Occ::Dead);
49            if all_dead {
50                // Drop the entire group, keep just body
51                let body_tree = expr.extract_subtree(*body);
52                Some(replace_subtree(expr, idx, &body_tree))
53            } else {
54                try_children(expr, idx, occ_map)
55            }
56        }
57        _ => try_children(expr, idx, occ_map),
58    }
59}
60
61fn try_children(expr: &CoreExpr, idx: usize, occ_map: &crate::occ::OccMap) -> Option<CoreExpr> {
62    let children = get_children(&expr.nodes[idx]);
63    for child in children {
64        if let Some(result) = try_dce_at(expr, child, occ_map) {
65            return Some(result);
66        }
67    }
68    None
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74    use tidepool_eval::{eval, Env, VecHeap};
75    use tidepool_repr::{Literal, VarId};
76
77    // Helper to build a tree
78    fn tree(nodes: Vec<CoreFrame<usize>>) -> CoreExpr {
79        CoreExpr { nodes }
80    }
81
82    // 1. test_dce_dead_let: let x = 42 in 0 -> 0. Binder Dead, dropped.
83    #[test]
84    fn test_dce_dead_let() {
85        let x = VarId(1);
86        let expr = tree(vec![
87            CoreFrame::Lit(Literal::LitInt(42)), // 0: rhs
88            CoreFrame::Lit(Literal::LitInt(0)),  // 1: body
89            CoreFrame::LetNonRec {
90                binder: x,
91                rhs: 0,
92                body: 1,
93            }, // 2: root
94        ]);
95        let mut dce_expr = expr;
96        let changed = Dce.run(&mut dce_expr);
97        assert!(changed);
98        assert_eq!(dce_expr.nodes.len(), 1);
99        assert_eq!(dce_expr.nodes[0], CoreFrame::Lit(Literal::LitInt(0)));
100    }
101
102    // 2. test_dce_live_let_preserved: let x = 42 in x -> unchanged. Binder Once, kept.
103    #[test]
104    fn test_dce_live_let_preserved() {
105        let x = VarId(1);
106        let expr = tree(vec![
107            CoreFrame::Lit(Literal::LitInt(42)), // 0: rhs
108            CoreFrame::Var(x),                   // 1: body
109            CoreFrame::LetNonRec {
110                binder: x,
111                rhs: 0,
112                body: 1,
113            }, // 2: root
114        ]);
115        let mut dce_expr = expr.clone();
116        let changed = Dce.run(&mut dce_expr);
117        assert!(!changed);
118        assert_eq!(dce_expr, expr);
119    }
120
121    // 3. test_dce_letrec_all_dead: letrec { f = 1; g = 2 } in 0 -> 0. All Dead, drop entire group.
122    #[test]
123    fn test_dce_letrec_all_dead() {
124        let f = VarId(1);
125        let g = VarId(2);
126        let expr = tree(vec![
127            CoreFrame::Lit(Literal::LitInt(1)), // 0: f's rhs
128            CoreFrame::Lit(Literal::LitInt(2)), // 1: g's rhs
129            CoreFrame::Lit(Literal::LitInt(0)), // 2: body
130            CoreFrame::LetRec {
131                bindings: vec![(f, 0), (g, 1)],
132                body: 2,
133            }, // 3: root
134        ]);
135        let mut dce_expr = expr;
136        let changed = Dce.run(&mut dce_expr);
137        assert!(changed);
138        assert_eq!(dce_expr.nodes.len(), 1);
139        assert_eq!(dce_expr.nodes[0], CoreFrame::Lit(Literal::LitInt(0)));
140    }
141
142    // 4. test_dce_letrec_one_live: letrec { f = g; g = 1 } in f -> unchanged.
143    // f is Once (live), keep entire group even though g might be referenced only by f.
144    #[test]
145    fn test_dce_letrec_one_live() {
146        let f = VarId(1);
147        let g = VarId(2);
148        let expr = tree(vec![
149            CoreFrame::Var(g),                  // 0: f's rhs
150            CoreFrame::Lit(Literal::LitInt(1)), // 1: g's rhs
151            CoreFrame::Var(f),                  // 2: body
152            CoreFrame::LetRec {
153                bindings: vec![(f, 0), (g, 1)],
154                body: 2,
155            }, // 3: root
156        ]);
157        let mut dce_expr = expr.clone();
158        let changed = Dce.run(&mut dce_expr);
159        assert!(!changed);
160        assert_eq!(dce_expr, expr);
161    }
162
163    // 5. test_dce_nested: let x = 42 in let y = 0 in x -> after DCE drops y's let, result is let x = 42 in x.
164    #[test]
165    fn test_dce_nested() {
166        let x = VarId(1);
167        let y = VarId(2);
168        let expr = tree(vec![
169            CoreFrame::Lit(Literal::LitInt(42)), // 0: x's rhs
170            CoreFrame::Lit(Literal::LitInt(0)),  // 1: y's rhs
171            CoreFrame::Var(x),                   // 2: y's body
172            CoreFrame::LetNonRec {
173                binder: y,
174                rhs: 1,
175                body: 2,
176            }, // 3: x's body
177            CoreFrame::LetNonRec {
178                binder: x,
179                rhs: 0,
180                body: 3,
181            }, // 4: root
182        ]);
183        let mut dce_expr = expr;
184        let changed = Dce.run(&mut dce_expr);
185        assert!(changed);
186        // Should have dropped y
187        // let x = 42 in x
188        assert_eq!(dce_expr.nodes.len(), 3);
189        // The root should be a LetNonRec for x
190        let root_idx = dce_expr.nodes.len() - 1;
191        if let CoreFrame::LetNonRec { binder, .. } = &dce_expr.nodes[root_idx] {
192            assert_eq!(*binder, x);
193        } else {
194            panic!(
195                "Root should be LetNonRec for x, got {:?}",
196                dce_expr.nodes[root_idx]
197            );
198        }
199    }
200
201    // 6. test_dce_preserves_eval: let x = 42 in let y = 99 in x -> eval before/after, verify match.
202    #[test]
203    fn test_dce_preserves_eval() {
204        let x = VarId(1);
205        let y = VarId(2);
206        let expr = tree(vec![
207            CoreFrame::Lit(Literal::LitInt(42)), // 0: x's rhs
208            CoreFrame::Lit(Literal::LitInt(99)), // 1: y's rhs
209            CoreFrame::Var(x),                   // 2: y's body
210            CoreFrame::LetNonRec {
211                binder: y,
212                rhs: 1,
213                body: 2,
214            }, // 3: x's body
215            CoreFrame::LetNonRec {
216                binder: x,
217                rhs: 0,
218                body: 3,
219            }, // 4: root
220        ]);
221        let mut dce_expr = expr.clone();
222
223        let mut heap = VecHeap::new();
224        let env = Env::new();
225
226        let val_orig = eval(&expr, &env, &mut heap).expect("Original eval failed");
227
228        let changed = Dce.run(&mut dce_expr);
229        assert!(changed);
230
231        let val_dce = eval(&dce_expr, &env, &mut heap).expect("DCE eval failed");
232
233        match (val_orig, val_dce) {
234            (tidepool_eval::Value::Lit(l1), tidepool_eval::Value::Lit(l2)) => assert_eq!(l1, l2),
235            _ => panic!("Expected literals"),
236        }
237    }
238}
tidepool_optimize/dce.rs

tidepool_optimize/
dce.rs
