nxpu_opt/
const_fold.rs

1//! Constant folding pass.
2//!
3//! Evaluates binary and unary operations on literal operands at compile time,
4//! replacing them with the resulting literal in the expression arena.
5
6use nxpu_ir::{Arena, BinaryOp, Expression, Handle, Literal, MathFunction, Module, UnaryOp};
7
8use crate::Pass;
9
10/// Folds constant expressions (literal operands) at compile time.
11#[derive(Debug)]
12pub struct ConstantFolding;
13
14impl Pass for ConstantFolding {
15    fn name(&self) -> &str {
16        "const-fold"
17    }
18
19    fn run(&self, module: &mut Module) -> bool {
20        let mut changed = false;
21        changed |= fold_expression_arena(&mut module.global_expressions);
22        for (_, func) in module.functions.iter_mut() {
23            changed |= fold_expression_arena(&mut func.expressions);
24        }
25        for ep in &mut module.entry_points {
26            changed |= fold_expression_arena(&mut ep.function.expressions);
27        }
28        changed
29    }
30}
31
32fn fold_expression_arena(arena: &mut Arena<Expression>) -> bool {
33    let mut changed = false;
34
35    let handles: Vec<Handle<Expression>> = arena.iter().map(|(h, _)| h).collect();
36
37    for handle in handles {
38        let replacement = match &arena[handle] {
39            Expression::Binary { op, left, right } => {
40                let left_val = &arena[*left];
41                let right_val = &arena[*right];
42                if let (Expression::Literal(l), Expression::Literal(r)) = (left_val, right_val) {
43                    fold_binary(*op, *l, *r).map(Expression::Literal)
44                } else {
45                    None
46                }
47            }
48            Expression::Unary { op, expr } => {
49                if let Expression::Literal(lit) = &arena[*expr] {
50                    fold_unary(*op, *lit).map(Expression::Literal)
51                } else {
52                    None
53                }
54            }
55            Expression::Math {
56                fun,
57                arg,
58                arg1,
59                arg2,
60                arg3: _,
61            } => fold_math(
62                *fun,
63                &arena[*arg],
64                arg1.map(|h| &arena[h]),
65                arg2.map(|h| &arena[h]),
66            )
67            .map(Expression::Literal),
68            _ => None,
69        };
70
71        if let Some(new_expr) = replacement {
72            arena[handle] = new_expr;
73            changed = true;
74        }
75    }
76
77    changed
78}
79
80fn fold_binary(op: BinaryOp, left: Literal, right: Literal) -> Option<Literal> {
81    match (left, right) {
82        (Literal::F32(l), Literal::F32(r)) => fold_f32(op, l, r),
83        (Literal::I32(l), Literal::I32(r)) => fold_i32(op, l, r),
84        (Literal::U32(l), Literal::U32(r)) => fold_u32(op, l, r),
85        (Literal::Bool(l), Literal::Bool(r)) => fold_bool(op, l, r),
86        _ => None,
87    }
88}
89
90/// Returns `Some(v)` only if `v` is finite (not NaN or Infinity).
91/// Non-finite results are left unfolded for correct runtime evaluation.
92fn finite(v: f32) -> Option<f32> {
93    if v.is_finite() { Some(v) } else { None }
94}
95
96fn fold_f32(op: BinaryOp, l: f32, r: f32) -> Option<Literal> {
97    match op {
98        BinaryOp::Add => finite(l + r).map(Literal::F32),
99        BinaryOp::Subtract => finite(l - r).map(Literal::F32),
100        BinaryOp::Multiply => finite(l * r).map(Literal::F32),
101        BinaryOp::Divide => finite(l / r).map(Literal::F32),
102        BinaryOp::Modulo => finite(l % r).map(Literal::F32),
103        BinaryOp::Equal => Some(Literal::Bool(l == r)),
104        BinaryOp::NotEqual => Some(Literal::Bool(l != r)),
105        BinaryOp::Less => Some(Literal::Bool(l < r)),
106        BinaryOp::LessEqual => Some(Literal::Bool(l <= r)),
107        BinaryOp::Greater => Some(Literal::Bool(l > r)),
108        BinaryOp::GreaterEqual => Some(Literal::Bool(l >= r)),
109        _ => None,
110    }
111}
112
113fn fold_i32(op: BinaryOp, l: i32, r: i32) -> Option<Literal> {
114    match op {
115        BinaryOp::Add => Some(Literal::I32(l.wrapping_add(r))),
116        BinaryOp::Subtract => Some(Literal::I32(l.wrapping_sub(r))),
117        BinaryOp::Multiply => Some(Literal::I32(l.wrapping_mul(r))),
118        BinaryOp::Divide if r != 0 => Some(Literal::I32(l.wrapping_div(r))),
119        BinaryOp::Modulo if r != 0 => Some(Literal::I32(l.wrapping_rem(r))),
120        BinaryOp::Equal => Some(Literal::Bool(l == r)),
121        BinaryOp::NotEqual => Some(Literal::Bool(l != r)),
122        BinaryOp::Less => Some(Literal::Bool(l < r)),
123        BinaryOp::LessEqual => Some(Literal::Bool(l <= r)),
124        BinaryOp::Greater => Some(Literal::Bool(l > r)),
125        BinaryOp::GreaterEqual => Some(Literal::Bool(l >= r)),
126        BinaryOp::BitwiseAnd => Some(Literal::I32(l & r)),
127        BinaryOp::BitwiseOr => Some(Literal::I32(l | r)),
128        BinaryOp::BitwiseXor => Some(Literal::I32(l ^ r)),
129        BinaryOp::ShiftLeft => Some(Literal::I32(l.wrapping_shl(r as u32))),
130        BinaryOp::ShiftRight => Some(Literal::I32(l.wrapping_shr(r as u32))),
131        _ => None,
132    }
133}
134
135fn fold_u32(op: BinaryOp, l: u32, r: u32) -> Option<Literal> {
136    match op {
137        BinaryOp::Add => Some(Literal::U32(l.wrapping_add(r))),
138        BinaryOp::Subtract => Some(Literal::U32(l.wrapping_sub(r))),
139        BinaryOp::Multiply => Some(Literal::U32(l.wrapping_mul(r))),
140        BinaryOp::Divide if r != 0 => Some(Literal::U32(l / r)),
141        BinaryOp::Modulo if r != 0 => Some(Literal::U32(l % r)),
142        BinaryOp::Equal => Some(Literal::Bool(l == r)),
143        BinaryOp::NotEqual => Some(Literal::Bool(l != r)),
144        BinaryOp::Less => Some(Literal::Bool(l < r)),
145        BinaryOp::LessEqual => Some(Literal::Bool(l <= r)),
146        BinaryOp::Greater => Some(Literal::Bool(l > r)),
147        BinaryOp::GreaterEqual => Some(Literal::Bool(l >= r)),
148        BinaryOp::BitwiseAnd => Some(Literal::U32(l & r)),
149        BinaryOp::BitwiseOr => Some(Literal::U32(l | r)),
150        BinaryOp::BitwiseXor => Some(Literal::U32(l ^ r)),
151        BinaryOp::ShiftLeft => Some(Literal::U32(l.wrapping_shl(r))),
152        BinaryOp::ShiftRight => Some(Literal::U32(l.wrapping_shr(r))),
153        _ => None,
154    }
155}
156
157fn fold_bool(op: BinaryOp, l: bool, r: bool) -> Option<Literal> {
158    match op {
159        BinaryOp::Equal => Some(Literal::Bool(l == r)),
160        BinaryOp::NotEqual => Some(Literal::Bool(l != r)),
161        BinaryOp::LogicalAnd => Some(Literal::Bool(l && r)),
162        BinaryOp::LogicalOr => Some(Literal::Bool(l || r)),
163        _ => None,
164    }
165}
166
167fn fold_unary(op: UnaryOp, lit: Literal) -> Option<Literal> {
168    match (op, lit) {
169        (UnaryOp::Negate, Literal::F32(v)) => finite(-v).map(Literal::F32),
170        (UnaryOp::Negate, Literal::I32(v)) => Some(Literal::I32(v.wrapping_neg())),
171        (UnaryOp::LogicalNot, Literal::Bool(v)) => Some(Literal::Bool(!v)),
172        (UnaryOp::BitwiseNot, Literal::I32(v)) => Some(Literal::I32(!v)),
173        (UnaryOp::BitwiseNot, Literal::U32(v)) => Some(Literal::U32(!v)),
174        _ => None,
175    }
176}
177
178fn fold_math(
179    fun: MathFunction,
180    arg: &Expression,
181    arg1: Option<&Expression>,
182    arg2: Option<&Expression>,
183) -> Option<Literal> {
184    // Only fold f32 scalar literals; skip vector ops.
185    let a = match arg {
186        Expression::Literal(Literal::F32(v)) => *v,
187        _ => return None,
188    };
189    let b = arg1.and_then(|e| match e {
190        Expression::Literal(Literal::F32(v)) => Some(*v),
191        _ => None,
192    });
193    let c = arg2.and_then(|e| match e {
194        Expression::Literal(Literal::F32(v)) => Some(*v),
195        _ => None,
196    });
197
198    let result = match fun {
199        // 1-arg
200        MathFunction::Abs => Some(a.abs()),
201        MathFunction::Floor => Some(a.floor()),
202        MathFunction::Ceil => Some(a.ceil()),
203        MathFunction::Round => Some(a.round()),
204        MathFunction::Fract => Some(a.fract()),
205        MathFunction::Trunc => Some(a.trunc()),
206        MathFunction::Sqrt => Some(a.sqrt()),
207        MathFunction::InverseSqrt => Some(1.0 / a.sqrt()),
208        MathFunction::Log => Some(a.ln()),
209        MathFunction::Log2 => Some(a.log2()),
210        MathFunction::Exp => Some(a.exp()),
211        MathFunction::Exp2 => Some(a.exp2()),
212        MathFunction::Sin => Some(a.sin()),
213        MathFunction::Cos => Some(a.cos()),
214        MathFunction::Tan => Some(a.tan()),
215        MathFunction::Asin => Some(a.asin()),
216        MathFunction::Acos => Some(a.acos()),
217        MathFunction::Atan => Some(a.atan()),
218        MathFunction::Sinh => Some(a.sinh()),
219        MathFunction::Cosh => Some(a.cosh()),
220        MathFunction::Tanh => Some(a.tanh()),
221        MathFunction::Saturate => Some(a.clamp(0.0, 1.0)),
222        // 2-arg
223        MathFunction::Min => Some(a.min(b?)),
224        MathFunction::Max => Some(a.max(b?)),
225        MathFunction::Pow => Some(a.powf(b?)),
226        MathFunction::Atan2 => Some(a.atan2(b?)),
227        MathFunction::Step => {
228            // step(edge, x) returns 0.0 if x < edge, else 1.0.
229            // In IR: arg=edge(a), arg1=x(b).
230            let x = b?;
231            Some(if x < a { 0.0 } else { 1.0 })
232        }
233        // 3-arg
234        MathFunction::Clamp => {
235            let lo = b?;
236            let hi = c?;
237            Some(a.clamp(lo, hi))
238        }
239        MathFunction::Mix => {
240            let y = b?;
241            let t = c?;
242            Some(a * (1.0 - t) + y * t)
243        }
244        MathFunction::Fma => {
245            let mb = b?;
246            let mc = c?;
247            Some(a.mul_add(mb, mc))
248        }
249        // Vector ops — skip
250        MathFunction::Dot
251        | MathFunction::Cross
252        | MathFunction::Normalize
253        | MathFunction::Length
254        | MathFunction::Distance
255        | MathFunction::SmoothStep => None,
256        // Bit manipulation — this folder works in f32, and these are integer
257        // operations. Folding them would mean deciding what extractBits of a
258        // float means, which is a question with no useful answer.
259        MathFunction::ExtractBits | MathFunction::InsertBits => None,
260    };
261
262    result.and_then(finite).map(Literal::F32)
263}
264
265#[cfg(test)]
266mod bit_op_folding_tests {
267    use super::*;
268    use nxpu_ir::{Expression, Function, Literal, MathFunction};
269
270    #[test]
271    fn bit_manipulation_is_not_folded() {
272        // This folder works in f32. `extractBits` of a float is a question
273        // with no useful answer, so it declines rather than inventing one.
274        let mut func = Function::new("test");
275        let a = func
276            .expressions
277            .append(Expression::Literal(Literal::F32(8.0)));
278        let h = func.expressions.append(Expression::Math {
279            fun: MathFunction::ExtractBits,
280            arg: a,
281            arg1: Some(a),
282            arg2: Some(a),
283            arg3: None,
284        });
285        let before = format!("{:?}", func.expressions[h]);
286        fold_expression_arena(&mut func.expressions);
287        assert_eq!(format!("{:?}", func.expressions[h]), before);
288    }
289}
290
291#[cfg(test)]
292mod tests {
293    use super::*;
294    use nxpu_ir::{Function, Literal};
295
296    #[test]
297    fn fold_f32_add() {
298        let mut func = Function::new("test");
299        let one = func
300            .expressions
301            .append(Expression::Literal(Literal::F32(1.0)));
302        let two = func
303            .expressions
304            .append(Expression::Literal(Literal::F32(2.0)));
305        let add = func.expressions.append(Expression::Binary {
306            op: BinaryOp::Add,
307            left: one,
308            right: two,
309        });
310
311        let changed = fold_expression_arena(&mut func.expressions);
312        assert!(changed);
313        match &func.expressions[add] {
314            Expression::Literal(Literal::F32(v)) => assert_eq!(*v, 3.0),
315            other => panic!("expected Literal(F32(3.0)), got {other:?}"),
316        }
317    }
318
319    #[test]
320    fn fold_i32_multiply() {
321        let mut func = Function::new("test");
322        let three = func
323            .expressions
324            .append(Expression::Literal(Literal::I32(3)));
325        let four = func
326            .expressions
327            .append(Expression::Literal(Literal::I32(4)));
328        let mul = func.expressions.append(Expression::Binary {
329            op: BinaryOp::Multiply,
330            left: three,
331            right: four,
332        });
333
334        let changed = fold_expression_arena(&mut func.expressions);
335        assert!(changed);
336        match &func.expressions[mul] {
337            Expression::Literal(Literal::I32(v)) => assert_eq!(*v, 12),
338            other => panic!("expected Literal(I32(12)), got {other:?}"),
339        }
340    }
341
342    #[test]
343    fn fold_unary_negate() {
344        let mut func = Function::new("test");
345        let five = func
346            .expressions
347            .append(Expression::Literal(Literal::F32(5.0)));
348        let neg = func.expressions.append(Expression::Unary {
349            op: UnaryOp::Negate,
350            expr: five,
351        });
352
353        let changed = fold_expression_arena(&mut func.expressions);
354        assert!(changed);
355        match &func.expressions[neg] {
356            Expression::Literal(Literal::F32(v)) => assert_eq!(*v, -5.0),
357            other => panic!("expected Literal(F32(-5.0)), got {other:?}"),
358        }
359    }
360
361    #[test]
362    fn no_fold_non_literal_operands() {
363        let mut func = Function::new("test");
364        let arg = func.expressions.append(Expression::FunctionArgument(0));
365        let lit = func
366            .expressions
367            .append(Expression::Literal(Literal::F32(2.0)));
368        let _add = func.expressions.append(Expression::Binary {
369            op: BinaryOp::Add,
370            left: arg,
371            right: lit,
372        });
373
374        let changed = fold_expression_arena(&mut func.expressions);
375        assert!(!changed);
376    }
377
378    #[test]
379    fn fold_comparison() {
380        let mut func = Function::new("test");
381        let a = func
382            .expressions
383            .append(Expression::Literal(Literal::U32(5)));
384        let b = func
385            .expressions
386            .append(Expression::Literal(Literal::U32(3)));
387        let cmp = func.expressions.append(Expression::Binary {
388            op: BinaryOp::Greater,
389            left: a,
390            right: b,
391        });
392
393        let changed = fold_expression_arena(&mut func.expressions);
394        assert!(changed);
395        match &func.expressions[cmp] {
396            Expression::Literal(Literal::Bool(v)) => assert!(*v),
397            other => panic!("expected Literal(Bool(true)), got {other:?}"),
398        }
399    }
400
401    #[test]
402    fn cascade_folding() {
403        // Tests that a*b -> literal, then (a*b)+c -> literal in the same pass.
404        let mut func = Function::new("test");
405        let a = func
406            .expressions
407            .append(Expression::Literal(Literal::F32(2.0)));
408        let b = func
409            .expressions
410            .append(Expression::Literal(Literal::F32(3.0)));
411        let mul = func.expressions.append(Expression::Binary {
412            op: BinaryOp::Multiply,
413            left: a,
414            right: b,
415        });
416        let c = func
417            .expressions
418            .append(Expression::Literal(Literal::F32(4.0)));
419        let add = func.expressions.append(Expression::Binary {
420            op: BinaryOp::Add,
421            left: mul,
422            right: c,
423        });
424
425        let changed = fold_expression_arena(&mut func.expressions);
426        assert!(changed);
427        // mul should be folded to 6.0, then add should be folded to 10.0.
428        match &func.expressions[add] {
429            Expression::Literal(Literal::F32(v)) => assert_eq!(*v, 10.0),
430            other => panic!("expected Literal(F32(10.0)), got {other:?}"),
431        }
432    }
433
434    // --- Math expression folding tests ---
435
436    fn make_math1(func: &mut Function, fun: MathFunction, val: f32) -> Handle<Expression> {
437        let arg = func
438            .expressions
439            .append(Expression::Literal(Literal::F32(val)));
440        func.expressions.append(Expression::Math {
441            fun,
442            arg,
443            arg1: None,
444            arg2: None,
445            arg3: None,
446        })
447    }
448
449    fn make_math2(func: &mut Function, fun: MathFunction, a: f32, b: f32) -> Handle<Expression> {
450        let arg = func
451            .expressions
452            .append(Expression::Literal(Literal::F32(a)));
453        let arg1 = func
454            .expressions
455            .append(Expression::Literal(Literal::F32(b)));
456        func.expressions.append(Expression::Math {
457            fun,
458            arg,
459            arg1: Some(arg1),
460            arg2: None,
461            arg3: None,
462        })
463    }
464
465    fn make_math3(
466        func: &mut Function,
467        fun: MathFunction,
468        a: f32,
469        b: f32,
470        c: f32,
471    ) -> Handle<Expression> {
472        let arg = func
473            .expressions
474            .append(Expression::Literal(Literal::F32(a)));
475        let arg1 = func
476            .expressions
477            .append(Expression::Literal(Literal::F32(b)));
478        let arg2 = func
479            .expressions
480            .append(Expression::Literal(Literal::F32(c)));
481        func.expressions.append(Expression::Math {
482            fun,
483            arg,
484            arg1: Some(arg1),
485            arg2: Some(arg2),
486            arg3: None,
487        })
488    }
489
490    #[test]
491    fn fold_math_abs() {
492        let mut func = Function::new("test");
493        let h = make_math1(&mut func, MathFunction::Abs, -3.0);
494        let changed = fold_expression_arena(&mut func.expressions);
495        assert!(changed);
496        match &func.expressions[h] {
497            Expression::Literal(Literal::F32(v)) => assert_eq!(*v, 3.0),
498            other => panic!("expected 3.0, got {other:?}"),
499        }
500    }
501
502    #[test]
503    fn fold_math_sqrt() {
504        let mut func = Function::new("test");
505        let h = make_math1(&mut func, MathFunction::Sqrt, 9.0);
506        let changed = fold_expression_arena(&mut func.expressions);
507        assert!(changed);
508        match &func.expressions[h] {
509            Expression::Literal(Literal::F32(v)) => assert_eq!(*v, 3.0),
510            other => panic!("expected 3.0, got {other:?}"),
511        }
512    }
513
514    #[test]
515    fn fold_math_min_max() {
516        let mut func = Function::new("test");
517        let min_h = make_math2(&mut func, MathFunction::Min, 2.0, 5.0);
518        let max_h = make_math2(&mut func, MathFunction::Max, 2.0, 5.0);
519        let changed = fold_expression_arena(&mut func.expressions);
520        assert!(changed);
521        match &func.expressions[min_h] {
522            Expression::Literal(Literal::F32(v)) => assert_eq!(*v, 2.0),
523            other => panic!("expected 2.0, got {other:?}"),
524        }
525        match &func.expressions[max_h] {
526            Expression::Literal(Literal::F32(v)) => assert_eq!(*v, 5.0),
527            other => panic!("expected 5.0, got {other:?}"),
528        }
529    }
530
531    #[test]
532    fn fold_math_clamp() {
533        let mut func = Function::new("test");
534        let h = make_math3(&mut func, MathFunction::Clamp, 10.0, 0.0, 5.0);
535        let changed = fold_expression_arena(&mut func.expressions);
536        assert!(changed);
537        match &func.expressions[h] {
538            Expression::Literal(Literal::F32(v)) => assert_eq!(*v, 5.0),
539            other => panic!("expected 5.0, got {other:?}"),
540        }
541    }
542
543    #[test]
544    fn fold_math_trig() {
545        let mut func = Function::new("test");
546        let h = make_math1(&mut func, MathFunction::Sin, 0.0);
547        let changed = fold_expression_arena(&mut func.expressions);
548        assert!(changed);
549        match &func.expressions[h] {
550            Expression::Literal(Literal::F32(v)) => assert!(v.abs() < 1e-6),
551            other => panic!("expected ~0.0, got {other:?}"),
552        }
553    }
554
555    #[test]
556    fn fold_math_fma() {
557        let mut func = Function::new("test");
558        // fma(2.0, 3.0, 4.0) = 2*3 + 4 = 10
559        let h = make_math3(&mut func, MathFunction::Fma, 2.0, 3.0, 4.0);
560        let changed = fold_expression_arena(&mut func.expressions);
561        assert!(changed);
562        match &func.expressions[h] {
563            Expression::Literal(Literal::F32(v)) => assert_eq!(*v, 10.0),
564            other => panic!("expected 10.0, got {other:?}"),
565        }
566    }
567
568    #[test]
569    fn no_fold_f32_div_by_zero() {
570        let mut func = Function::new("test");
571        let a = func
572            .expressions
573            .append(Expression::Literal(Literal::F32(1.0)));
574        let b = func
575            .expressions
576            .append(Expression::Literal(Literal::F32(0.0)));
577        let div = func.expressions.append(Expression::Binary {
578            op: BinaryOp::Divide,
579            left: a,
580            right: b,
581        });
582
583        let changed = fold_expression_arena(&mut func.expressions);
584        assert!(!changed);
585        // Should remain a Binary expression (not folded to Inf).
586        assert!(matches!(&func.expressions[div], Expression::Binary { .. }));
587    }
588
589    #[test]
590    fn no_fold_sqrt_negative() {
591        let mut func = Function::new("test");
592        let h = make_math1(&mut func, MathFunction::Sqrt, -1.0);
593        let changed = fold_expression_arena(&mut func.expressions);
594        assert!(!changed);
595        assert!(matches!(&func.expressions[h], Expression::Math { .. }));
596    }
597
598    #[test]
599    fn no_fold_log_zero() {
600        let mut func = Function::new("test");
601        let h = make_math1(&mut func, MathFunction::Log, 0.0);
602        let changed = fold_expression_arena(&mut func.expressions);
603        assert!(!changed);
604        assert!(matches!(&func.expressions[h], Expression::Math { .. }));
605    }
606
607    #[test]
608    fn no_fold_math_non_literal() {
609        let mut func = Function::new("test");
610        let arg = func.expressions.append(Expression::FunctionArgument(0));
611        let _h = func.expressions.append(Expression::Math {
612            fun: MathFunction::Abs,
613            arg,
614            arg1: None,
615            arg2: None,
616            arg3: None,
617        });
618        let changed = fold_expression_arena(&mut func.expressions);
619        assert!(!changed);
620    }
621
622    #[test]
623    fn fold_global_expression() {
624        use nxpu_ir::Module;
625
626        let mut module = Module::default();
627        let one = module
628            .global_expressions
629            .append(Expression::Literal(Literal::F32(2.0)));
630        let two = module
631            .global_expressions
632            .append(Expression::Literal(Literal::F32(3.0)));
633        let add = module.global_expressions.append(Expression::Binary {
634            op: BinaryOp::Add,
635            left: one,
636            right: two,
637        });
638
639        let pass = ConstantFolding;
640        let changed = pass.run(&mut module);
641        assert!(changed);
642        match &module.global_expressions[add] {
643            Expression::Literal(Literal::F32(v)) => assert_eq!(*v, 5.0),
644            other => panic!("expected Literal(F32(5.0)), got {other:?}"),
645        }
646    }
647}
nxpu_opt/const_fold.rs

nxpu_opt/
const_fold.rs
