DarkRoom/core/dr-pipeline/src/declared/expr.rs

//! TRACES: FR-PLG-2
//! The little expression language a node declaration derives its uniforms in.
//!
//! Uniforms are functions of parameters — `exp2(exposure)`, `blacks / 100 *
//! 0.02` — and that derivation is the one piece of a node that is genuinely
//! computation rather than description. The language is arithmetic over the
//! node's own parameters plus a fixed set of maths functions: enough for every
//! operation in the chain, and small enough that a reader of the YAML can see
//! exactly what will happen.
//!
//! # Why this file is compiled twice
//!
//! It is `#[path]`-included by `build.rs` as well as being a module of the
//! crate, and it deliberately depends on nothing but `std` so that it can be.
//!
//! There are two backends over one grammar. `build.rs` renders an [`Expr`] to
//! Rust source, so a built-in node's arithmetic costs nothing at run time and
//! an unknown name is a build error naming the file. [`Expr::eval`] evaluates
//! the same tree directly, which is what lets a declaration loaded at run time
//! produce uniforms without a compiler.
//!
//! **The two must agree bit for bit**, or a plugin is not the same kind of
//! thing as a built-in and `tests/declared_parity.rs` says so. Sharing the
//! tokeniser and the parser removes the larger half of the ways they could
//! drift; the remaining half is the pair of backends, and each entry in
//! [`FUNCTIONS`] below is written twice on purpose, once here and once in
//! `build.rs`, with the parity test standing between them.

use std::collections::BTreeSet;

/// A number in a declaration, as the `f32` the arithmetic will actually use.
///
/// **The route through the decimal string is load-bearing, not clumsiness.**
/// `build.rs` renders a number as a Rust literal — `0.02f32` — and `rustc`
/// rounds that decimal text to the nearest `f32` exactly once. Writing
/// `n as f32` here would instead round the `f64` YAML parsed to the nearest
/// `f32`, which is a *second* rounding on top of the one that produced the
/// `f64`, and double rounding does not always land where single rounding does.
///
/// So this reproduces what the compiler sees: `{:?}` is the shortest decimal
/// that round-trips the `f64`, which is exactly the literal `build.rs` emits,
/// and parsing it as `f32` is exactly what `rustc` does with it.
pub fn as_f32(n: f64) -> f32 {
    // Infallible in practice: `{:?}` on a finite `f64` is always a parseable
    // decimal, and the infinities and NaN it can also print all parse back.
    // The fallback is the direct cast rather than a panic, because a bad
    // number in a declaration is the reader's error to report, not this
    // function's to crash on.
    format!("{n:?}").parse::<f32>().unwrap_or(n as f32)
}

/// The maths functions a declaration may call, and how many arguments each
/// takes.
///
/// A closed list rather than a passthrough to `f32`: a node is a description,
/// and letting it name arbitrary Rust would make the YAML a second, worse
/// place to write code. Closed for the same reason `WidgetKind` and
/// `attributes` are closed (FR-PLG-2d) — a typo must be an error rather than a
/// silent new category of one.
///
/// Shared by both backends so that the *set* of callable functions cannot
/// drift even though the two translations of each must be written separately.
pub const FUNCTIONS: &[(&str, usize)] = &[
    ("exp2", 1),
    ("log2", 1),
    ("exp", 1),
    ("sqrt", 1),
    ("abs", 1),
    ("floor", 1),
    ("ceil", 1),
    ("round", 1),
    ("pow", 2),
    ("min", 2),
    ("max", 2),
    ("clamp", 3),
    ("mix", 3),
];

/// A parsed expression over a node's parameters.
#[derive(Debug, Clone, PartialEq)]
pub enum Expr {
    Num(f64),
    Param(String),
    Neg(Box<Expr>),
    Bin(char, Box<Expr>, Box<Expr>),
    Call(String, Vec<Expr>),
}

