PicoBot/src/tools/expression.rs

/// Evaluate a self-contained mathematical expression without executing code or
/// resolving external variables.
pub(super) fn evaluate(input: &str) -> Result<f64, String> {
    const MAX_EXPRESSION_BYTES: usize = 4096;
    if input.len() > MAX_EXPRESSION_BYTES {
        return Err(format!(
            "expression exceeds the {MAX_EXPRESSION_BYTES}-byte limit"
        ));
    }
    let mut parser = Parser {
        input,
        position: 0,
        depth: 0,
    };
    let value = parser.parse_expression()?;
    parser.skip_whitespace();
    if parser.position != input.len() {
        return Err(parser.error("unexpected trailing input"));
    }
    Ok(value)
}

struct Parser<'a> {
    input: &'a str,
    position: usize,
    depth: usize,
}

impl Parser<'_> {
    fn parse_expression(&mut self) -> Result<f64, String> {
        let mut value = self.parse_term()?;
        loop {
            if self.consume(b'+') {
                value += self.parse_term()?;
            } else if self.consume(b'-') {
                value -= self.parse_term()?;
            } else {
                return Ok(value);
            }
        }
    }

    fn parse_term(&mut self) -> Result<f64, String> {
        let mut value = self.parse_unary()?;
        loop {
            if self.consume(b'*') {
                value *= self.parse_unary()?;
            } else if self.consume(b'/') {
                value /= self.parse_unary()?;
            } else if self.consume(b'%') {
                value %= self.parse_unary()?;
            } else {
                return Ok(value);
            }
        }
    }

    fn parse_unary(&mut self) -> Result<f64, String> {
        if self.consume(b'+') {
            self.nested(Self::parse_unary)
        } else if self.consume(b'-') {
            Ok(-self.nested(Self::parse_unary)?)
        } else {
            self.parse_power()
        }
    }

    fn parse_power(&mut self) -> Result<f64, String> {
        let base = self.parse_primary()?;
        if self.consume(b'^') {
            Ok(base.powf(self.nested(Self::parse_unary)?))
        } else {
            Ok(base)
        }
    }

    fn parse_primary(&mut self) -> Result<f64, String> {
        self.skip_whitespace();
        match self.peek() {
            Some(b'(') => {
                self.position += 1;
                let value = self.nested(Self::parse_expression)?;
                if !self.consume(b')') {
                    return Err(self.error("expected ')'"));
                }
                Ok(value)
            }
            Some(byte) if byte.is_ascii_digit() || byte == b'.' => self.parse_number(),
            Some(byte) if byte.is_ascii_alphabetic() || byte == b'_' => self.parse_identifier(),
            Some(_) => Err(self.error("expected a number, constant, function, or '('")),
            None => Err(self.error("unexpected end of expression")),
        }
    }

    fn parse_number(&mut self) -> Result<f64, String> {
        self.skip_whitespace();
        let start = self.position;
        let mut digits = 0;
        while self.peek().is_some_and(|byte| byte.is_ascii_digit()) {
            self.position += 1;
            digits += 1;
        }
        if self.peek() == Some(b'.') {
            self.position += 1;
            while self.peek().is_some_and(|byte| byte.is_ascii_digit()) {
                self.position += 1;
                digits += 1;
            }
        }
        if digits == 0 {
            return Err(self.error("invalid number"));
        }
        if matches!(self.peek(), Some(b'e' | b'E')) {
            self.position += 1;
            if matches!(self.peek(), Some(b'+' | b'-')) {
                self.position += 1;
            }
            let exponent_start = self.position;
            while self.peek().is_some_and(|byte| byte.is_ascii_digit()) {
                self.position += 1;
            }
            if self.position == exponent_start {
                return Err(self.error("invalid numeric exponent"));
            }
        }

        self.input[start..self.position]
            .parse::<f64>()
            .map_err(|_| self.error("invalid number"))
    }

    fn parse_identifier(&mut self) -> Result<f64, String> {
        self.skip_whitespace();
        let start = self.position;
        while self
            .peek()
            .is_some_and(|byte| byte.is_ascii_alphanumeric() || byte == b'_')
        {
            self.position += 1;
        }
        let name = self.input[start..self.position].to_ascii_lowercase();
        self.skip_whitespace();
        if self.peek() != Some(b'(') {
            return match name.as_str() {
                "pi" => Ok(std::f64::consts::PI),
                "e" => Ok(std::f64::consts::E),
                _ => Err(self.error(&format!("unknown constant or variable '{name}'"))),
            };
        }

        self.position += 1;
        let mut arguments = Vec::new();
        self.skip_whitespace();
        if self.peek() != Some(b')') {
            loop {
                arguments.push(self.nested(Self::parse_expression)?);
                if self.consume(b',') {
                    continue;
                }
                break;
            }
        }
        if !self.consume(b')') {
            return Err(self.error("expected ')' after function arguments"));
        }
        apply_function(&name, &arguments).map_err(|message| self.error(&message))
    }

    fn consume(&mut self, expected: u8) -> bool {
        self.skip_whitespace();
        if self.peek() == Some(expected) {
            self.position += 1;
            true
        } else {
            false
        }
    }

    fn skip_whitespace(&mut self) {
        while self.peek().is_some_and(|byte| byte.is_ascii_whitespace()) {
            self.position += 1;
        }
    }

    fn peek(&self) -> Option<u8> {
        self.input.as_bytes().get(self.position).copied()
    }

    fn nested<T>(&mut self, parse: fn(&mut Self) -> Result<T, String>) -> Result<T, String> {
        const MAX_PARSE_DEPTH: usize = 128;
        if self.depth >= MAX_PARSE_DEPTH {
            return Err(self.error("expression nesting limit exceeded"));
        }
        self.depth += 1;
        let result = parse(self);
        self.depth -= 1;
        result
    }

    fn error(&self, message: &str) -> String {
        format!("{message} at byte {}", self.position)
    }
}

