eww/src/eval.rs

211 lines
8.5 KiB
Rust

use itertools::Itertools;
use crate::{
ast::{BinOp, SimplExpr, Span, UnaryOp},
dynval::{ConversionError, DynVal},
};
use std::collections::HashMap;
#[derive(Debug, thiserror::Error)]
pub enum EvalError {
#[error("Invalid regex: {0}")]
InvalidRegex(#[from] regex::Error),
#[error("got unresolved variable `{0}`")]
UnresolvedVariable(VarName),
#[error("Type error: {0}")]
ConversionError(#[from] ConversionError),
#[error("Incorrect number of arguments given to function: {0}")]
WrongArgCount(String),
#[error("Unknown function {0}")]
UnknownFunction(String),
#[error("Unknown variable {0}")]
UnknownVariable(String),
#[error("Unable to index into value {0}")]
CannotIndex(String),
#[error("{1}")]
Spanned(Span, Box<EvalError>),
}
impl EvalError {
pub fn span(&self) -> Option<Span> {
match self {
EvalError::Spanned(span, _) => Some(*span),
EvalError::ConversionError(err) => err.span(),
_ => None,
}
}
pub fn at(self, span: Span) -> Self {
Self::Spanned(span, Box::new(self))
}
}
type VarName = String;
pub trait FunctionSource {
type Err;
fn run_fn(&self, name: &str, args: &Vec<DynVal>) -> Result<DynVal, Self::Err>;
}
impl SimplExpr {
pub fn map_terminals_into(self, f: impl Fn(Self) -> Self) -> Self {
use SimplExpr::*;
match self {
BinOp(span, box a, op, box b) => BinOp(span, box f(a), op, box f(b)),
UnaryOp(span, op, box a) => UnaryOp(span, op, box f(a)),
IfElse(span, box a, box b, box c) => IfElse(span, box f(a), box f(b), box f(c)),
JsonAccess(span, box a, box b) => JsonAccess(span, box f(a), box f(b)),
FunctionCall(span, name, args) => FunctionCall(span, name, args.into_iter().map(f).collect()),
other => f(other),
}
}
/// resolve variable references in the expression. Fails if a variable cannot be resolved.
pub fn resolve_refs(self, variables: &HashMap<VarName, DynVal>) -> Result<Self, EvalError> {
use SimplExpr::*;
match self {
// Literal(x) => Ok(Literal(AttrValue::from_primitive(x.resolve_fully(&variables)?))),
Literal(span, x) => Ok(Literal(span, x)),
BinOp(span, box a, op, box b) => Ok(BinOp(span, box a.resolve_refs(variables)?, op, box b.resolve_refs(variables)?)),
UnaryOp(span, op, box x) => Ok(UnaryOp(span, op, box x.resolve_refs(variables)?)),
IfElse(span, box a, box b, box c) => {
Ok(IfElse(span, box a.resolve_refs(variables)?, box b.resolve_refs(variables)?, box c.resolve_refs(variables)?))
}
JsonAccess(span, box a, box b) => {
Ok(JsonAccess(span, box a.resolve_refs(variables)?, box b.resolve_refs(variables)?))
}
FunctionCall(span, function_name, args) => Ok(FunctionCall(
span,
function_name,
args.into_iter().map(|a| a.resolve_refs(variables)).collect::<Result<_, EvalError>>()?,
)),
VarRef(span, ref name) => match variables.get(name) {
Some(value) => Ok(Literal(span, value.clone())),
None => Err(EvalError::UnknownVariable(name.to_string()).at(span)),
},
}
}
pub fn var_refs(&self) -> Vec<&String> {
use SimplExpr::*;
match self {
Literal(..) => Vec::new(),
VarRef(_, name) => vec![name],
BinOp(_, box a, _, box b) | JsonAccess(_, box a, box b) => {
let mut refs = a.var_refs();
refs.append(&mut b.var_refs());
refs
}
UnaryOp(_, _, box x) => x.var_refs(),
IfElse(_, box a, box b, box c) => {
let mut refs = a.var_refs();
refs.append(&mut b.var_refs());
refs.append(&mut c.var_refs());
refs
}
FunctionCall(_, _, args) => args.iter().flat_map(|a| a.var_refs()).collect_vec(),
}
}
pub fn eval(self, values: &HashMap<VarName, DynVal>) -> Result<DynVal, EvalError> {
let span = self.span();
let value = match self {
SimplExpr::Literal(_, x) => Ok(x),
SimplExpr::VarRef(span, ref name) => {
Ok(values.get(name).cloned().ok_or_else(|| EvalError::UnresolvedVariable(name.to_string()).at(span))?.at(span))
}
SimplExpr::BinOp(_, a, op, b) => {
let a = a.eval(values)?;
let b = b.eval(values)?;
Ok(match op {
BinOp::Equals => DynVal::from(a == b),
BinOp::NotEquals => DynVal::from(a != b),
BinOp::And => DynVal::from(a.as_bool()? && b.as_bool()?),
BinOp::Or => DynVal::from(a.as_bool()? || b.as_bool()?),
BinOp::Plus => match a.as_f64() {
Ok(num) => DynVal::from(num + b.as_f64()?),
Err(_) => DynVal::from(format!("{}{}", a.as_string()?, b.as_string()?)),
},
BinOp::Minus => DynVal::from(a.as_f64()? - b.as_f64()?),
BinOp::Times => DynVal::from(a.as_f64()? * b.as_f64()?),
BinOp::Div => DynVal::from(a.as_f64()? / b.as_f64()?),
BinOp::Mod => DynVal::from(a.as_f64()? % b.as_f64()?),
BinOp::GT => DynVal::from(a.as_f64()? > b.as_f64()?),
BinOp::LT => DynVal::from(a.as_f64()? < b.as_f64()?),
BinOp::Elvis => DynVal::from(if a.0.is_empty() { b } else { a }),
BinOp::RegexMatch => {
let regex = regex::Regex::new(&b.as_string()?)?;
DynVal::from(regex.is_match(&a.as_string()?))
}
})
}
SimplExpr::UnaryOp(_, op, a) => {
let a = a.eval(values)?;
Ok(match op {
UnaryOp::Not => DynVal::from(!a.as_bool()?),
})
}
SimplExpr::IfElse(_, cond, yes, no) => {
if cond.eval(values)?.as_bool()? {
yes.eval(values)
} else {
no.eval(values)
}
}
SimplExpr::JsonAccess(span, val, index) => {
let val = val.eval(values)?;
let index = index.eval(values)?;
match val.as_json_value()? {
serde_json::Value::Array(val) => {
let index = index.as_i32()?;
let indexed_value = val.get(index as usize).unwrap_or(&serde_json::Value::Null);
Ok(DynVal::from(indexed_value))
}
serde_json::Value::Object(val) => {
let indexed_value = val
.get(&index.as_string()?)
.or_else(|| val.get(&index.as_i32().ok()?.to_string()))
.unwrap_or(&serde_json::Value::Null);
Ok(DynVal::from(indexed_value))
}
_ => Err(EvalError::CannotIndex(format!("{}", val)).at(span)),
}
}
SimplExpr::FunctionCall(span, function_name, args) => {
let args = args.into_iter().map(|a| a.eval(values)).collect::<Result<_, EvalError>>()?;
call_expr_function(&function_name, args).map_err(|e| e.at(span))
}
};
Ok(value?.at(span))
}
}
fn call_expr_function(name: &str, args: Vec<DynVal>) -> Result<DynVal, EvalError> {
match name {
"round" => match args.as_slice() {
[num, digits] => {
let num = num.as_f64()?;
let digits = digits.as_i32()?;
Ok(DynVal::from(format!("{:.1$}", num, digits as usize)))
}
_ => Err(EvalError::WrongArgCount(name.to_string())),
},
"replace" => match args.as_slice() {
[string, pattern, replacement] => {
let string = string.as_string()?;
let pattern = regex::Regex::new(&pattern.as_string()?)?;
let replacement = replacement.as_string()?;
Ok(DynVal::from(pattern.replace_all(&string, replacement.replace("$", "$$").replace("\\", "$")).into_owned()))
}
_ => Err(EvalError::WrongArgCount(name.to_string())),
},
_ => Err(EvalError::UnknownFunction(name.to_string())),
}
}