use std::sync::LazyLock;
use nodedb_types::Value;
use sonic_rs;
use crate::error::SqlError;
use crate::functions::registry::{FunctionCategory, FunctionRegistry};
use crate::types::{BinaryOp, SqlExpr, SqlValue, UnaryOp};
pub type FoldResult = Result<Option<SqlValue>, SqlError>;
static DEFAULT_REGISTRY: LazyLock<FunctionRegistry> = LazyLock::new(FunctionRegistry::new);
pub fn default_registry() -> &'static FunctionRegistry {
&DEFAULT_REGISTRY
}
pub fn fold_constant_default(expr: &SqlExpr) -> FoldResult {
fold_constant(expr, default_registry())
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FoldScope {
Reuse,
Once,
}
pub fn fold_constant(expr: &SqlExpr, registry: &FunctionRegistry) -> FoldResult {
fold_constant_scoped(expr, registry, FoldScope::Reuse)
}
pub fn fold_constant_scoped(
expr: &SqlExpr,
registry: &FunctionRegistry,
scope: FoldScope,
) -> FoldResult {
match expr {
SqlExpr::Literal(v) => Ok(Some(v.clone())),
SqlExpr::ArrayLiteral(items) => {
let mut folded = Vec::with_capacity(items.len());
for item in items {
match fold_constant_scoped(item, registry, scope)? {
Some(v) => folded.push(v),
None => return Ok(None),
}
}
Ok(Some(SqlValue::Array(folded)))
}
SqlExpr::UnaryOp {
op: UnaryOp::Neg,
expr,
} => Ok(match fold_constant_scoped(expr, registry, scope)? {
Some(SqlValue::Int(i)) => i.checked_neg().map(SqlValue::Int),
Some(SqlValue::Float(f)) => Some(SqlValue::Float(-f)),
Some(SqlValue::Decimal(d))
if d.scale() == 0 && -d == rust_decimal::Decimal::from(i64::MIN) =>
{
Some(SqlValue::Int(i64::MIN))
}
Some(SqlValue::Decimal(d)) => Some(SqlValue::Decimal(-d)),
_ => None,
}),
SqlExpr::BinaryOp { left, op, right } => {
let (Some(l), Some(r)) = (
fold_constant_scoped(left, registry, scope)?,
fold_constant_scoped(right, registry, scope)?,
) else {
return Ok(None);
};
fold_binary(l, *op, r)
}
SqlExpr::Function { name, args, .. } => {
fold_function_call_scoped(name, args, registry, scope)
}
SqlExpr::Cast { expr, to_type } => Ok(fold_constant_scoped(expr, registry, scope)?
.and_then(|inner| fold_cast(inner, to_type))),
_ => Ok(None),
}
}
fn fold_cast(inner: SqlValue, to_type: &str) -> Option<SqlValue> {
let upper = to_type.to_uppercase();
let base = upper
.split('(')
.next()
.map(str::trim)
.unwrap_or(upper.as_str());
match base {
"REGCLASS" => None,
"NUMERIC" | "DECIMAL" => match inner {
SqlValue::Decimal(d) => Some(SqlValue::Decimal(d)),
SqlValue::Int(i) => Some(SqlValue::Decimal(rust_decimal::Decimal::from(i))),
SqlValue::Float(f) => rust_decimal::Decimal::try_from(f)
.ok()
.map(SqlValue::Decimal),
SqlValue::String(s) => rust_decimal::Decimal::from_str_exact(&s)
.ok()
.map(SqlValue::Decimal),
_ => None,
},
"INTEGER" | "INT" | "BIGINT" | "SMALLINT" | "INT2" | "INT4" | "INT8" => match inner {
SqlValue::Int(i) => Some(SqlValue::Int(i)),
SqlValue::Decimal(d) => {
rust_decimal::prelude::ToPrimitive::to_i64(&d).map(SqlValue::Int)
}
SqlValue::Float(f) => {
if f.is_finite() {
Some(SqlValue::Int(f as i64))
} else {
None
}
}
SqlValue::String(s) => s.parse::<i64>().ok().map(SqlValue::Int),
_ => None,
},
"FLOAT" | "DOUBLE" | "REAL" | "FLOAT4" | "FLOAT8" | "DOUBLE PRECISION" => match inner {
SqlValue::Float(f) => Some(SqlValue::Float(f)),
SqlValue::Int(i) => Some(SqlValue::Float(i as f64)),
SqlValue::Decimal(d) => {
rust_decimal::prelude::ToPrimitive::to_f64(&d).map(SqlValue::Float)
}
SqlValue::String(s) => s.parse::<f64>().ok().map(SqlValue::Float),
_ => None,
},
"TEXT" | "VARCHAR" | "CHAR" | "CHARACTER VARYING" | "CHARACTER" | "BPCHAR" => match inner {
SqlValue::String(s) => Some(SqlValue::String(s)),
SqlValue::Int(i) => Some(SqlValue::String(i.to_string())),
SqlValue::Float(f) => Some(SqlValue::String(f.to_string())),
SqlValue::Decimal(d) => Some(SqlValue::String(d.to_string())),
SqlValue::Bool(b) => Some(SqlValue::String(b.to_string())),
_ => None,
},
"BOOL" | "BOOLEAN" => match inner {
SqlValue::Bool(b) => Some(SqlValue::Bool(b)),
SqlValue::Int(i) => Some(SqlValue::Bool(i != 0)),
SqlValue::String(s) => match s.to_lowercase().as_str() {
"true" | "t" | "yes" | "1" | "on" => Some(SqlValue::Bool(true)),
"false" | "f" | "no" | "0" | "off" => Some(SqlValue::Bool(false)),
_ => None,
},
_ => None,
},
"JSON" | "JSONB" => match inner {
SqlValue::String(s) => Some(SqlValue::String(s)),
SqlValue::Int(i) => Some(SqlValue::String(i.to_string())),
SqlValue::Float(f) => Some(SqlValue::String(f.to_string())),
SqlValue::Decimal(d) => Some(SqlValue::String(d.to_string())),
SqlValue::Bool(b) => Some(SqlValue::String(b.to_string())),
SqlValue::Null => Some(SqlValue::Null),
_ => None,
},
_ => None,
}
}
fn is_zero(v: &SqlValue) -> bool {
match v {
SqlValue::Int(0) => true,
SqlValue::Float(f) => *f == 0.0,
SqlValue::Decimal(d) => d.is_zero(),
_ => false,
}
}
fn overflowed(detail: &str) -> SqlError {
SqlError::ConstantOverflow {
detail: detail.to_owned(),
}
}
fn fold_binary(l: SqlValue, op: BinaryOp, r: SqlValue) -> FoldResult {
if matches!(op, BinaryOp::Div | BinaryOp::Mod) && is_zero(&r) {
return Err(SqlError::DivisionByZero);
}
Ok(Some(match (l, op, r) {
(SqlValue::Int(a), BinaryOp::Add, SqlValue::Int(b)) => {
SqlValue::Int(a.checked_add(b).ok_or_else(|| overflowed("integer add"))?)
}
(SqlValue::Int(a), BinaryOp::Sub, SqlValue::Int(b)) => SqlValue::Int(
a.checked_sub(b)
.ok_or_else(|| overflowed("integer subtract"))?,
),
(SqlValue::Int(a), BinaryOp::Mul, SqlValue::Int(b)) => SqlValue::Int(
a.checked_mul(b)
.ok_or_else(|| overflowed("integer multiply"))?,
),
(SqlValue::Int(a), BinaryOp::Div, SqlValue::Int(b)) => SqlValue::Int(
a.checked_div(b)
.ok_or_else(|| overflowed("integer divide"))?,
),
(SqlValue::Int(a), BinaryOp::Mod, SqlValue::Int(b)) => SqlValue::Int(
a.checked_rem(b)
.ok_or_else(|| overflowed("integer modulo"))?,
),
(SqlValue::Float(a), BinaryOp::Add, SqlValue::Float(b)) => SqlValue::Float(a + b),
(SqlValue::Float(a), BinaryOp::Sub, SqlValue::Float(b)) => SqlValue::Float(a - b),
(SqlValue::Float(a), BinaryOp::Mul, SqlValue::Float(b)) => SqlValue::Float(a * b),
(SqlValue::Float(a), BinaryOp::Div, SqlValue::Float(b)) => SqlValue::Float(a / b),
(SqlValue::Float(a), BinaryOp::Mod, SqlValue::Float(b)) => SqlValue::Float(a % b),
(SqlValue::Decimal(a), BinaryOp::Add, SqlValue::Decimal(b)) => {
SqlValue::Decimal(a.checked_add(b).ok_or_else(|| overflowed("decimal add"))?)
}
(SqlValue::Decimal(a), BinaryOp::Sub, SqlValue::Decimal(b)) => SqlValue::Decimal(
a.checked_sub(b)
.ok_or_else(|| overflowed("decimal subtract"))?,
),
(SqlValue::Decimal(a), BinaryOp::Mul, SqlValue::Decimal(b)) => SqlValue::Decimal(
a.checked_mul(b)
.ok_or_else(|| overflowed("decimal multiply"))?,
),
(SqlValue::Decimal(a), BinaryOp::Div, SqlValue::Decimal(b)) => SqlValue::Decimal(
a.checked_div(b)
.ok_or_else(|| overflowed("decimal divide"))?,
),
(SqlValue::Decimal(a), BinaryOp::Mod, SqlValue::Decimal(b)) => SqlValue::Decimal(
a.checked_rem(b)
.ok_or_else(|| overflowed("decimal modulo"))?,
),
(SqlValue::Decimal(a), BinaryOp::Add, SqlValue::Int(b)) => SqlValue::Decimal(
a.checked_add(rust_decimal::Decimal::from(b))
.ok_or_else(|| overflowed("decimal add"))?,
),
(SqlValue::Int(a), BinaryOp::Add, SqlValue::Decimal(b)) => SqlValue::Decimal(
rust_decimal::Decimal::from(a)
.checked_add(b)
.ok_or_else(|| overflowed("decimal add"))?,
),
(SqlValue::Decimal(a), BinaryOp::Sub, SqlValue::Int(b)) => SqlValue::Decimal(
a.checked_sub(rust_decimal::Decimal::from(b))
.ok_or_else(|| overflowed("decimal subtract"))?,
),
(SqlValue::Int(a), BinaryOp::Sub, SqlValue::Decimal(b)) => SqlValue::Decimal(
rust_decimal::Decimal::from(a)
.checked_sub(b)
.ok_or_else(|| overflowed("decimal subtract"))?,
),
(SqlValue::Decimal(a), BinaryOp::Mul, SqlValue::Int(b)) => SqlValue::Decimal(
a.checked_mul(rust_decimal::Decimal::from(b))
.ok_or_else(|| overflowed("decimal multiply"))?,
),
(SqlValue::Int(a), BinaryOp::Mul, SqlValue::Decimal(b)) => SqlValue::Decimal(
rust_decimal::Decimal::from(a)
.checked_mul(b)
.ok_or_else(|| overflowed("decimal multiply"))?,
),
(SqlValue::Decimal(a), BinaryOp::Div, SqlValue::Int(b)) => SqlValue::Decimal(
a.checked_div(rust_decimal::Decimal::from(b))
.ok_or_else(|| overflowed("decimal divide"))?,
),
(SqlValue::Int(a), BinaryOp::Div, SqlValue::Decimal(b)) => SqlValue::Decimal(
rust_decimal::Decimal::from(a)
.checked_div(b)
.ok_or_else(|| overflowed("decimal divide"))?,
),
(SqlValue::Decimal(a), BinaryOp::Mod, SqlValue::Int(b)) => SqlValue::Decimal(
a.checked_rem(rust_decimal::Decimal::from(b))
.ok_or_else(|| overflowed("decimal modulo"))?,
),
(SqlValue::Int(a), BinaryOp::Mod, SqlValue::Decimal(b)) => SqlValue::Decimal(
rust_decimal::Decimal::from(a)
.checked_rem(b)
.ok_or_else(|| overflowed("decimal modulo"))?,
),
(SqlValue::String(a), BinaryOp::Concat, SqlValue::String(b)) => {
SqlValue::String(format!("{a}{b}"))
}
_ => return Ok(None),
}))
}
pub fn fold_function_call(name: &str, args: &[SqlExpr], registry: &FunctionRegistry) -> FoldResult {
fold_function_call_scoped(name, args, registry, FoldScope::Reuse)
}
pub fn fold_function_call_scoped(
name: &str,
args: &[SqlExpr],
registry: &FunctionRegistry,
scope: FoldScope,
) -> FoldResult {
let name_lower = name.to_lowercase();
let Some(meta) = registry.lookup(&name_lower) else {
return Ok(None);
};
if matches!(
meta.category,
FunctionCategory::Aggregate | FunctionCategory::Window
) {
return Ok(None);
}
if meta.volatility.is_volatile() && scope == FoldScope::Reuse {
return Ok(None);
}
if matches!(name_lower.as_str(), "nextval" | "currval" | "setval") {
return Ok(None);
}
let mut folded_args = Vec::with_capacity(args.len());
for arg in args {
match fold_constant_scoped(arg, registry, scope)? {
Some(v) => folded_args.push(sql_to_ndb_value(v)),
None => return Ok(None),
}
}
match nodedb_query::functions::eval_function(&name_lower, &folded_args) {
Ok(result) => Ok(Some(ndb_to_sql_value(result))),
Err(nodedb_query::EvalError::DivisionByZero) => Err(SqlError::DivisionByZero),
Err(
e @ (nodedb_query::EvalError::VectorDimensionMismatch { .. }
| nodedb_query::EvalError::ArgumentType { .. }
| nodedb_query::EvalError::InvalidJsonPath { .. }
| nodedb_query::EvalError::NumericOverflow { .. }),
) => Err(SqlError::DataException {
detail: e.to_string(),
}),
Err(nodedb_query::EvalError::UnknownFunction { .. }) => Ok(None),
}
}
fn sql_to_ndb_value(v: SqlValue) -> Value {
match v {
SqlValue::Null => Value::Null,
SqlValue::Bool(b) => Value::Bool(b),
SqlValue::Int(i) => Value::Integer(i),
SqlValue::Float(f) => Value::Float(f),
SqlValue::Decimal(d) => Value::Decimal(d),
SqlValue::String(s) => Value::String(s),
SqlValue::Bytes(b) => Value::Bytes(b),
SqlValue::Array(a) => Value::Array(a.into_iter().map(sql_to_ndb_value).collect()),
SqlValue::Timestamp(dt) => Value::NaiveDateTime(dt),
SqlValue::Timestamptz(dt) => Value::DateTime(dt),
}
}
fn ndb_to_sql_value(v: Value) -> SqlValue {
match v {
Value::Null => SqlValue::Null,
Value::Bool(b) => SqlValue::Bool(b),
Value::Integer(i) => SqlValue::Int(i),
Value::Float(f) => SqlValue::Float(f),
Value::String(s) => SqlValue::String(s),
Value::Bytes(b) => SqlValue::Bytes(b),
Value::Array(a) => SqlValue::Array(a.into_iter().map(ndb_to_sql_value).collect()),
Value::DateTime(dt) => SqlValue::Timestamptz(dt),
Value::NaiveDateTime(dt) => SqlValue::Timestamp(dt),
Value::Uuid(s) | Value::Ulid(s) | Value::Regex(s) => SqlValue::String(s),
Value::Duration(d) => SqlValue::String(d.to_human()),
Value::Decimal(d) => SqlValue::Decimal(d),
Value::Geometry(g) => sonic_rs::to_string(&g)
.map(SqlValue::String)
.unwrap_or(SqlValue::Null),
Value::Object(map) => sonic_rs::to_string(&map)
.map(SqlValue::String)
.unwrap_or(SqlValue::Null),
Value::Set(_) | Value::Range { .. } | Value::Record { .. } | Value::ArrayCell(_) => {
SqlValue::Null
}
_ => SqlValue::Null,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn fold_now_produces_timestamptz_once_only() {
let registry = FunctionRegistry::new();
let expr = SqlExpr::Function {
name: "now".into(),
args: vec![],
distinct: false,
};
assert!(
matches!(fold_constant(&expr, ®istry), Ok(None)),
"now() must not fold into a reusable plan"
);
let val = fold_constant_scoped(&expr, ®istry, FoldScope::Once)
.expect("fold must not error")
.expect("now() must fold for a single execution");
match val {
SqlValue::Timestamptz(dt) => {
assert!(dt.micros > 0, "expected post-epoch timestamp, got micros=0");
}
other => panic!("expected SqlValue::Timestamptz, got {other:?}"),
}
}
#[test]
fn negated_integer_literal_past_i64_max_stays_exact() {
let registry = FunctionRegistry::new();
let neg = |digits: &str| SqlExpr::UnaryOp {
op: UnaryOp::Neg,
expr: Box::new(SqlExpr::Literal(SqlValue::Decimal(
rust_decimal::Decimal::from_str_exact(digits).unwrap(),
))),
};
assert_eq!(
fold_constant(&neg("9223372036854775808"), ®istry).unwrap(),
Some(SqlValue::Int(i64::MIN))
);
assert_eq!(
fold_constant(&neg("18446744073709551615"), ®istry).unwrap(),
Some(SqlValue::Decimal(-rust_decimal::Decimal::from(u64::MAX)))
);
}
#[test]
fn fold_of_a_malformed_json_path_is_a_data_exception() {
let registry = FunctionRegistry::new();
let expr = SqlExpr::Function {
name: "doc_get".into(),
args: vec![
SqlExpr::Literal(SqlValue::String("{\"a\":1}".into())),
SqlExpr::Literal(SqlValue::String("$..a".into())),
],
distinct: false,
};
let err = fold_constant(&expr, ®istry).expect_err("malformed path must fail");
assert!(
matches!(&err, SqlError::DataException { detail } if detail.contains("invalid JSONPath")),
"{err:?}"
);
}
#[test]
fn fold_current_timestamp_produces_timestamptz_once_only() {
let registry = FunctionRegistry::new();
let expr = SqlExpr::Function {
name: "current_timestamp".into(),
args: vec![],
distinct: false,
};
assert!(
matches!(fold_constant(&expr, ®istry), Ok(None)),
"current_timestamp must not fold into a reusable plan"
);
assert!(matches!(
fold_constant_scoped(&expr, ®istry, FoldScope::Once),
Ok(Some(SqlValue::Timestamptz(_)))
));
}
#[test]
fn fold_unknown_function_returns_none() {
let registry = FunctionRegistry::new();
let expr = SqlExpr::Function {
name: "definitely_not_a_real_function".into(),
args: vec![],
distinct: false,
};
assert!(
fold_constant(&expr, ®istry)
.expect("fold must not error")
.is_none()
);
}
#[test]
fn fold_array_literal_recursively() {
let registry = FunctionRegistry::new();
let expr = SqlExpr::ArrayLiteral(vec![
SqlExpr::Literal(SqlValue::String("public".into())),
SqlExpr::Literal(SqlValue::Int(42)),
]);
assert_eq!(
fold_constant(&expr, ®istry),
Ok(Some(SqlValue::Array(vec![
SqlValue::String("public".into()),
SqlValue::Int(42),
])))
);
}
#[test]
fn fold_literal_arithmetic_still_works() {
let registry = FunctionRegistry::new();
let expr = SqlExpr::BinaryOp {
left: Box::new(SqlExpr::Literal(SqlValue::Int(2))),
op: BinaryOp::Add,
right: Box::new(SqlExpr::Literal(SqlValue::Int(3))),
};
assert_eq!(fold_constant(&expr, ®istry), Ok(Some(SqlValue::Int(5))));
}
#[test]
fn fold_column_ref_returns_none() {
let registry = FunctionRegistry::new();
let expr = SqlExpr::Column {
table: None,
name: "name".into(),
};
assert!(
fold_constant(&expr, ®istry)
.expect("fold must not error")
.is_none()
);
}
#[test]
fn fold_decimal_literal() {
let registry = FunctionRegistry::new();
let d = rust_decimal::Decimal::new(12345, 2); let expr = SqlExpr::Literal(SqlValue::Decimal(d));
assert_eq!(
fold_constant(&expr, ®istry),
Ok(Some(SqlValue::Decimal(d)))
);
}
#[test]
fn fold_decimal_addition() {
use rust_decimal::Decimal;
let registry = FunctionRegistry::new();
let a = Decimal::new(12345, 2); let b = Decimal::new(45678, 2); let expr = SqlExpr::BinaryOp {
left: Box::new(SqlExpr::Literal(SqlValue::Decimal(a))),
op: BinaryOp::Add,
right: Box::new(SqlExpr::Literal(SqlValue::Decimal(b))),
};
let expected = Decimal::new(58023, 2); assert_eq!(
fold_constant(&expr, ®istry),
Ok(Some(SqlValue::Decimal(expected)))
);
}
#[test]
fn fold_decimal_negation() {
use rust_decimal::Decimal;
let registry = FunctionRegistry::new();
let d = Decimal::new(100, 0);
let expr = SqlExpr::UnaryOp {
op: crate::types::UnaryOp::Neg,
expr: Box::new(SqlExpr::Literal(SqlValue::Decimal(d))),
};
assert_eq!(
fold_constant(&expr, ®istry),
Ok(Some(SqlValue::Decimal(-d)))
);
}
#[test]
fn fold_st_geohash() {
let registry = FunctionRegistry::new();
let expr = SqlExpr::Function {
name: "st_geohash".into(),
args: vec![
SqlExpr::UnaryOp {
op: UnaryOp::Neg,
expr: Box::new(SqlExpr::Literal(SqlValue::Float(122.4))),
},
SqlExpr::Literal(SqlValue::Float(37.8)),
SqlExpr::Literal(SqlValue::Int(6)),
],
distinct: false,
};
let v = fold_constant(&expr, ®istry).expect("fold must not error");
match v {
Some(SqlValue::String(ref s)) if !s.is_empty() => {}
other => panic!("expected non-empty SqlValue::String, got {other:?}"),
}
}
#[test]
fn fold_st_distance_nested_st_point() {
let registry = FunctionRegistry::new();
let make_point = |lng: f64, lat: f64| SqlExpr::Function {
name: "st_point".into(),
args: vec![
SqlExpr::Literal(SqlValue::Float(lng)),
SqlExpr::Literal(SqlValue::Float(lat)),
],
distinct: false,
};
let expr = SqlExpr::Function {
name: "st_distance".into(),
args: vec![make_point(-122.4, 37.8), make_point(-87.6, 41.8)],
distinct: false,
};
let v = fold_constant(&expr, ®istry).expect("fold must not error");
match v {
Some(SqlValue::Float(d)) => {
assert!(d > 0.0, "distance should be positive, got {d}");
}
other => panic!("expected SqlValue::Float, got {other:?}"),
}
}
#[test]
fn fold_decimal_int_widening() {
use rust_decimal::Decimal;
let registry = FunctionRegistry::new();
let d = Decimal::new(100, 0); let expr = SqlExpr::BinaryOp {
left: Box::new(SqlExpr::Literal(SqlValue::Decimal(d))),
op: BinaryOp::Add,
right: Box::new(SqlExpr::Literal(SqlValue::Int(50))),
};
let expected = Decimal::new(150, 0);
assert_eq!(
fold_constant(&expr, ®istry),
Ok(Some(SqlValue::Decimal(expected)))
);
}
}