/// Parse and validate an expression against the parameters a node declares.
///
/// Validation happens here rather than in either backend, so that "this names
/// a parameter that does not exist" is one error message in one place and
/// cannot be reported at build time but missed at load time.
pub fn parse(src: &str, params: &BTreeSet<&str>) -> Result<Expr, String> {
    let tokens = tokenise(src)?;
    let mut parser = Parser { tokens, at: 0 };
    let expr = parser.expr()?;
    if parser.at < parser.tokens.len() {
        return Err(format!(
            "unexpected `{}` after the end of the expression",
            parser.tokens[parser.at]
        ));
    }
    check(&expr, params)?;
    Ok(expr)
}

/// Every name and arity in the tree resolves.
fn check(expr: &Expr, params: &BTreeSet<&str>) -> Result<(), String> {
    match expr {
        Expr::Num(_) => Ok(()),
        Expr::Param(name) => {
            if params.contains(name.as_str()) {
                return Ok(());
            }
            let known: Vec<&str> = params.iter().copied().collect();
            Err(format!(
                "`{name}` is not a parameter of this node. Its parameters \
                 are: {}",
                known.join(", ")
            ))
        }
        Expr::Neg(inner) => check(inner, params),
        Expr::Bin(_, l, r) => {
            check(l, params)?;
            check(r, params)
        }
        Expr::Call(name, args) => {
            let Some((_, arity)) = FUNCTIONS.iter().find(|(f, _)| *f == name) else {
                let known: Vec<&str> = FUNCTIONS.iter().map(|(f, _)| *f).collect();
                return Err(format!(
                    "`{name}` is not one of the maths functions a node may \
                     call. Available: {}",
                    known.join(", ")
                ));
            };
            if args.len() != *arity {
                return Err(format!(
                    "`{name}` takes {arity} argument(s), given {}",
                    args.len()
                ));
            }
            for a in args {
                check(a, params)?;
            }
            Ok(())
        }
    }
}

