nxpu_opt/
dce.rs

1//! Dead code elimination pass.
2//!
3//! Removes `Emit` statements whose expression ranges contain no
4//! expressions referenced (directly or transitively) by any other statement.
5
6use std::collections::{HashMap, HashSet};
7
8use nxpu_ir::{Arena, Expression, Function, Handle, LocalVariable, Module, Statement};
9
10use crate::Pass;
11
12/// Removes unused `Emit` statements and unreferenced local variables from
13/// function bodies.
14#[derive(Debug)]
15pub struct DeadCodeElimination;
16
17impl Pass for DeadCodeElimination {
18    fn name(&self) -> &str {
19        "dce"
20    }
21
22    fn run(&self, module: &mut Module) -> bool {
23        let mut changed = false;
24        for (_, func) in module.functions.iter_mut() {
25            changed |= run_on_function(func);
26            changed |= remove_dead_locals(func);
27        }
28        for ep in &mut module.entry_points {
29            changed |= run_on_function(&mut ep.function);
30            changed |= remove_dead_locals(&mut ep.function);
31        }
32        changed
33    }
34}
35
36fn run_on_function(func: &mut Function) -> bool {
37    // Collect locals that are loaded (read) somewhere.
38    let loaded_locals = collect_loaded_locals(func);
39
40    // 1. Collect root expression handles from non-Emit statements.
41    let mut used: HashSet<Handle<Expression>> = HashSet::new();
42    collect_used_from_block(&func.body, &mut used, &loaded_locals, func);
43
44    // Also mark local variable init expressions as used.
45    for (_, local) in func.local_variables.iter() {
46        if let Some(init) = local.init {
47            used.insert(init);
48        }
49    }
50
51    // 2. Transitively mark operands of used expressions.
52    let mut worklist: Vec<Handle<Expression>> = used.iter().copied().collect();
53    while let Some(handle) = worklist.pop() {
54        if let Some(expr) = func.expressions.try_get(handle) {
55            for operand in expression_operands(expr) {
56                if used.insert(operand) {
57                    worklist.push(operand);
58                }
59            }
60        }
61    }
62
63    // 3. Pre-compute which pointers target dead locals.
64    let dead_store_ptrs: HashSet<Handle<Expression>> = func
65        .expressions
66        .iter()
67        .filter(|(_, expr)| {
68            matches!(expr, Expression::LocalVariable(lv) if !loaded_locals.contains(lv))
69        })
70        .map(|(h, _)| h)
71        .collect();
72
73    // 4. Filter out dead Emit/Store/Call statements.
74    filter_dead_in_block(&mut func.body, &used, &dead_store_ptrs)
75}
76
77/// Collect all local variables that are loaded (read) anywhere in the function.
78#[allow(clippy::collapsible_if)] // nested if-let for MSRV 1.87 compat (no let chains)
79fn collect_loaded_locals(func: &Function) -> HashSet<Handle<LocalVariable>> {
80    let mut loaded = HashSet::new();
81    for (_, expr) in func.expressions.iter() {
82        if let Expression::Load { pointer } = expr {
83            if let Some(Expression::LocalVariable(lv)) = func.expressions.try_get(*pointer) {
84                loaded.insert(*lv);
85            }
86        }
87    }
88    loaded
89}
90
91fn collect_used_from_block(
92    block: &[Statement],
93    used: &mut HashSet<Handle<Expression>>,
94    loaded_locals: &HashSet<Handle<LocalVariable>>,
95    func: &Function,
96) {
97    for stmt in block {
98        match stmt {
99            Statement::Emit(_) => {}
100            Statement::Store { pointer, value } => {
101                // Check if this store targets an unread local variable.
102                let is_dead_store = if let Some(Expression::LocalVariable(lv)) =
103                    func.expressions.try_get(*pointer)
104                {
105                    !loaded_locals.contains(lv)
106                } else {
107                    false
108                };
109                if !is_dead_store {
110                    used.insert(*pointer);
111                    used.insert(*value);
112                }
113            }
114            Statement::If {
115                condition,
116                accept,
117                reject,
118            } => {
119                used.insert(*condition);
120                collect_used_from_block(accept, used, loaded_locals, func);
121                collect_used_from_block(reject, used, loaded_locals, func);
122            }
123            Statement::Loop {
124                body,
125                continuing,
126                break_if,
127            } => {
128                collect_used_from_block(body, used, loaded_locals, func);
129                collect_used_from_block(continuing, used, loaded_locals, func);
130                if let Some(brk) = break_if {
131                    used.insert(*brk);
132                }
133            }
134            Statement::Call {
135                arguments, result, ..
136            } => {
137                // Always mark call arguments and results as used (conservative).
138                for arg in arguments {
139                    used.insert(*arg);
140                }
141                if let Some(r) = result {
142                    used.insert(*r);
143                }
144            }
145            Statement::Atomic {
146                pointer,
147                fun,
148                value,
149                result,
150            } => {
151                used.insert(*pointer);
152                used.insert(*value);
153                if let Some(r) = result {
154                    used.insert(*r);
155                }
156                if let nxpu_ir::AtomicFunction::Exchange {
157                    compare: Some(cmp), ..
158                } = fun
159                {
160                    used.insert(*cmp);
161                }
162            }
163            Statement::Return { value } => {
164                if let Some(v) = value {
165                    used.insert(*v);
166                }
167            }
168            Statement::Barrier(_) | Statement::Break | Statement::Continue => {}
169        }
170    }
171}
172
173/// Returns all expression handles directly referenced by an expression.
174pub(crate) fn expression_operands(expr: &Expression) -> Vec<Handle<Expression>> {
175    match expr {
176        Expression::Literal(_)
177        | Expression::FunctionArgument(_)
178        | Expression::GlobalVariable(_)
179        | Expression::LocalVariable(_)
180        | Expression::CallResult(_)
181        | Expression::AtomicResult { .. }
182        | Expression::ZeroValue(_) => vec![],
183
184        Expression::Load { pointer } => vec![*pointer],
185        Expression::Unary { expr, .. } => vec![*expr],
186        Expression::ArrayLength(e) => vec![*e],
187        Expression::Splat { value, .. } => vec![*value],
188        Expression::As { expr, .. } => vec![*expr],
189
190        Expression::Binary { left, right, .. } => vec![*left, *right],
191        Expression::Access { base, index } => vec![*base, *index],
192        Expression::AccessIndex { base, .. } => vec![*base],
193        Expression::Select {
194            condition,
195            accept,
196            reject,
197        } => vec![*condition, *accept, *reject],
198        Expression::Swizzle { vector, .. } => vec![*vector],
199
200        Expression::Compose { components, .. } => components.clone(),
201        Expression::Math {
202            arg,
203            arg1,
204            arg2,
205            arg3,
206            ..
207        } => {
208            let mut ops = vec![*arg];
209            if let Some(a) = arg1 {
210                ops.push(*a);
211            }
212            if let Some(a) = arg2 {
213                ops.push(*a);
214            }
215            if let Some(a) = arg3 {
216                ops.push(*a);
217            }
218            ops
219        }
220    }
221}
222
223fn filter_dead_in_block(
224    block: &mut Vec<Statement>,
225    used: &HashSet<Handle<Expression>>,
226    dead_store_ptrs: &HashSet<Handle<Expression>>,
227) -> bool {
228    let mut changed = false;
229    block.retain_mut(|stmt| match stmt {
230        Statement::Emit(range) => {
231            let has_used = used
232                .iter()
233                .any(|h| range.index_range().contains(&(h.index() as u32)));
234            if !has_used {
235                changed = true;
236                return false;
237            }
238            true
239        }
240        Statement::Store { pointer, .. } => {
241            // Remove stores to unread local variables.
242            if dead_store_ptrs.contains(pointer) {
243                changed = true;
244                return false;
245            }
246            true
247        }
248        Statement::If { accept, reject, .. } => {
249            changed |= filter_dead_in_block(accept, used, dead_store_ptrs);
250            changed |= filter_dead_in_block(reject, used, dead_store_ptrs);
251            true
252        }
253        Statement::Loop {
254            body, continuing, ..
255        } => {
256            changed |= filter_dead_in_block(body, used, dead_store_ptrs);
257            changed |= filter_dead_in_block(continuing, used, dead_store_ptrs);
258            true
259        }
260        _ => true,
261    });
262    changed
263}
264
265/// Removes local variables that are never referenced in any expression,
266/// and remaps `Expression::LocalVariable` handles accordingly.
267#[allow(clippy::collapsible_if)] // nested if-let for MSRV 1.87 compat (no let chains)
268fn remove_dead_locals(func: &mut Function) -> bool {
269    // Collect referenced local variable handles.
270    let mut referenced: HashSet<Handle<LocalVariable>> = HashSet::new();
271    for (_, expr) in func.expressions.iter() {
272        if let Expression::LocalVariable(lv) = expr {
273            referenced.insert(*lv);
274        }
275    }
276
277    let old_len = func.local_variables.len();
278    // Check if every local variable in the arena is referenced.
279    let all_referenced = func
280        .local_variables
281        .iter()
282        .all(|(h, _)| referenced.contains(&h));
283    if all_referenced {
284        return false;
285    }
286
287    // Build a new arena and a handle remap table.
288    let mut new_arena: Arena<LocalVariable> = Arena::new();
289    let mut remap: HashMap<Handle<LocalVariable>, Handle<LocalVariable>> = HashMap::new();
290
291    for (old_handle, local) in func.local_variables.iter() {
292        if referenced.contains(&old_handle) {
293            let new_handle = new_arena.append(local.clone());
294            remap.insert(old_handle, new_handle);
295        }
296    }
297
298    // Remap all Expression::LocalVariable handles.
299    let expr_handles: Vec<Handle<Expression>> = func.expressions.iter().map(|(h, _)| h).collect();
300    for h in expr_handles {
301        if let Expression::LocalVariable(lv) = &func.expressions[h] {
302            if let Some(&new_lv) = remap.get(lv) {
303                if new_lv != *lv {
304                    func.expressions[h] = Expression::LocalVariable(new_lv);
305                }
306            }
307        }
308    }
309
310    func.local_variables = new_arena;
311    func.local_variables.len() < old_len
312}
313
314#[cfg(test)]
315mod tests {
316    use super::*;
317    use nxpu_ir::{BinaryOp, Literal, Range, Statement};
318
319    #[test]
320    fn removes_unused_emit() {
321        let mut func = Function::new("test");
322        let lit_a = func
323            .expressions
324            .append(Expression::Literal(Literal::F32(1.0)));
325        let _lit_b = func
326            .expressions
327            .append(Expression::Literal(Literal::F32(2.0)));
328        let _lit_c = func
329            .expressions
330            .append(Expression::Literal(Literal::F32(3.0)));
331
332        // Emit all three.
333        func.body
334            .push(Statement::Emit(Range::from_index_range(0..3)));
335        // Only lit_a is used.
336        func.body.push(Statement::Return { value: Some(lit_a) });
337        // Emit a fourth unused expression.
338        let _lit_d = func
339            .expressions
340            .append(Expression::Literal(Literal::F32(99.0)));
341        func.body
342            .push(Statement::Emit(Range::from_index_range(3..4)));
343
344        let changed = run_on_function(&mut func);
345        assert!(changed);
346        // The second Emit (range 3..4) should be removed; the first stays because lit_a is used.
347        assert_eq!(func.body.len(), 2);
348    }
349
350    #[test]
351    fn keeps_transitively_used() {
352        let mut func = Function::new("test");
353        let a = func
354            .expressions
355            .append(Expression::Literal(Literal::F32(1.0)));
356        let b = func
357            .expressions
358            .append(Expression::Literal(Literal::F32(2.0)));
359        let add = func.expressions.append(Expression::Binary {
360            op: BinaryOp::Add,
361            left: a,
362            right: b,
363        });
364
365        func.body
366            .push(Statement::Emit(Range::from_index_range(0..3)));
367        func.body.push(Statement::Return { value: Some(add) });
368
369        let changed = run_on_function(&mut func);
370        // No change — all expressions are transitively used via `add`.
371        assert!(!changed);
372        assert_eq!(func.body.len(), 2);
373    }
374
375    #[test]
376    fn no_change_on_empty_function() {
377        let mut func = Function::new("test");
378        let changed = run_on_function(&mut func);
379        assert!(!changed);
380    }
381
382    fn dummy_type_handle() -> Handle<nxpu_ir::Type> {
383        let mut types = nxpu_ir::UniqueArena::new();
384        types.insert(nxpu_ir::Type {
385            name: None,
386            inner: nxpu_ir::TypeInner::Scalar(nxpu_ir::Scalar::F32),
387        })
388    }
389
390    fn dummy_gv_handle() -> Handle<nxpu_ir::GlobalVariable> {
391        let mut arena = Arena::new();
392        arena.append(nxpu_ir::GlobalVariable {
393            name: None,
394            space: nxpu_ir::AddressSpace::Storage {
395                access: nxpu_ir::StorageAccess::LOAD | nxpu_ir::StorageAccess::STORE,
396            },
397            binding: None,
398            ty: dummy_type_handle(),
399            init: None,
400            layout: None,
401        })
402    }
403
404    #[test]
405    fn removes_dead_store_to_unread_local() {
406        let mut func = Function::new("test");
407        let lv = func.local_variables.append(LocalVariable {
408            name: Some("temp".into()),
409            ty: dummy_type_handle(),
410            init: None,
411        });
412        let ptr = func.expressions.append(Expression::LocalVariable(lv));
413        let val = func
414            .expressions
415            .append(Expression::Literal(Literal::F32(42.0)));
416        // Store to the local (never read).
417        func.body.push(Statement::Store {
418            pointer: ptr,
419            value: val,
420        });
421        func.body.push(Statement::Return { value: None });
422
423        let changed = run_on_function(&mut func);
424        assert!(changed);
425        // The store should have been removed.
426        assert_eq!(func.body.len(), 1); // only Return remains
427    }
428
429    #[test]
430    fn keeps_store_to_read_local() {
431        let mut func = Function::new("test");
432        let lv = func.local_variables.append(LocalVariable {
433            name: Some("temp".into()),
434            ty: dummy_type_handle(),
435            init: None,
436        });
437        let ptr = func.expressions.append(Expression::LocalVariable(lv));
438        let val = func
439            .expressions
440            .append(Expression::Literal(Literal::F32(42.0)));
441        // Load the local (read).
442        let loaded = func.expressions.append(Expression::Load { pointer: ptr });
443        func.body.push(Statement::Emit(Range::from_index_range(
444            ptr.index() as u32..val.index() as u32 + 1,
445        )));
446        func.body.push(Statement::Store {
447            pointer: ptr,
448            value: val,
449        });
450        func.body.push(Statement::Return {
451            value: Some(loaded),
452        });
453
454        let changed = run_on_function(&mut func);
455        let store_count = func
456            .body
457            .iter()
458            .filter(|s| matches!(s, Statement::Store { .. }))
459            .count();
460        assert_eq!(store_count, 1);
461        let _ = changed;
462    }
463
464    #[test]
465    fn keeps_store_to_global() {
466        let mut func = Function::new("test");
467        let gv_handle = dummy_gv_handle();
468        let ptr = func
469            .expressions
470            .append(Expression::GlobalVariable(gv_handle));
471        let val = func
472            .expressions
473            .append(Expression::Literal(Literal::F32(42.0)));
474        func.body.push(Statement::Store {
475            pointer: ptr,
476            value: val,
477        });
478        func.body.push(Statement::Return { value: None });
479
480        let changed = run_on_function(&mut func);
481        let store_count = func
482            .body
483            .iter()
484            .filter(|s| matches!(s, Statement::Store { .. }))
485            .count();
486        assert_eq!(store_count, 1);
487        let _ = changed;
488    }
489
490    #[test]
491    fn removes_unreferenced_local() {
492        let mut func = Function::new("test");
493        // Add two locals; only reference one in expressions.
494        let _lv_unused = func.local_variables.append(LocalVariable {
495            name: Some("unused".into()),
496            ty: dummy_type_handle(),
497            init: None,
498        });
499        let lv_used = func.local_variables.append(LocalVariable {
500            name: Some("used".into()),
501            ty: dummy_type_handle(),
502            init: None,
503        });
504        // Reference only lv_used.
505        let _ptr = func.expressions.append(Expression::LocalVariable(lv_used));
506
507        assert_eq!(func.local_variables.len(), 2);
508        let changed = remove_dead_locals(&mut func);
509        assert!(changed);
510        assert_eq!(func.local_variables.len(), 1);
511    }
512
513    #[test]
514    fn keeps_all_referenced_locals() {
515        let mut func = Function::new("test");
516        let lv_a = func.local_variables.append(LocalVariable {
517            name: Some("a".into()),
518            ty: dummy_type_handle(),
519            init: None,
520        });
521        let lv_b = func.local_variables.append(LocalVariable {
522            name: Some("b".into()),
523            ty: dummy_type_handle(),
524            init: None,
525        });
526        let _ptr_a = func.expressions.append(Expression::LocalVariable(lv_a));
527        let _ptr_b = func.expressions.append(Expression::LocalVariable(lv_b));
528
529        assert_eq!(func.local_variables.len(), 2);
530        let changed = remove_dead_locals(&mut func);
531        assert!(!changed);
532        assert_eq!(func.local_variables.len(), 2);
533    }
534}
nxpu_opt/dce.rs

nxpu_opt/
dce.rs
