oxigdal_algorithms/dsl/
optimizer.rs

1//! Expression optimizer for the DSL
2//!
3//! This module provides various optimization passes:
4//! - Constant folding
5//! - Common subexpression elimination
6//! - Dead code elimination
7//! - Algebraic simplifications
8
9use super::ast::{BinaryOp, Expr, Program, Statement, UnaryOp};
10
11#[cfg(not(feature = "std"))]
12use alloc::{boxed::Box, collections::BTreeMap as HashMap, string::String, vec::Vec};
13
14#[cfg(feature = "std")]
15use std::collections::HashMap;
16
17/// Optimization level
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum OptLevel {
20    /// No optimization
21    None,
22    /// Basic optimizations (constant folding)
23    Basic,
24    /// Standard optimizations (basic + algebraic simplifications)
25    Standard,
26    /// Aggressive optimizations (standard + CSE + DCE)
27    Aggressive,
28}
29
30/// Optimizer for DSL programs
31pub struct Optimizer {
32    level: OptLevel,
33}
34
35impl Default for Optimizer {
36    fn default() -> Self {
37        Self::new(OptLevel::Standard)
38    }
39}
40
41impl Optimizer {
42    /// Creates a new optimizer with the given optimization level
43    pub fn new(level: OptLevel) -> Self {
44        Self { level }
45    }
46
47    /// Optimizes a program
48    pub fn optimize_program(&self, mut program: Program) -> Program {
49        if self.level == OptLevel::None {
50            return program;
51        }
52
53        program.statements = program
54            .statements
55            .into_iter()
56            .map(|stmt| self.optimize_statement(stmt))
57            .collect();
58
59        program
60    }
61
62    /// Optimizes a single statement
63    pub fn optimize_statement(&self, stmt: Statement) -> Statement {
64        match stmt {
65            Statement::VariableDecl { name, value } => Statement::VariableDecl {
66                name,
67                value: Box::new(self.optimize_expr(*value)),
68            },
69            Statement::FunctionDecl { name, params, body } => Statement::FunctionDecl {
70                name,
71                params,
72                body: Box::new(self.optimize_expr(*body)),
73            },
74            Statement::Return(expr) => Statement::Return(Box::new(self.optimize_expr(*expr))),
75            Statement::Expr(expr) => Statement::Expr(Box::new(self.optimize_expr(*expr))),
76        }
77    }
78
79    /// Optimizes an expression
80    pub fn optimize_expr(&self, expr: Expr) -> Expr {
81        if self.level == OptLevel::None {
82            return expr;
83        }
84
85        let mut optimized = expr;
86
87        // Apply constant folding
88        optimized = self.constant_fold(optimized);
89
90        // Apply algebraic simplifications
91        if matches!(self.level, OptLevel::Standard | OptLevel::Aggressive) {
92            optimized = self.algebraic_simplify(optimized);
93        }
94
95        // Apply common subexpression elimination
96        if self.level == OptLevel::Aggressive {
97            optimized = self.eliminate_common_subexpressions(optimized);
98        }
99
100        optimized
101    }
102
103    /// Performs constant folding
104    fn constant_fold(&self, expr: Expr) -> Expr {
105        match expr {
106            Expr::Binary {
107                left,
108                op,
109                right,
110                ty,
111            } => {
112                let left_opt = self.constant_fold(*left);
113                let right_opt = self.constant_fold(*right);
114
115                if let (Expr::Number(l), Expr::Number(r)) = (&left_opt, &right_opt) {
116                    if let Some(result) = self.eval_const_binary(*l, op, *r) {
117                        return Expr::Number(result);
118                    }
119                }
120
121                Expr::Binary {
122                    left: Box::new(left_opt),
123                    op,
124                    right: Box::new(right_opt),
125                    ty,
126                }
127            }
128            Expr::Unary {
129                op,
130                expr: inner,
131                ty,
132            } => {
133                let inner_opt = self.constant_fold(*inner);
134
135                if let Expr::Number(n) = &inner_opt {
136                    if let Some(result) = self.eval_const_unary(op, *n) {
137                        return Expr::Number(result);
138                    }
139                }
140
141                Expr::Unary {
142                    op,
143                    expr: Box::new(inner_opt),
144                    ty,
145                }
146            }
147            Expr::Conditional {
148                condition,
149                then_expr,
150                else_expr,
151                ty,
152            } => {
153                let cond_opt = self.constant_fold(*condition);
154
155                // If condition is constant, return only the taken branch
156                if let Expr::Number(n) = &cond_opt {
157                    if n.abs() > f64::EPSILON {
158                        return self.constant_fold(*then_expr);
159                    } else {
160                        return self.constant_fold(*else_expr);
161                    }
162                }
163
164                Expr::Conditional {
165                    condition: Box::new(cond_opt),
166                    then_expr: Box::new(self.constant_fold(*then_expr)),
167                    else_expr: Box::new(self.constant_fold(*else_expr)),
168                    ty,
169                }
170            }
171            Expr::Call { name, args, ty } => Expr::Call {
172                name,
173                args: args
174                    .into_iter()
175                    .map(|arg| self.constant_fold(arg))
176                    .collect(),
177                ty,
178            },
179            Expr::Block {
180                statements,
181                result,
182                ty,
183            } => Expr::Block {
184                statements: statements
185                    .into_iter()
186                    .map(|stmt| self.optimize_statement(stmt))
187                    .collect(),
188                result: result.map(|r| Box::new(self.constant_fold(*r))),
189                ty,
190            },
191            _ => expr,
192        }
193    }
194
195    /// Evaluates a constant binary operation
196    fn eval_const_binary(&self, left: f64, op: BinaryOp, right: f64) -> Option<f64> {
197        let result = match op {
198            BinaryOp::Add => left + right,
199            BinaryOp::Subtract => left - right,
200            BinaryOp::Multiply => left * right,
201            BinaryOp::Divide => {
202                if right.abs() < f64::EPSILON {
203                    return None;
204                }
205                left / right
206            }
207            BinaryOp::Modulo => left % right,
208            BinaryOp::Power => left.powf(right),
209            BinaryOp::Equal => {
210                if (left - right).abs() < f64::EPSILON {
211                    1.0
212                } else {
213                    0.0
214                }
215            }
216            BinaryOp::NotEqual => {
217                if (left - right).abs() >= f64::EPSILON {
218                    1.0
219                } else {
220                    0.0
221                }
222            }
223            BinaryOp::Less => {
224                if left < right {
225                    1.0
226                } else {
227                    0.0
228                }
229            }
230            BinaryOp::LessEqual => {
231                if left <= right {
232                    1.0
233                } else {
234                    0.0
235                }
236            }
237            BinaryOp::Greater => {
238                if left > right {
239                    1.0
240                } else {
241                    0.0
242                }
243            }
244            BinaryOp::GreaterEqual => {
245                if left >= right {
246                    1.0
247                } else {
248                    0.0
249                }
250            }
251            BinaryOp::And => {
252                if left != 0.0 && right != 0.0 {
253                    1.0
254                } else {
255                    0.0
256                }
257            }
258            BinaryOp::Or => {
259                if left != 0.0 || right != 0.0 {
260                    1.0
261                } else {
262                    0.0
263                }
264            }
265        };
266
267        Some(result)
268    }
269
270    /// Evaluates a constant unary operation
271    fn eval_const_unary(&self, op: UnaryOp, operand: f64) -> Option<f64> {
272        let result = match op {
273            UnaryOp::Negate => -operand,
274            UnaryOp::Plus => operand,
275            UnaryOp::Not => {
276                if operand.abs() < f64::EPSILON {
277                    1.0
278                } else {
279                    0.0
280                }
281            }
282        };
283
284        Some(result)
285    }
286
287    /// Performs algebraic simplifications
288    fn algebraic_simplify(&self, expr: Expr) -> Expr {
289        match expr {
290            Expr::Binary {
291                left,
292                op,
293                right,
294                ty,
295            } => {
296                let left_opt = self.algebraic_simplify(*left);
297                let right_opt = self.algebraic_simplify(*right);
298
299                // x + 0 = x
300                if op == BinaryOp::Add {
301                    if let Expr::Number(n) = &right_opt {
302                        if n.abs() < f64::EPSILON {
303                            return left_opt;
304                        }
305                    }
306                    if let Expr::Number(n) = &left_opt {
307                        if n.abs() < f64::EPSILON {
308                            return right_opt;
309                        }
310                    }
311                }
312
313                // x - 0 = x
314                if op == BinaryOp::Subtract {
315                    if let Expr::Number(n) = &right_opt {
316                        if n.abs() < f64::EPSILON {
317                            return left_opt;
318                        }
319                    }
320                }
321
322                // NOTE: `x * 0 = 0` and `0 * x = 0` are intentionally NOT simplified here.
323                // Per IEEE-754, `NaN * 0.0 == NaN` and `Inf * 0.0 == NaN`, not `0.0`. This
324                // crate's raster-algebra evaluator uses NaN as a NoData sentinel (e.g. from
325                // division by near-zero), so folding `x * 0` to a constant `0.0` would
326                // silently turn per-pixel NoData into a false constant zero. See the
327                // `test_algebraic_simplify_mul_zero_preserves_nan_semantics` regression test.
328
329                // x * 1 = x
330                if op == BinaryOp::Multiply {
331                    if let Expr::Number(n) = &right_opt {
332                        if (n - 1.0).abs() < f64::EPSILON {
333                            return left_opt;
334                        }
335                    }
336                    if let Expr::Number(n) = &left_opt {
337                        if (n - 1.0).abs() < f64::EPSILON {
338                            return right_opt;
339                        }
340                    }
341                }
342
343                // x / 1 = x
344                if op == BinaryOp::Divide {
345                    if let Expr::Number(n) = &right_opt {
346                        if (n - 1.0).abs() < f64::EPSILON {
347                            return left_opt;
348                        }
349                    }
350                }
351
352                // x ^ 0 = 1
353                if op == BinaryOp::Power {
354                    if let Expr::Number(n) = &right_opt {
355                        if n.abs() < f64::EPSILON {
356                            return Expr::Number(1.0);
357                        }
358                    }
359                }
360
361                // x ^ 1 = x
362                if op == BinaryOp::Power {
363                    if let Expr::Number(n) = &right_opt {
364                        if (n - 1.0).abs() < f64::EPSILON {
365                            return left_opt;
366                        }
367                    }
368                }
369
370                Expr::Binary {
371                    left: Box::new(left_opt),
372                    op,
373                    right: Box::new(right_opt),
374                    ty,
375                }
376            }
377            Expr::Unary {
378                op,
379                expr: inner,
380                ty,
381            } => {
382                let inner_opt = self.algebraic_simplify(*inner);
383
384                // --x = x
385                if op == UnaryOp::Negate {
386                    if let Expr::Unary {
387                        op: UnaryOp::Negate,
388                        expr: double_neg,
389                        ..
390                    } = &inner_opt
391                    {
392                        return *double_neg.clone();
393                    }
394                }
395
396                // +x = x
397                if op == UnaryOp::Plus {
398                    return inner_opt;
399                }
400
401                Expr::Unary {
402                    op,
403                    expr: Box::new(inner_opt),
404                    ty,
405                }
406            }
407            Expr::Conditional {
408                condition,
409                then_expr,
410                else_expr,
411                ty,
412            } => Expr::Conditional {
413                condition: Box::new(self.algebraic_simplify(*condition)),
414                then_expr: Box::new(self.algebraic_simplify(*then_expr)),
415                else_expr: Box::new(self.algebraic_simplify(*else_expr)),
416                ty,
417            },
418            Expr::Call { name, args, ty } => Expr::Call {
419                name,
420                args: args
421                    .into_iter()
422                    .map(|arg| self.algebraic_simplify(arg))
423                    .collect(),
424                ty,
425            },
426            Expr::Block {
427                statements,
428                result,
429                ty,
430            } => Expr::Block {
431                statements: statements
432                    .into_iter()
433                    .map(|stmt| self.optimize_statement(stmt))
434                    .collect(),
435                result: result.map(|r| Box::new(self.algebraic_simplify(*r))),
436                ty,
437            },
438            _ => expr,
439        }
440    }
441
442    /// Eliminates common subexpressions
443    fn eliminate_common_subexpressions(&self, expr: Expr) -> Expr {
444        let mut seen: HashMap<String, usize> = HashMap::new();
445        self.cse_pass(&expr, &mut seen);
446        // Note: Full CSE implementation would require more complex analysis
447        // This is a simplified version that just counts occurrences
448        expr
449    }
450
451    fn cse_pass(&self, expr: &Expr, seen: &mut HashMap<String, usize>) {
452        match expr {
453            Expr::Binary { left, right, .. } => {
454                self.cse_pass(left, seen);
455                self.cse_pass(right, seen);
456                let key = format!("{:?}", expr);
457                *seen.entry(key).or_insert(0) += 1;
458            }
459            Expr::Unary { expr: inner, .. } => {
460                self.cse_pass(inner, seen);
461            }
462            Expr::Call { args, .. } => {
463                for arg in args {
464                    self.cse_pass(arg, seen);
465                }
466            }
467            Expr::Conditional {
468                condition,
469                then_expr,
470                else_expr,
471                ..
472            } => {
473                self.cse_pass(condition, seen);
474                self.cse_pass(then_expr, seen);
475                self.cse_pass(else_expr, seen);
476            }
477            _ => {}
478        }
479    }
480}
481
482#[cfg(test)]
483mod tests {
484    use super::*;
485    use crate::dsl::Type;
486
487    #[test]
488    fn test_constant_fold_add() {
489        let expr = Expr::Binary {
490            left: Box::new(Expr::Number(2.0)),
491            op: BinaryOp::Add,
492            right: Box::new(Expr::Number(3.0)),
493            ty: Type::Number,
494        };
495
496        let opt = Optimizer::new(OptLevel::Basic);
497        let result = opt.optimize_expr(expr);
498
499        assert!(matches!(result, Expr::Number(n) if (n - 5.0).abs() < 1e-10));
500    }
501
502    #[test]
503    fn test_constant_fold_nested() {
504        let expr = Expr::Binary {
505            left: Box::new(Expr::Binary {
506                left: Box::new(Expr::Number(2.0)),
507                op: BinaryOp::Multiply,
508                right: Box::new(Expr::Number(3.0)),
509                ty: Type::Number,
510            }),
511            op: BinaryOp::Add,
512            right: Box::new(Expr::Number(4.0)),
513            ty: Type::Number,
514        };
515
516        let opt = Optimizer::new(OptLevel::Basic);
517        let result = opt.optimize_expr(expr);
518
519        assert!(matches!(result, Expr::Number(n) if (n - 10.0).abs() < 1e-10));
520    }
521
522    #[test]
523    fn test_algebraic_simplify_add_zero() {
524        let expr = Expr::Binary {
525            left: Box::new(Expr::Band(1)),
526            op: BinaryOp::Add,
527            right: Box::new(Expr::Number(0.0)),
528            ty: Type::Raster,
529        };
530
531        let opt = Optimizer::new(OptLevel::Standard);
532        let result = opt.optimize_expr(expr);
533
534        assert!(matches!(result, Expr::Band(1)));
535    }
536
537    #[test]
538    fn test_algebraic_simplify_mul_one() {
539        let expr = Expr::Binary {
540            left: Box::new(Expr::Band(1)),
541            op: BinaryOp::Multiply,
542            right: Box::new(Expr::Number(1.0)),
543            ty: Type::Raster,
544        };
545
546        let opt = Optimizer::new(OptLevel::Standard);
547        let result = opt.optimize_expr(expr);
548
549        assert!(matches!(result, Expr::Band(1)));
550    }
551
552    #[test]
553    fn test_algebraic_simplify_mul_zero_not_folded() {
554        // `x * 0` and `0 * x` must NOT be simplified to a constant `Number(0.0)`.
555        // Per IEEE-754, NaN * 0.0 == NaN and Inf * 0.0 == NaN (not 0.0), and this
556        // crate's raster-algebra evaluator relies on NaN as a NoData sentinel.
557        // Folding this to a constant would silently turn per-pixel NoData into 0.0.
558        let expr = Expr::Binary {
559            left: Box::new(Expr::Band(1)),
560            op: BinaryOp::Multiply,
561            right: Box::new(Expr::Number(0.0)),
562            ty: Type::Raster,
563        };
564
565        let opt = Optimizer::new(OptLevel::Standard);
566        let result = opt.optimize_expr(expr);
567
568        // Must remain a Binary(Band(1) * 0), not collapse to Number(0.0).
569        assert!(matches!(
570            result,
571            Expr::Binary {
572                op: BinaryOp::Multiply,
573                ..
574            }
575        ));
576        assert!(!matches!(result, Expr::Number(_)));
577
578        // Same for the commuted form `0 * x`.
579        let expr_commuted = Expr::Binary {
580            left: Box::new(Expr::Number(0.0)),
581            op: BinaryOp::Multiply,
582            right: Box::new(Expr::Band(1)),
583            ty: Type::Raster,
584        };
585        let result_commuted = opt.optimize_expr(expr_commuted);
586        assert!(!matches!(result_commuted, Expr::Number(_)));
587    }
588
589    #[test]
590    fn test_algebraic_simplify_mul_zero_preserves_nan_semantics() {
591        // Regression test: `(B1 / B2) * 0` must NOT silently turn per-pixel NoData
592        // (NaN, produced by division-by-near-zero) into a false constant `0.0`.
593        //
594        // The evaluator treats NaN as the NoData sentinel; the algebraic simplifier
595        // must not fold `x * 0` to `0` since `NaN * 0.0 == NaN` (and `Inf * 0.0 ==
596        // NaN`) under IEEE-754, not `0.0`.
597        use crate::dsl::RasterDsl;
598        use oxigdal_core::buffer::RasterBuffer;
599        use oxigdal_core::types::RasterDataType;
600
601        let mut b1 = RasterBuffer::zeros(1, 1, RasterDataType::Float32);
602        let mut b2 = RasterBuffer::zeros(1, 1, RasterDataType::Float32);
603        // B1 = 1.0, B2 = 0.0 => B1 / B2 evaluates to NaN (NoData) in this evaluator.
604        assert!(b1.set_pixel(0, 0, 1.0).is_ok());
605        assert!(b2.set_pixel(0, 0, 0.0).is_ok());
606
607        let dsl = RasterDsl::new();
608        let result = dsl
609            .execute("(B1 / B2) * 0", &[b1, b2])
610            .expect("DSL expression should execute successfully");
611
612        let pixel = result
613            .get_pixel(0, 0)
614            .expect("pixel (0, 0) should be readable");
615
616        assert!(
617            pixel.is_nan(),
618            "expected NoData (NaN) to survive `* 0`, got {pixel} instead"
619        );
620    }
621
622    #[test]
623    fn test_double_negation() {
624        let expr = Expr::Unary {
625            op: UnaryOp::Negate,
626            expr: Box::new(Expr::Unary {
627                op: UnaryOp::Negate,
628                expr: Box::new(Expr::Band(1)),
629                ty: Type::Raster,
630            }),
631            ty: Type::Raster,
632        };
633
634        let opt = Optimizer::new(OptLevel::Standard);
635        let result = opt.optimize_expr(expr);
636
637        assert!(matches!(result, Expr::Band(1)));
638    }
639
640    #[test]
641    fn test_unary_plus() {
642        let expr = Expr::Unary {
643            op: UnaryOp::Plus,
644            expr: Box::new(Expr::Band(1)),
645            ty: Type::Raster,
646        };
647
648        let opt = Optimizer::new(OptLevel::Standard);
649        let result = opt.optimize_expr(expr);
650
651        assert!(matches!(result, Expr::Band(1)));
652    }
653}
oxigdal_algorithms/dsl/optimizer.rs

oxigdal_algorithms/dsl/
optimizer.rs