impl Expr {
    /// TRACES: FR-PLG-2
    /// Evaluate this expression for a set of parameter values.
    ///
    /// **Every step is an `f32` operation in the same order `build.rs` renders
    /// it**, which is what makes the interpreted result bit-identical to the
    /// compiled one rather than merely close. Rust's `f32` arithmetic is IEEE
    /// 754 with no excess precision, so `(a * b) + c` here and `(a * b) + c`
    /// in generated source are the same number down to the last bit — and the
    /// parity test asserts exactly that rather than an epsilon, because a
    /// tolerance is how a real divergence gets to hide.
    ///
    /// `param` is asked for a parameter's current value. It is a closure
    /// rather than a map so the caller can serve the values out of whatever it
    /// already has, which for [`super::DeclaredOp`] is a plain `Vec<f32>`
    /// indexed in declaration order.
    ///
    /// Infallible: [`parse`] has already established that every name resolves
    /// and every call has the right arity. An unknown parameter reaching here
    /// would be a reader that let one through, so `param` decides what to do
    /// about it rather than this returning a `Result` every caller would
    /// unwrap.
    pub fn eval(&self, param: &dyn Fn(&str) -> f32) -> f32 {
        match self {
            Expr::Num(n) => as_f32(*n),
            Expr::Param(name) => param(name),
            Expr::Neg(inner) => -inner.eval(param),
            Expr::Bin(op, l, r) => {
                let (l, r) = (l.eval(param), r.eval(param));
                match op {
                    '+' => l + r,
                    '-' => l - r,
                    '*' => l * r,
                    '/' => l / r,
                    // `tokenise` only ever produces these four as binary
                    // operators, and `Parser` only ever builds `Bin` from what
                    // `tokenise` produced.
                    _ => unreachable!("`{op}` is not a binary operator"),
                }
            }
            Expr::Call(name, args) => {
                let a = |i: usize| args[i].eval(param);
                match name.as_str() {
                    "exp2" => f32::exp2(a(0)),
                    "log2" => f32::log2(a(0)),
                    "exp" => f32::exp(a(0)),
                    "sqrt" => f32::sqrt(a(0)),
                    "abs" => f32::abs(a(0)),
                    "floor" => f32::floor(a(0)),
                    "ceil" => f32::ceil(a(0)),
                    "round" => f32::round(a(0)),
                    "pow" => f32::powf(a(0), a(1)),
                    "min" => f32::min(a(0), a(1)),
                    "max" => f32::max(a(0), a(1)),
                    "clamp" => f32::clamp(a(0), a(1), a(2)),
                    // Spelled out rather than called, matching what `build.rs`
                    // renders: Rust has no `mix`, and this linear form is what
                    // WGSL's `mix` means. The association matters — `a + (b -
                    // a) * t` and `a * (1 - t) + b * t` are the same value in
                    // real arithmetic and different ones in `f32`.
                    "mix" => {
                        let (x, y, t) = (a(0), a(1), a(2));
                        x + (y - x) * t
                    }
                    // `parse` rejects anything not in `FUNCTIONS`.
                    _ => unreachable!("`{name}` is not a declared maths function"),
                }
            }
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum Tok {
    Num(f64),
    Ident(String),
    Sym(char),
}

impl std::fmt::Display for Tok {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Tok::Num(n) => write!(f, "{n}"),
            Tok::Ident(s) => write!(f, "{s}"),
            Tok::Sym(c) => write!(f, "{c}"),
        }
    }
}

fn tokenise(src: &str) -> Result<Vec<Tok>, String> {
    let bytes: Vec<char> = src.chars().collect();
    let mut out = Vec::new();
    let mut i = 0;
    while i < bytes.len() {
        let c = bytes[i];
        if c.is_whitespace() {
            i += 1;
        } else if c.is_ascii_digit()
            || (c == '.' && bytes.get(i + 1).is_some_and(char::is_ascii_digit))
        {
            let start = i;
            while i < bytes.len() && (bytes[i].is_ascii_digit() || bytes[i] == '.') {
                i += 1;
            }
            let text: String = bytes[start..i].iter().collect();
            let n = text
                .parse::<f64>()
                .map_err(|_| format!("`{text}` is not a number"))?;
            out.push(Tok::Num(n));
        } else if c.is_ascii_alphabetic() || c == '_' {
            let start = i;
            while i < bytes.len() && (bytes[i].is_ascii_alphanumeric() || bytes[i] == '_') {
                i += 1;
            }
            out.push(Tok::Ident(bytes[start..i].iter().collect()));
        } else if "+-*/(),".contains(c) {
            out.push(Tok::Sym(c));
            i += 1;
        } else {
            return Err(format!(
                "`{c}` is not valid in an expression; the language is \
                 arithmetic (+ - * /), parentheses, numbers, this node's \
                 parameters, and the maths functions"
            ));
        }
    }
    if out.is_empty() {
        return Err("the expression is empty".into());
    }
    Ok(out)
}

struct Parser {
    tokens: Vec<Tok>,
    at: usize,
}

impl Parser {
    fn peek(&self) -> Option<&Tok> {
        self.tokens.get(self.at)
    }

    fn eat(&mut self, sym: char) -> bool {
        if self.peek() == Some(&Tok::Sym(sym)) {
            self.at += 1;
            return true;
        }
        false
    }

    fn expr(&mut self) -> Result<Expr, String> {
        let mut left = self.term()?;
        loop {
            if self.eat('+') {
                left = Expr::Bin('+', Box::new(left), Box::new(self.term()?));
            } else if self.eat('-') {
                left = Expr::Bin('-', Box::new(left), Box::new(self.term()?));
            } else {
                return Ok(left);
            }
        }
    }

    fn term(&mut self) -> Result<Expr, String> {
        let mut left = self.unary()?;
        loop {
            if self.eat('*') {
                left = Expr::Bin('*', Box::new(left), Box::new(self.unary()?));
            } else if self.eat('/') {
                left = Expr::Bin('/', Box::new(left), Box::new(self.unary()?));
            } else {
                return Ok(left);
            }
        }
    }

