Skip to main content

runmat_vm/compiler/
end_expr.rs

1//! Lowering support for `end` expressions used in indexing.
2
3use crate::instr::EndExpr;
4use runmat_hir::{HirExpr, HirExprKind};
5
6pub(crate) fn expr_is_one(expr: &HirExpr) -> bool {
7    parse_number(expr)
8        .map(|v| (v - 1.0).abs() < 1e-9)
9        .unwrap_or(false)
10}
11
12fn parse_number(expr: &HirExpr) -> Option<f64> {
13    if let HirExprKind::Number(raw) = &expr.kind {
14        raw.parse::<f64>().ok()
15    } else {
16        None
17    }
18}
19
20pub(crate) fn end_numeric_expr(expr: &HirExpr) -> Option<EndExpr> {
21    end_numeric_expr_from_expr(&expr.kind)
22}
23
24pub(crate) fn numeric_expr_any(expr: &HirExpr) -> Option<EndExpr> {
25    end_numeric_expr_internal(&expr.kind).map(|(e, _)| e)
26}
27
28pub(crate) fn range_dynamic_end_spec(
29    expr: &HirExpr,
30) -> Option<(Option<EndExpr>, Option<EndExpr>, EndExpr)> {
31    let HirExprKind::Range(start, step, end) = &expr.kind else {
32        return None;
33    };
34    let start_end = end_numeric_expr(start);
35    let step_end = step.as_ref().and_then(|s| end_numeric_expr(s));
36    let end_end = end_numeric_expr(end);
37    if start_end.is_none() && step_end.is_none() && end_end.is_none() {
38        return None;
39    }
40    let resolved_end = if let Some(e) = end_end {
41        e
42    } else {
43        numeric_expr_any(end)?
44    };
45    Some((start_end, step_end, resolved_end))
46}
47
48fn end_numeric_expr_from_expr(kind: &HirExprKind) -> Option<EndExpr> {
49    let (expr, has_end) = end_numeric_expr_internal(kind)?;
50    if has_end {
51        Some(expr)
52    } else {
53        None
54    }
55}
56
57fn end_numeric_expr_internal(kind: &HirExprKind) -> Option<(EndExpr, bool)> {
58    match kind {
59        HirExprKind::End => Some((EndExpr::End, true)),
60        HirExprKind::Number(s) => s.parse::<f64>().ok().map(|v| (EndExpr::Const(v), false)),
61        HirExprKind::Var(id) => Some((EndExpr::Var(id.0), false)),
62        HirExprKind::Unary(op, inner) => {
63            let (child, has_end) = end_numeric_expr_internal(&inner.kind)?;
64            match op {
65                runmat_parser::UnOp::Plus => Some((EndExpr::Pos(Box::new(child)), has_end)),
66                runmat_parser::UnOp::Minus => Some((EndExpr::Neg(Box::new(child)), has_end)),
67                _ => None,
68            }
69        }
70        HirExprKind::FuncCall(name, args) => {
71            let mut out_args = Vec::with_capacity(args.len());
72            let mut has_end = false;
73            for arg in args {
74                let (e, h) = end_numeric_expr_internal(&arg.kind)?;
75                out_args.push(e);
76                has_end |= h;
77            }
78            let lname = name.to_ascii_lowercase();
79            if args.len() == 1 {
80                let single = out_args.into_iter().next().unwrap_or(EndExpr::Const(0.0));
81                let out = match lname.as_str() {
82                    "floor" => EndExpr::Floor(Box::new(single)),
83                    "ceil" | "ceiling" => EndExpr::Ceil(Box::new(single)),
84                    "round" => EndExpr::Round(Box::new(single)),
85                    "fix" => EndExpr::Fix(Box::new(single)),
86                    _ => EndExpr::Call(name.clone(), vec![single]),
87                };
88                Some((out, has_end))
89            } else {
90                Some((EndExpr::Call(name.clone(), out_args), has_end))
91            }
92        }
93        HirExprKind::Binary(left, op, right) => {
94            let (lhs, left_has_end) = end_numeric_expr_internal(&left.kind)?;
95            let (rhs, right_has_end) = end_numeric_expr_internal(&right.kind)?;
96            let has_end = left_has_end || right_has_end;
97            match op {
98                runmat_parser::BinOp::Add => {
99                    Some((EndExpr::Add(Box::new(lhs), Box::new(rhs)), has_end))
100                }
101                runmat_parser::BinOp::Sub => {
102                    Some((EndExpr::Sub(Box::new(lhs), Box::new(rhs)), has_end))
103                }
104                runmat_parser::BinOp::Mul | runmat_parser::BinOp::ElemMul => {
105                    Some((EndExpr::Mul(Box::new(lhs), Box::new(rhs)), has_end))
106                }
107                runmat_parser::BinOp::RightDiv | runmat_parser::BinOp::ElemDiv => {
108                    Some((EndExpr::Div(Box::new(lhs), Box::new(rhs)), has_end))
109                }
110                runmat_parser::BinOp::LeftDiv | runmat_parser::BinOp::ElemLeftDiv => {
111                    Some((EndExpr::LeftDiv(Box::new(lhs), Box::new(rhs)), has_end))
112                }
113                runmat_parser::BinOp::Pow | runmat_parser::BinOp::ElemPow => {
114                    Some((EndExpr::Pow(Box::new(lhs), Box::new(rhs)), has_end))
115                }
116                _ => None,
117            }
118        }
119        _ => None,
120    }
121}