fn apply_function(name: &str, arguments: &[f64]) -> Result<f64, String> {
    let unary = |function: fn(f64) -> f64| match arguments {
        [value] => Ok(function(*value)),
        _ => Err(format!("function '{name}' expects one argument")),
    };
    match name {
        "sqrt" => unary(f64::sqrt),
        "abs" => unary(f64::abs),
        "exp" => unary(f64::exp),
        "ln" => unary(f64::ln),
        "log2" => unary(f64::log2),
        "log10" => unary(f64::log10),
        "sin" => unary(f64::sin),
        "cos" => unary(f64::cos),
        "tan" => unary(f64::tan),
        "asin" => unary(f64::asin),
        "acos" => unary(f64::acos),
        "atan" => unary(f64::atan),
        "sinh" => unary(f64::sinh),
        "cosh" => unary(f64::cosh),
        "tanh" => unary(f64::tanh),
        "asinh" => unary(f64::asinh),
        "acosh" => unary(f64::acosh),
        "atanh" => unary(f64::atanh),
        "floor" => unary(f64::floor),
        "ceil" => unary(f64::ceil),
        "round" => unary(f64::round),
        "signum" => unary(f64::signum),
        "atan2" => match arguments {
            [y, x] => Ok(y.atan2(*x)),
            _ => Err("function 'atan2' expects two arguments".to_string()),
        },
        "min" => arguments
            .iter()
            .copied()
            .reduce(f64::min)
            .ok_or_else(|| "function 'min' expects at least one argument".to_string()),
        "max" => arguments
            .iter()
            .copied()
            .reduce(f64::max)
            .ok_or_else(|| "function 'max' expects at least one argument".to_string()),
        _ => Err(format!("unknown function '{name}'")),
    }
}

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

    #[test]
    fn respects_precedence_and_right_associative_power() {
        assert_eq!(evaluate("15*3+5^(2+1)").unwrap(), 170.0);
        assert_eq!(evaluate("2^3^2").unwrap(), 512.0);
        assert_eq!(evaluate("-2^2").unwrap(), -4.0);
    }

    #[test]
    fn supports_constants_functions_and_scientific_notation() {
        assert_eq!(evaluate("sqrt(1.44e2)").unwrap(), 12.0);
        assert_eq!(evaluate("max(1, 2, 3) + min(4, 5)").unwrap(), 7.0);
        assert!((evaluate("sin(pi / 2)").unwrap() - 1.0).abs() < f64::EPSILON);
    }

    #[test]
    fn rejects_unknown_names_and_trailing_input() {
        assert!(evaluate("unknown").is_err());
        assert!(evaluate("1 + 2 garbage").is_err());
        assert!(evaluate("sqrt() ").is_err());
        assert!(evaluate(&"(".repeat(129)).is_err());
        assert!(evaluate(&"1+".repeat(3000)).is_err());
    }
}