    fn unary(&mut self) -> Result<Expr, String> {
        if self.eat('-') {
            return Ok(Expr::Neg(Box::new(self.unary()?)));
        }
        self.primary()
    }

    fn primary(&mut self) -> Result<Expr, String> {
        match self.peek().cloned() {
            Some(Tok::Num(n)) => {
                self.at += 1;
                Ok(Expr::Num(n))
            }
            Some(Tok::Ident(name)) => {
                self.at += 1;
                if !self.eat('(') {
                    return Ok(Expr::Param(name));
                }
                let mut args = Vec::new();
                if !self.eat(')') {
                    loop {
                        args.push(self.expr()?);
                        if self.eat(',') {
                            continue;
                        }
                        if self.eat(')') {
                            break;
                        }
                        return Err(format!("expected `,` or `)` in the call to `{name}`"));
                    }
                }
                Ok(Expr::Call(name, args))
            }
            Some(Tok::Sym('(')) => {
                self.at += 1;
                let inner = self.expr()?;
                if !self.eat(')') {
                    return Err("unclosed `(`".into());
                }
                Ok(inner)
            }
            Some(t) => Err(format!("unexpected `{t}`")),
            None => Err("the expression ends early".into()),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn params() -> BTreeSet<&'static str> {
        ["a", "b"].into_iter().collect()
    }

    fn eval(src: &str, a: f32, b: f32) -> f32 {
        parse(src, &params())
            .expect("parses")
            .eval(&|name| match name {
                "a" => a,
                "b" => b,
                other => panic!("no parameter {other}"),
            })
    }

    #[test]
    fn arithmetic_follows_the_usual_precedence() {
        assert_eq!(eval("a + b * 2", 1.0, 3.0), 7.0);
        assert_eq!(eval("(a + b) * 2", 1.0, 3.0), 8.0);
    }

    #[test]
    fn unary_minus_binds_tighter_than_addition() {
        // `-(a + b)` and `-a + b` are different numbers, and the parser has to
        // agree with the renderer about which one `-a + b` is.
        assert_eq!(eval("-a + b", 1.0, 3.0), 2.0);
        assert_eq!(eval("-(a + b)", 1.0, 3.0), -4.0);
    }

    #[test]
    fn a_number_is_rounded_once_the_way_the_compiler_rounds_a_literal() {
        // The whole reason `as_f32` goes through the decimal string. If this
        // ever regresses to `n as f32`, the values a declared node produces
        // drift from the generated one in the last bit, and every parity
        // assertion has to become a tolerance to keep passing — which is
        // exactly the silent disagreement the test exists to catch.
        assert_eq!(as_f32(0.02), 0.02f32);
        assert_eq!(as_f32(0.1), 0.1f32);
        // A third: eight significant digits, which is past what an `f32`
        // resolves, so this is the case where a second rounding could land
        // somewhere the compiler's single one does not.
        assert_eq!(as_f32(1.0 / 3.0), 0.333_333_34_f32);
    }

    #[test]
    fn mix_is_the_linear_form_wgsl_means() {
        assert_eq!(eval("mix(a, b, 0.25)", 0.0, 4.0), 1.0);
    }

    #[test]
    fn an_unknown_parameter_names_the_ones_that_exist() {
        // The error is the whole value of validating in the parser: whoever
        // wrote the typo needs the list, and needs it identically whether the
        // declaration was read by the build script or at load time.
        let err = parse("c * 2", &params()).unwrap_err();
        assert!(err.contains("`c` is not a parameter"), "{err}");
        assert!(err.contains("a, b"), "{err}");
    }

    #[test]
    fn an_unknown_function_is_rejected_rather_than_passed_through() {
        let err = parse("tan(a)", &params()).unwrap_err();
        assert!(err.contains("not one of the maths functions"), "{err}");
    }

    #[test]
    fn a_wrong_arity_is_caught_where_it_is_written() {
        let err = parse("pow(a)", &params()).unwrap_err();
        assert!(err.contains("takes 2 argument(s), given 1"), "{err}");
    }
}