use std::collections::HashMap;
use crate::ast::{
Expr, Follow, FollowDestination, OpNode, OrderSpec, Query, SelectItem, Subquery, Where,
};
use crate::errors::{OqxError, Result};
pub fn resolve_aliases(q: &Query) -> Result<Query> {
let r#where = match &q.r#where {
None => None,
Some(w) => Some(resolve_nested_where(
resolve_where(&q.select, w.clone()).map_err(OqxError::parse)?,
)?),
};
Ok(Query {
source: resolve_expr(&q.source)?,
from: q.from.iter().map(resolve_expr).collect::<Result<_>>()?,
r#where,
select: resolve_select(&q.select).map_err(OqxError::parse)?,
order_by: resolve_order(q.order_by.as_deref())?,
consumer: q.consumer,
follow: q.follow.as_ref().map(resolve_follow).transpose()?,
frontier: q.frontier.as_ref().map(resolve_expr).transpose()?,
distinct: q.distinct,
values: q.values,
limit: q.limit.as_ref().map(resolve_expr).transpose()?,
offset: q.offset.as_ref().map(resolve_expr).transpose()?,
span: q.span,
})
}
pub fn resolve_subquery(s: &Subquery) -> Result<Subquery> {
let r#where = match &s.r#where {
None => None,
Some(w) => Some(resolve_nested_where(
resolve_where(&s.select, w.clone()).map_err(OqxError::parse)?,
)?),
};
Ok(Subquery {
from: s.from.iter().map(resolve_expr).collect::<Result<_>>()?,
r#where,
select: resolve_select(&s.select).map_err(OqxError::parse)?,
order_by: resolve_order(s.order_by.as_deref())?,
follow: s.follow.as_ref().map(resolve_follow).transpose()?,
frontier: s.frontier.as_ref().map(resolve_expr).transpose()?,
depth: s.depth,
values: s.values,
limit: s.limit.as_ref().map(resolve_expr).transpose()?,
offset: s.offset.as_ref().map(resolve_expr).transpose()?,
span: s.span,
})
}
fn resolve_order(o: Option<&[OrderSpec]>) -> Result<Option<Vec<OrderSpec>>> {
o.map(|specs| {
specs
.iter()
.map(|s| {
Ok(OrderSpec {
expr: resolve_expr(&s.expr)?,
desc: s.desc,
span: s.span,
})
})
.collect::<Result<Vec<_>>>()
})
.transpose()
}
fn resolve_op(op: &OpNode) -> Result<OpNode> {
Ok(OpNode {
receiver: resolve_expr(&op.receiver)?,
op: op.op,
sub: resolve_subquery(&op.sub)?,
count_cmp: op.count_cmp.clone(),
distinct: op.distinct,
span: op.span,
})
}
fn resolve_follow(f: &Follow) -> Result<Follow> {
Ok(Follow {
destinations: f
.destinations
.iter()
.map(|d| match d {
FollowDestination::Block(op) => {
Ok(FollowDestination::Block(Box::new(resolve_op(op)?)))
}
FollowDestination::Relation(e) => Ok(FollowDestination::Relation(resolve_expr(e)?)),
})
.collect::<Result<_>>()?,
distinct: f.distinct,
span: f.span,
})
}
fn resolve_nested_where(w: Where) -> Result<Where> {
Ok(match w {
Where::And { parts, span } => Where::And {
parts: parts
.into_iter()
.map(resolve_nested_where)
.collect::<Result<_>>()?,
span,
},
Where::Or { parts, span } => Where::Or {
parts: parts
.into_iter()
.map(resolve_nested_where)
.collect::<Result<_>>()?,
span,
},
Where::Not { expr, span } => Where::Not {
expr: Box::new(resolve_nested_where(*expr)?),
span,
},
Where::Scalar { expr, span } => Where::Scalar {
expr: resolve_expr(&expr)?,
span,
},
Where::Op(op) => Where::Op(Box::new(resolve_op(&op)?)),
})
}
fn resolve_expr(e: &Expr) -> Result<Expr> {
map_expr(e.clone(), &mut |x| match x {
Expr::Op(op) => Some(resolve_op(&op).map(|op| Expr::Op(Box::new(op)))),
_ => None,
})
}
fn map_expr<F>(e: Expr, f: &mut F) -> Result<Expr>
where
F: FnMut(Expr) -> Option<Result<Expr>>,
{
if let Some(hit) = f(e.clone()) {
return hit;
}
let mut go = |x: Box<Expr>| -> Result<Box<Expr>> { map_expr(*x, f).map(Box::new) };
Ok(match e {
Expr::Member { recv, name, span } => Expr::Member {
recv: go(recv)?,
name,
span,
},
Expr::Call {
recv,
name,
args,
span,
} => {
let recv = match recv {
Some(r) => Some(go(r)?),
None => None,
};
let mut mapped = Vec::with_capacity(args.len());
for a in args {
mapped.push(*go(Box::new(a))?);
}
Expr::Call {
recv,
name,
args: mapped,
span,
}
}
Expr::Unary { op, expr, span } => Expr::Unary {
op,
expr: go(expr)?,
span,
},
Expr::Required { expr, span } => Expr::Required {
expr: go(expr)?,
span,
},
Expr::Binary {
op,
left,
right,
span,
} => Expr::Binary {
op,
left: go(left)?,
right: go(right)?,
span,
},
Expr::Logical {
op,
left,
right,
span,
} => Expr::Logical {
op,
left: go(left)?,
right: go(right)?,
span,
},
Expr::In { left, right, span } => Expr::In {
left: go(left)?,
right: go(right)?,
span,
},
Expr::Range {
lo,
hi,
exclusive_end,
span,
} => Expr::Range {
lo: match lo {
Some(b) => Some(go(b)?),
None => None,
},
hi: match hi {
Some(b) => Some(go(b)?),
None => None,
},
exclusive_end,
span,
},
Expr::Tuple { items, span } => {
let mut mapped = Vec::with_capacity(items.len());
for x in items {
mapped.push(*go(Box::new(x))?);
}
Expr::Tuple {
items: mapped,
span,
}
}
Expr::Op(mut op) => {
op.receiver = *go(Box::new(op.receiver))?;
Expr::Op(op)
}
other @ (Expr::Lit { .. }
| Expr::Ident { .. }
| Expr::Binding { .. }
| Expr::Outer { .. }) => other,
})
}
fn alias_name(it: &SelectItem) -> Option<&str> {
match it {
SelectItem::Collect { name, .. } if !name.is_empty() => Some(name),
SelectItem::Field { name, lift: 0, .. } if !name.is_empty() => Some(name),
_ => None,
}
}
fn alias_refs(it: &SelectItem, index: &HashMap<&str, usize>) -> Vec<String> {
let own = alias_name(it);
let mut out: Vec<String> = Vec::new();
fn visit(e: &Expr, own: Option<&str>, index: &HashMap<&str, usize>, out: &mut Vec<String>) {
match e {
Expr::Ident { name, .. } => {
if Some(name.as_str()) != own
&& index.contains_key(name.as_str())
&& !out.iter().any(|n| n == name)
{
out.push(name.clone());
}
}
Expr::Lit { .. } | Expr::Outer { .. } | Expr::Binding { .. } => {}
Expr::Member { recv, .. } => visit(recv, own, index, out),
Expr::Call { recv, args, .. } => {
if let Some(r) = recv {
visit(r, own, index, out);
}
args.iter().for_each(|a| visit(a, own, index, out));
}
Expr::Unary { expr, .. } | Expr::Required { expr, .. } => visit(expr, own, index, out),
Expr::Binary { left, right, .. }
| Expr::Logical { left, right, .. }
| Expr::In { left, right, .. } => {
visit(left, own, index, out);
visit(right, own, index, out);
}
Expr::Range { lo, hi, .. } => {
for b in [lo, hi].into_iter().flatten() {
visit(b, own, index, out);
}
}
Expr::Tuple { items, .. } => items.iter().for_each(|x| visit(x, own, index, out)),
Expr::Op(op) => visit(&op.receiver, own, index, out),
}
}
match it {
SelectItem::Collect { op, .. } => visit(&op.receiver, own, index, &mut out),
SelectItem::Field { expr, .. } => visit(expr, own, index, &mut out),
}
out
}
fn cycle_path(
select: &[SelectItem],
index: &HashMap<&str, usize>,
from: usize,
to: &str,
) -> Option<Vec<String>> {
fn go(
select: &[SelectItem],
index: &HashMap<&str, usize>,
i: usize,
to: &str,
seen: &mut Vec<usize>,
path: Vec<String>,
) -> Option<Vec<String>> {
if seen.contains(&i) {
return None;
}
seen.push(i);
let it = &select[i];
let name = alias_name(it)
.expect("indexed items are aliasable")
.to_string();
for r in alias_refs(it, index) {
let mut next = path.clone();
next.push(name.clone());
if r == to {
next.push(to.to_string());
return Some(next);
}
if let Some(hit) = go(select, index, index[r.as_str()], to, seen, next) {
return Some(hit);
}
}
None
}
go(select, index, from, to, &mut Vec::new(), Vec::new())
}
pub fn resolve_select(select: &[SelectItem]) -> std::result::Result<Vec<SelectItem>, String> {
let mut index: HashMap<&str, usize> = HashMap::new();
for (i, it) in select.iter().enumerate() {
if let Some(n) = alias_name(it) {
index.insert(n, i);
}
}
let wrap = |r: Result<Expr>| r.map_err(|e| e.message);
let mut resolved: HashMap<String, Expr> = HashMap::new();
let mut out = Vec::with_capacity(select.len());
for (i, it) in select.iter().enumerate() {
let own = alias_name(it);
let mut failure: Option<String> = None;
let mut subst = |e: Expr| -> Option<Result<Expr>> {
let Expr::Ident { name, .. } = &e else {
return None;
};
if Some(name.as_str()) == own {
return Some(Ok(e));
}
if let Some(r) = resolved.get(name.as_str()) {
return Some(Ok(r.clone()));
}
let j = *index.get(name.as_str())?;
if j <= i {
return None;
}
let msg = match own.and_then(|o| cycle_path(select, &index, j, o)) {
Some(back) => format!(
"select aliases form a cycle: {} — an alias cannot depend on itself",
std::iter::once(own.expect("a cycle has an owner").to_string())
.chain(back)
.collect::<Vec<_>>()
.join(" → ")
),
None => format!(
"select alias '{name}' is used before it is defined — a select item may reference only the items to its left"
),
};
failure.get_or_insert(msg);
Some(Ok(e))
};
let item = match it {
SelectItem::Collect { name, op, span } => {
let receiver = map_expr(op.receiver.clone(), &mut subst).map_err(|e| e.message)?;
if let Some(msg) = failure {
return Err(msg);
}
let op = resolve_op(&OpNode {
receiver,
op: op.op,
sub: op.sub.clone(),
count_cmp: op.count_cmp.clone(),
distinct: op.distinct,
span: op.span,
})
.map_err(|e| e.message)?;
if let Some(o) = own {
resolved.insert(o.to_string(), Expr::Op(Box::new(op.clone())));
}
SelectItem::Collect {
name: name.clone(),
op: Box::new(op),
span: *span,
}
}
SelectItem::Field {
name,
expr,
lift,
span,
} => {
let substituted = map_expr(expr.clone(), &mut subst).map_err(|e| e.message)?;
if let Some(msg) = failure {
return Err(msg);
}
let expr = wrap(resolve_expr(&substituted))?;
if let Some(o) = own {
resolved.insert(o.to_string(), expr.clone());
}
match expr {
Expr::Op(op) => SelectItem::Collect {
name: name.clone(),
op,
span: *span,
},
expr => SelectItem::Field {
name: name.clone(),
expr,
lift: *lift,
span: *span,
},
}
}
};
out.push(item);
}
Ok(out)
}
pub fn resolve_where(select: &[SelectItem], r#where: Where) -> std::result::Result<Where, String> {
if select.is_empty() {
return Ok(r#where);
}
let mut aliases: HashMap<&str, &SelectItem> = HashMap::new();
for it in select {
if let Some(n) = alias_name(it) {
aliases.insert(n, it);
}
}
if aliases.is_empty() {
return Ok(r#where);
}
let mut resolving: Vec<String> = Vec::new();
walk_where(&aliases, &mut resolving, r#where)
}
fn walk_where(
aliases: &HashMap<&str, &SelectItem>,
resolving: &mut Vec<String>,
w: Where,
) -> std::result::Result<Where, String> {
Ok(match w {
Where::And { parts, span } => Where::And {
parts: parts
.into_iter()
.map(|p| walk_where(aliases, resolving, p))
.collect::<std::result::Result<_, _>>()?,
span,
},
Where::Or { parts, span } => Where::Or {
parts: parts
.into_iter()
.map(|p| walk_where(aliases, resolving, p))
.collect::<std::result::Result<_, _>>()?,
span,
},
Where::Not { expr, span } => Where::Not {
expr: Box::new(walk_where(aliases, resolving, *expr)?),
span,
},
Where::Scalar { expr, span } => {
if let Expr::Ident { name, .. } = &expr {
if let Some(SelectItem::Collect { op, .. }) = aliases.get(name.as_str()) {
return Ok(Where::Op(op.clone()));
}
}
Where::Scalar {
expr: subst(aliases, resolving, expr)?,
span,
}
}
Where::Op(op) => {
let mut op = *op;
op.receiver = subst(aliases, resolving, op.receiver)?;
Where::Op(Box::new(op))
}
})
}
fn subst(
aliases: &HashMap<&str, &SelectItem>,
resolving: &mut Vec<String>,
e: Expr,
) -> std::result::Result<Expr, String> {
let mut failure: Option<String> = None;
let out = map_expr(e, &mut |x| {
let Expr::Ident { name, .. } = &x else {
return None;
};
let a = *aliases.get(name.as_str())?;
if resolving.last() == Some(name) {
return Some(Ok(x));
}
if let Some(i) = resolving.iter().position(|r| r == name) {
let cycle = resolving[i..]
.iter()
.map(String::as_str)
.chain(std::iter::once(name.as_str()))
.collect::<Vec<_>>()
.join(" → ");
failure.get_or_insert(format!(
"select aliases form a cycle: {cycle} — an alias used in `where` cannot depend on itself"
));
return Some(Ok(x));
}
match a {
SelectItem::Collect { op, .. } => {
failure.get_or_insert(format!(
"select alias '{name}' is a {} {{ … }} block — in `where` it can only stand alone as a non-empty test, not inside an expression",
op.op.as_str()
));
Some(Ok(x))
}
SelectItem::Field { expr, .. } => {
resolving.push(name.clone());
let r = subst(aliases, resolving, expr.clone());
resolving.pop();
match r {
Ok(e) => Some(Ok(e)),
Err(msg) => {
failure.get_or_insert(msg);
Some(Ok(x))
}
}
}
}
})
.map_err(|e| e.message)?;
match failure {
Some(msg) => Err(msg),
None => Ok(out),
}
}