use reblessive::Stk;
use crate::expr::match_plan::{BindingId, BindingKind};
use crate::gql::ast::{BinaryOp, GqlExpr, GqlLiteral, Ident, SetQuantifier, TruthValue, UnaryOp};
use crate::gql::lower::binding::Registry;
use crate::gql::lower::naming;
use crate::sql::function::{Function, FunctionCall};
use crate::sql::literal::ObjectEntry;
use crate::sql::{BinaryOperator, Expr, Idiom, Literal, Param, Part, PrefixOperator};
use crate::syn::error::{SyntaxError, bail, syntax_error};
use crate::syn::token::Span;
pub(super) struct Scope<'a> {
pub(super) registry: &'a Registry,
pub(super) allow_aggregates: bool,
}
impl<'a> Scope<'a> {
fn no_aggregates(&self) -> Scope<'a> {
Scope {
registry: self.registry,
allow_aggregates: false,
}
}
}
impl Scope<'_> {
fn variable_is_optional(&self, ident: &Ident) -> bool {
self.registry.resolve(ident).map(|id| self.registry.optional_depth(id) > 0).unwrap_or(false)
}
fn binding_expr(&self, ident: &Ident, props: &[&Ident]) -> Result<Expr, SyntaxError> {
let id = self.registry.resolve(ident)?;
let kind = self.registry.kind(id);
let prop_span = props.first().map(|p| p.span);
self.binding_idiom(&ident.name, kind, props, prop_span)
}
fn binding_idiom(
&self,
name: &str,
kind: BindingKind,
props: &[&Ident],
prop_span: Option<Span>,
) -> Result<Expr, SyntaxError> {
if !props.is_empty() {
match kind {
BindingKind::EdgeGroup | BindingKind::Path => {
bail!(
"Property access on a group or path variable is not supported yet",
@prop_span.unwrap_or_else(Span::empty) => "return the variable itself"
);
}
BindingKind::Node | BindingKind::Edge => {}
}
}
let mut parts = Vec::with_capacity(props.len() + 1);
parts.push(Part::Field(name.to_owned().into()));
parts.extend(props.iter().map(|p| Part::Field(p.name.clone().into())));
Ok(Expr::Idiom(Idiom(parts)))
}
}
pub(super) async fn lower_value(
stk: &mut Stk,
expr: &GqlExpr,
scope: &Scope<'_>,
) -> Result<Expr, SyntaxError> {
match expr {
GqlExpr::Literal(lit, _) => Ok(lower_literal(lit)),
GqlExpr::Param {
name,
span,
} => lower_param(name, *span),
GqlExpr::Variable(ident) => scope.binding_expr(ident, &[]),
GqlExpr::Property(..) => lower_property(stk, expr, scope).await,
GqlExpr::Unary {
op,
expr: operand,
..
} => {
let operand = stk.run(|stk| lower_value(stk, operand, scope)).await?;
let op = match op {
UnaryOp::Not => PrefixOperator::Not,
UnaryOp::Neg => PrefixOperator::Negate,
UnaryOp::Plus => PrefixOperator::Positive,
};
Ok(Expr::Prefix {
op,
expr: Box::new(operand),
})
}
GqlExpr::Binary {
left,
op,
right,
span,
} => {
let op = binary_op(*op, *span)?;
let left = stk.run(|stk| lower_value(stk, left, scope)).await?;
let right = stk.run(|stk| lower_value(stk, right, scope)).await?;
Ok(Expr::Binary {
left: Box::new(left),
op,
right: Box::new(right),
})
}
GqlExpr::IsBool {
..
}
| GqlExpr::IsNull {
..
} => lower_truth_test(stk, expr, false, scope).await,
GqlExpr::FunctionCall {
name,
quantifier,
star,
args,
span,
} => lower_function(stk, name, *quantifier, *star, args, *span, scope).await,
GqlExpr::List(items, _) => {
let mut exprs = Vec::with_capacity(items.len());
for item in items {
exprs.push(stk.run(|stk| lower_value(stk, item, scope)).await?);
}
Ok(Expr::Literal(Literal::Array(exprs)))
}
GqlExpr::Map(fields, _) => {
let mut entries = Vec::with_capacity(fields.len());
for (key, value) in fields {
let value = stk.run(|stk| lower_value(stk, value, scope)).await?;
entries.push(ObjectEntry {
key: key.name.clone().into(),
value,
});
}
Ok(Expr::Literal(Literal::Object(entries)))
}
}
}
pub(super) async fn lower_predicate(
stk: &mut Stk,
expr: &GqlExpr,
negated: bool,
scope: &Scope<'_>,
) -> Result<Expr, SyntaxError> {
match expr {
GqlExpr::Unary {
op: UnaryOp::Not,
expr: inner,
..
} => stk.run(|stk| lower_predicate(stk, inner, !negated, scope)).await,
GqlExpr::Binary {
op: op @ (BinaryOp::And | BinaryOp::Or),
left,
right,
..
} => {
let lowered_left = stk.run(|stk| lower_predicate(stk, left, negated, scope)).await?;
let lowered_right = stk.run(|stk| lower_predicate(stk, right, negated, scope)).await?;
let and = (*op == BinaryOp::And) != negated;
Ok(Expr::Binary {
left: Box::new(lowered_left),
op: if and {
BinaryOperator::And
} else {
BinaryOperator::Or
},
right: Box::new(lowered_right),
})
}
GqlExpr::Binary {
op: BinaryOp::Xor,
span,
..
} => reject_xor(*span),
GqlExpr::Binary {
op:
op @ (BinaryOp::Eq
| BinaryOp::Neq
| BinaryOp::Lt
| BinaryOp::Lte
| BinaryOp::Gt
| BinaryOp::Gte),
left,
right,
span,
} => lower_comparison(stk, *op, negated, left, right, *span, scope).await,
GqlExpr::IsBool {
..
}
| GqlExpr::IsNull {
..
} => lower_truth_test(stk, expr, negated, scope).await,
GqlExpr::Literal(GqlLiteral::Bool(b), _) => Ok(Expr::Literal(Literal::Bool(*b != negated))),
other => {
let value = stk.run(|stk| lower_value(stk, other, scope)).await?;
Ok(equality(value, Literal::Bool(!negated), false))
}
}
}
pub(super) async fn lower_prop_equality(
stk: &mut Stk,
binding: BindingId,
binding_name: &str,
key: &Ident,
value: &GqlExpr,
scope: &Scope<'_>,
) -> Result<Expr, SyntaxError> {
let kind = scope.registry.kind(binding);
let left = scope.binding_idiom(binding_name, kind, &[key], Some(key.span))?;
let guards = if nullable(value, scope) {
let mut atoms = vec![left.clone()];
for atom in nullable_atoms(value, scope) {
let atom = stk.run(|stk| lower_value(stk, atom, scope)).await?;
if !atoms.contains(&atom) {
atoms.push(atom);
}
}
guard_conjuncts(&atoms)
} else {
Vec::new()
};
let right = stk.run(|stk| lower_value(stk, value, scope)).await?;
let comparison = Expr::Binary {
left: Box::new(left),
op: BinaryOperator::Equal,
right: Box::new(right),
};
Ok(match and_chain(guards) {
Some(guards) => Expr::Binary {
left: Box::new(guards),
op: BinaryOperator::And,
right: Box::new(comparison),
},
None => comparison,
})
}
pub(super) fn and_chain(exprs: impl IntoIterator<Item = Expr>) -> Option<Expr> {
exprs.into_iter().reduce(|left, right| Expr::Binary {
left: Box::new(left),
op: BinaryOperator::And,
right: Box::new(right),
})
}
async fn lower_comparison(
stk: &mut Stk,
op: BinaryOp,
negated: bool,
left: &GqlExpr,
right: &GqlExpr,
span: Span,
scope: &Scope<'_>,
) -> Result<Expr, SyntaxError> {
let op = if negated {
complement(op)
} else {
op
};
let guard_atoms: Vec<&GqlExpr> = match op {
BinaryOp::Lt | BinaryOp::Lte | BinaryOp::Gt | BinaryOp::Gte => {
let mut atoms = nullable_atoms(left, scope);
atoms.extend(nullable_atoms(right, scope));
atoms
}
BinaryOp::Eq if nullable(left, scope) && nullable(right, scope) => {
let mut atoms = nullable_atoms(left, scope);
atoms.extend(nullable_atoms(right, scope));
atoms
}
BinaryOp::Neq if nullable(left, scope) || nullable(right, scope) => {
let mut atoms = nullable_atoms(left, scope);
atoms.extend(nullable_atoms(right, scope));
atoms
}
_ => Vec::new(),
};
let mut atoms: Vec<Expr> = Vec::new();
for atom in guard_atoms {
let atom = stk.run(|stk| lower_value(stk, atom, scope)).await?;
if !atoms.contains(&atom) {
atoms.push(atom);
}
}
let lowered_left = stk.run(|stk| lower_value(stk, left, scope)).await?;
let lowered_right = stk.run(|stk| lower_value(stk, right, scope)).await?;
let comparison = Expr::Binary {
left: Box::new(lowered_left),
op: binary_op(op, span)?,
right: Box::new(lowered_right),
};
Ok(match and_chain(guard_conjuncts(&atoms)) {
Some(guards) => Expr::Binary {
left: Box::new(guards),
op: BinaryOperator::And,
right: Box::new(comparison),
},
None => comparison,
})
}
fn guard_conjuncts(atoms: &[Expr]) -> Vec<Expr> {
let mut out = Vec::with_capacity(atoms.len() * 2);
for atom in atoms {
out.push(Expr::Binary {
left: Box::new(atom.clone()),
op: BinaryOperator::NotEqual,
right: Box::new(Expr::Literal(Literal::None)),
});
out.push(Expr::Binary {
left: Box::new(atom.clone()),
op: BinaryOperator::NotEqual,
right: Box::new(Expr::Literal(Literal::Null)),
});
}
out
}
async fn lower_truth_test(
stk: &mut Stk,
expr: &GqlExpr,
outer_negated: bool,
scope: &Scope<'_>,
) -> Result<Expr, SyntaxError> {
match expr {
GqlExpr::IsNull {
expr: operand,
negated,
..
} => {
let value = stk.run(|stk| lower_value(stk, operand, scope)).await?;
Ok(null_test(value, outer_negated != *negated))
}
GqlExpr::IsBool {
expr: operand,
value,
negated,
..
} => {
let operand = stk.run(|stk| lower_value(stk, operand, scope)).await?;
let negated = outer_negated != *negated;
match value {
TruthValue::True => Ok(equality(operand, Literal::Bool(true), negated)),
TruthValue::False => Ok(equality(operand, Literal::Bool(false), negated)),
TruthValue::Unknown => Ok(null_test(operand, negated)),
}
}
other => Err(syntax_error!("Internal error: not a truth test", @other.span())),
}
}
fn null_test(value: Expr, negated: bool) -> Expr {
let (cmp, combine) = if negated {
(BinaryOperator::NotEqual, BinaryOperator::And)
} else {
(BinaryOperator::Equal, BinaryOperator::Or)
};
Expr::Binary {
left: Box::new(Expr::Binary {
left: Box::new(value.clone()),
op: cmp.clone(),
right: Box::new(Expr::Literal(Literal::Null)),
}),
op: combine,
right: Box::new(Expr::Binary {
left: Box::new(value),
op: cmp,
right: Box::new(Expr::Literal(Literal::None)),
}),
}
}
fn equality(left: Expr, literal: Literal, negated: bool) -> Expr {
Expr::Binary {
left: Box::new(left),
op: if negated {
BinaryOperator::NotEqual
} else {
BinaryOperator::Equal
},
right: Box::new(Expr::Literal(literal)),
}
}
fn nullable_atoms<'a>(expr: &'a GqlExpr, scope: &Scope<'_>) -> Vec<&'a GqlExpr> {
let mut out = Vec::new();
let mut stack = vec![expr];
while let Some(e) = stack.pop() {
match e {
GqlExpr::Property(..)
| GqlExpr::Param {
..
} => out.push(e),
GqlExpr::Variable(ident) if scope.variable_is_optional(ident) => out.push(e),
GqlExpr::Unary {
op: UnaryOp::Neg | UnaryOp::Plus,
expr,
..
} => stack.push(expr),
GqlExpr::Binary {
op: BinaryOp::Concat | BinaryOp::Add | BinaryOp::Sub | BinaryOp::Mul | BinaryOp::Div,
left,
right,
..
} => {
stack.push(right);
stack.push(left);
}
_ => {}
}
}
out
}
fn nullable(expr: &GqlExpr, scope: &Scope<'_>) -> bool {
let mut stack = vec![expr];
while let Some(e) = stack.pop() {
match e {
GqlExpr::Property(..)
| GqlExpr::Param {
..
}
| GqlExpr::Literal(GqlLiteral::Null, _) => return true,
GqlExpr::Variable(ident) if scope.variable_is_optional(ident) => return true,
GqlExpr::Unary {
op: UnaryOp::Neg | UnaryOp::Plus,
expr,
..
} => stack.push(expr),
GqlExpr::Binary {
op: BinaryOp::Concat | BinaryOp::Add | BinaryOp::Sub | BinaryOp::Mul | BinaryOp::Div,
left,
right,
..
} => {
stack.push(right);
stack.push(left);
}
_ => {}
}
}
false
}
fn complement(op: BinaryOp) -> BinaryOp {
match op {
BinaryOp::Eq => BinaryOp::Neq,
BinaryOp::Neq => BinaryOp::Eq,
BinaryOp::Lt => BinaryOp::Gte,
BinaryOp::Lte => BinaryOp::Gt,
BinaryOp::Gt => BinaryOp::Lte,
BinaryOp::Gte => BinaryOp::Lt,
other => other,
}
}
fn binary_op(op: BinaryOp, span: Span) -> Result<BinaryOperator, SyntaxError> {
Ok(match op {
BinaryOp::Or => BinaryOperator::Or,
BinaryOp::Xor => return reject_xor(span),
BinaryOp::And => BinaryOperator::And,
BinaryOp::Eq => BinaryOperator::Equal,
BinaryOp::Neq => BinaryOperator::NotEqual,
BinaryOp::Lt => BinaryOperator::LessThan,
BinaryOp::Lte => BinaryOperator::LessThanEqual,
BinaryOp::Gt => BinaryOperator::MoreThan,
BinaryOp::Gte => BinaryOperator::MoreThanEqual,
BinaryOp::Concat | BinaryOp::Add => BinaryOperator::Add,
BinaryOp::Sub => BinaryOperator::Subtract,
BinaryOp::Mul => BinaryOperator::Multiply,
BinaryOp::Div => BinaryOperator::Divide,
})
}
fn reject_xor<T>(span: Span) -> Result<T, SyntaxError> {
bail!(
"`XOR` is not supported yet",
@span => "rewrite `a XOR b` as `(a OR b) AND NOT (a AND b)`"
);
}
fn lower_param(name: &str, span: Span) -> Result<Expr, SyntaxError> {
naming::validate_param_name(name, span)?;
Ok(Expr::Param(Param::new(name)))
}
fn lower_literal(literal: &GqlLiteral) -> Expr {
Expr::Literal(match literal {
GqlLiteral::Null => Literal::Null,
GqlLiteral::Bool(b) => Literal::Bool(*b),
GqlLiteral::Integer(i) => Literal::Integer(*i),
GqlLiteral::Float(f) => Literal::Float(*f),
GqlLiteral::String(s) => Literal::String(s.clone().into()),
})
}
async fn lower_property(
stk: &mut Stk,
expr: &GqlExpr,
scope: &Scope<'_>,
) -> Result<Expr, SyntaxError> {
let mut names: Vec<&Ident> = Vec::new();
let mut base = expr;
while let GqlExpr::Property(inner, name, _) = base {
names.push(name);
base = inner;
}
names.reverse();
match base {
GqlExpr::Variable(ident) => scope.binding_expr(ident, &names),
other => {
let start = stk.run(|stk| lower_value(stk, other, scope)).await?;
let mut parts = match start {
Expr::Idiom(idiom) => idiom.0,
start => vec![Part::Start(start)],
};
parts.extend(names.iter().map(|name| Part::Field(name.name.clone().into())));
Ok(Expr::Idiom(Idiom(parts)))
}
}
}
enum AggregateTarget {
Count,
Mapped(&'static str),
}
fn aggregate_target(name: &str) -> Option<AggregateTarget> {
Some(match name {
"count" => AggregateTarget::Count,
"sum" => AggregateTarget::Mapped("math::sum"),
"collect" | "collect_list" => AggregateTarget::Mapped("array::group"),
"min" => AggregateTarget::Mapped("math::min"),
"max" => AggregateTarget::Mapped("math::max"),
"avg" => AggregateTarget::Mapped("math::mean"),
_ => return None,
})
}
const UNSUPPORTED_AGGREGATES: &[&str] =
&["percentile_cont", "percentile_disc", "stddev_pop", "stddev_samp"];
async fn lower_function(
stk: &mut Stk,
name: &Ident,
quantifier: Option<SetQuantifier>,
star: Option<Span>,
args: &[GqlExpr],
span: Span,
scope: &Scope<'_>,
) -> Result<Expr, SyntaxError> {
let lowered = name.name.to_ascii_lowercase();
let Some(target) = aggregate_target(&lowered) else {
if star.is_some()
|| quantifier.is_some()
|| UNSUPPORTED_AGGREGATES.contains(&lowered.as_str())
{
bail!("Aggregate functions are not supported yet", @span);
}
bail!("The function `{}` is not supported yet", name.name, @span);
};
if !scope.allow_aggregates {
bail!(
"Aggregate functions are only allowed in RETURN items and ORDER BY keys",
@span => "an aggregate cannot appear in WHERE, GROUP BY, or inside another aggregate"
);
}
if quantifier.is_some() {
bail!(
"DISTINCT/ALL inside an aggregate is not supported yet",
@span => "remove the set quantifier"
);
}
match target {
AggregateTarget::Count => {
if star.is_some() {
return Ok(function_call("count", Vec::new()));
}
let arg = single_arg(name, args, span)?;
let arg_scope = scope.no_aggregates();
let lowered = stk.run(|stk| lower_value(stk, arg, &arg_scope)).await?;
Ok(function_call("count", vec![null_test(lowered, true)]))
}
AggregateTarget::Mapped(surreal_name) => {
if let Some(star) = star {
bail!(
"`{}(*)` is not supported; only `count(*)` takes `*`",
name.name,
@star => "pass an expression to aggregate"
);
}
let arg = single_arg(name, args, span)?;
let arg_scope = scope.no_aggregates();
let lowered = stk.run(|stk| lower_value(stk, arg, &arg_scope)).await?;
Ok(function_call(surreal_name, vec![lowered]))
}
}
}
fn single_arg<'a>(
name: &Ident,
args: &'a [GqlExpr],
span: Span,
) -> Result<&'a GqlExpr, SyntaxError> {
match args {
[arg] => Ok(arg),
_ => bail!(
"`{}` takes exactly one argument",
name.name,
@span => "aggregate over a single expression"
),
}
}
fn function_call(name: &str, arguments: Vec<Expr>) -> Expr {
Expr::FunctionCall(Box::new(FunctionCall {
receiver: Function::Normal(name.to_owned()),
arguments,
}))
}
pub(super) fn gql_contains_aggregate(expr: &GqlExpr) -> bool {
let mut stack = vec![expr];
while let Some(e) = stack.pop() {
match e {
GqlExpr::FunctionCall {
name,
star,
args,
..
} => {
if star.is_some() || aggregate_target(&name.name.to_ascii_lowercase()).is_some() {
return true;
}
stack.extend(args.iter());
}
GqlExpr::Property(inner, _, _) => stack.push(inner),
GqlExpr::Unary {
expr,
..
}
| GqlExpr::IsBool {
expr,
..
}
| GqlExpr::IsNull {
expr,
..
} => stack.push(expr),
GqlExpr::Binary {
left,
right,
..
} => {
stack.push(left);
stack.push(right);
}
GqlExpr::List(items, _) => stack.extend(items.iter()),
GqlExpr::Map(fields, _) => stack.extend(fields.iter().map(|(_, v)| v)),
GqlExpr::Literal(..)
| GqlExpr::Param {
..
}
| GqlExpr::Variable(_) => {}
}
}
false
}