use super::*;
use plexus_parser::{
AggFn as AstAggFn, ArithOp as AstArithOp, Clause, CmpOp as AstCmpOp, Dir as AstDir,
Expr as AstExpr, Literal as AstLiteral, Query, SortDir as AstSortDir,
};
fn expr_to_plan(expr: &AstExpr, env: &HashMap<String, u32>) -> Expr {
match expr {
AstExpr::Cmp { op, lhs, rhs } => Expr::Cmp {
op: match op {
AstCmpOp::Eq => CmpOp::Eq,
AstCmpOp::Ne => CmpOp::Ne,
AstCmpOp::Lt => CmpOp::Lt,
AstCmpOp::Gt => CmpOp::Gt,
AstCmpOp::Le => CmpOp::Le,
AstCmpOp::Ge => CmpOp::Ge,
},
lhs: Box::new(expr_to_plan(lhs, env)),
rhs: Box::new(expr_to_plan(rhs, env)),
},
AstExpr::And(l, r) => Expr::And {
lhs: Box::new(expr_to_plan(l, env)),
rhs: Box::new(expr_to_plan(r, env)),
},
AstExpr::Or(l, r) => Expr::Or {
lhs: Box::new(expr_to_plan(l, env)),
rhs: Box::new(expr_to_plan(r, env)),
},
AstExpr::Not(e) => Expr::Not {
expr: Box::new(expr_to_plan(e, env)),
},
AstExpr::IsNull(e) => Expr::IsNull {
expr: Box::new(expr_to_plan(e, env)),
},
AstExpr::IsNotNull(e) => Expr::IsNotNull {
expr: Box::new(expr_to_plan(e, env)),
},
AstExpr::StartsWith { expr, pattern } => Expr::StartsWith {
expr: Box::new(expr_to_plan(expr, env)),
pattern: pattern.clone(),
},
AstExpr::EndsWith { expr, pattern } => Expr::EndsWith {
expr: Box::new(expr_to_plan(expr, env)),
pattern: pattern.clone(),
},
AstExpr::Contains { expr, pattern } => Expr::Contains {
expr: Box::new(expr_to_plan(expr, env)),
pattern: pattern.clone(),
},
AstExpr::PropAccess { var, prop } => Expr::PropAccess {
col: *env.get(var).expect("var bound"),
prop: prop.clone(),
},
AstExpr::VarRef(var) => Expr::ColRef {
idx: *env.get(var).expect("var bound"),
},
AstExpr::Lit(l) => match l {
AstLiteral::Int(v) => Expr::IntLiteral(*v),
AstLiteral::Float(v) => Expr::FloatLiteral(*v),
AstLiteral::Str(v) => Expr::StringLiteral(v.clone()),
AstLiteral::Bool(v) => Expr::BoolLiteral(*v),
AstLiteral::Null => Expr::NullLiteral,
},
AstExpr::Agg { fn_, expr } => Expr::Agg {
fn_: match fn_ {
AstAggFn::CountStar => AggFn::CountStar,
AstAggFn::Count => AggFn::Count,
AstAggFn::Sum => AggFn::Sum,
AstAggFn::Avg => AggFn::Avg,
AstAggFn::Min => AggFn::Min,
AstAggFn::Max => AggFn::Max,
AstAggFn::Collect => AggFn::Collect,
},
expr: expr.as_ref().map(|e| Box::new(expr_to_plan(e, env))),
},
AstExpr::Arith { op, lhs, rhs } => Expr::Arith {
op: match op {
AstArithOp::Add => plexus_serde::ArithOp::Add,
AstArithOp::Sub => plexus_serde::ArithOp::Sub,
AstArithOp::Mul => plexus_serde::ArithOp::Mul,
AstArithOp::Div => plexus_serde::ArithOp::Div,
},
lhs: Box::new(expr_to_plan(lhs, env)),
rhs: Box::new(expr_to_plan(rhs, env)),
},
AstExpr::Param { name } => Expr::Param {
name: name.clone(),
expected_type: None,
},
AstExpr::Case { arms, else_expr } => Expr::Case {
arms: arms
.iter()
.map(|arm| {
(
expr_to_plan(&arm.when_expr, env),
expr_to_plan(&arm.then_expr, env),
)
})
.collect(),
else_expr: else_expr.as_ref().map(|e| Box::new(expr_to_plan(e, env))),
},
AstExpr::In { .. }
| AstExpr::ListLiteral { .. }
| AstExpr::MapLiteral { .. }
| AstExpr::Exists { .. }
| AstExpr::ListComprehension { .. } => {
unimplemented!(
"expr_to_plan: Phase 12.1 expression form not supported in test compiler"
)
}
}
}
pub(super) fn compile_query_to_plan(query: &Query) -> Plan {
let mut ops = Vec::new();
let mut env: HashMap<String, u32> = HashMap::new();
for clause in &query.clauses {
match clause {
Clause::Match(m) => {
let mut row_width = 1u32;
let first = &m.pattern.nodes[0];
let first_var = first.var.clone().expect("test query needs bound node vars");
ops.push(Op::ScanNodes {
labels: first.labels.clone(),
schema: vec![ColDef {
name: first_var.clone(),
kind: ColKind::Node,
logical_type: plexus_serde::LogicalType::Unknown,
}],
must_labels: first.labels.clone(),
forbidden_labels: vec![],
graph_ref: None,
est_rows: -1,
selectivity: 1.0,
});
env.insert(first_var, 0);
for (i, rel) in m.pattern.rels.iter().enumerate() {
let src_var = m.pattern.nodes[i]
.var
.as_ref()
.expect("test query needs bound node vars");
let rel_var = rel.var.as_ref().expect("test query needs bound rel vars");
let dst_var = m.pattern.nodes[i + 1]
.var
.as_ref()
.expect("test query needs bound node vars");
let src_col = *env.get(src_var).expect("src bound");
let input = (ops.len() - 1) as u32;
ops.push(Op::Expand {
input,
src_col,
types: rel.types.clone(),
dir: match rel.dir {
AstDir::Out => ExpandDir::Out,
AstDir::In => ExpandDir::In,
AstDir::Both => ExpandDir::Both,
},
schema: vec![],
src_var: src_var.clone(),
rel_var: rel_var.clone(),
dst_var: dst_var.clone(),
legal_src_labels: m.pattern.nodes[i].labels.clone(),
legal_dst_labels: m.pattern.nodes[i + 1].labels.clone(),
graph_ref: None,
est_degree: -1.0,
});
let _ = src_var;
env.insert(rel_var.clone(), row_width);
env.insert(dst_var.clone(), row_width + 1);
row_width += 2;
}
if let Some(pred) = &m.where_ {
let input = (ops.len() - 1) as u32;
ops.push(Op::Filter {
input,
predicate: expr_to_plan(pred, &env),
});
}
}
Clause::Return(r) => {
let input = (ops.len() - 1) as u32;
let mut project_env: HashMap<String, u32> = HashMap::new();
let mut exprs = Vec::new();
let mut schema = Vec::new();
for (i, item) in r.items.iter().enumerate() {
exprs.push(expr_to_plan(&item.expr, &env));
let alias = item.alias.clone().unwrap_or_else(|| match &item.expr {
AstExpr::VarRef(v) => v.clone(),
AstExpr::PropAccess { prop, .. } => prop.clone(),
_ => format!("col_{i}"),
});
project_env.insert(alias.clone(), i as u32);
schema.push(ColDef {
name: alias,
kind: ColKind::Value,
logical_type: plexus_serde::LogicalType::Unknown,
});
}
ops.push(Op::Project {
input,
exprs,
schema,
});
if !r.order_by.is_empty() {
let mut keys = Vec::new();
let mut dirs = Vec::new();
for s in &r.order_by {
let key = match &s.expr {
AstExpr::VarRef(v) => *project_env.get(v).expect("sort key bound"),
AstExpr::PropAccess { prop, .. } => {
*project_env.get(prop).expect("sort key bound")
}
_ => panic!("unsupported sort expression in test compiler"),
};
keys.push(key);
dirs.push(match s.dir {
AstSortDir::Asc => SortDir::Asc,
AstSortDir::Desc => SortDir::Desc,
});
}
let input = (ops.len() - 1) as u32;
ops.push(Op::Sort { input, keys, dirs });
}
if r.limit.is_some() || r.skip.is_some() {
let input = (ops.len() - 1) as u32;
ops.push(Op::Limit {
input,
count: r.limit.unwrap_or(-1),
skip: r.skip.unwrap_or(0),
cursor: None,
emit_cursor: false,
});
}
let input = (ops.len() - 1) as u32;
ops.push(Op::Return { input });
}
_ => panic!("test compiler supports MATCH ... RETURN only"),
}
}
Plan {
version: version(),
root_op: (ops.len() - 1) as u32,
ops,
}
}