use std::collections::BTreeMap;
use std::iter::once;
use egg::{Id, Language, RecExpr};
use crate::ast::BinOp;
use crate::ast::Column;
use crate::ast::DeleteStatement;
use crate::ast::Expr;
use crate::ast::SelectStatement;
use crate::ast::Set;
use crate::ast::UnOp;
use crate::ast::UpdateStatement;
use crate::executor::Delete;
use crate::executor::Filter;
use crate::executor::Join;
use crate::executor::PlanNode;
use crate::executor::RangeScan;
use crate::executor::Select;
use crate::executor::Update;
use crate::kv::Range;
use crate::kv::RangeEnd;
use crate::value::Value;
use crate::Db;
use super::Binding;
use super::ColumnRef;
use super::IndexRef;
use super::Null;
use super::EPlanNode;
use super::TableRef;
pub fn compile_expr<T, F>(expr: &mut RecExpr<T>, ast: &Expr, f: F) -> Id
where
T: Language,
F: Fn(EPlanNode) -> T + Copy,
{
match ast {
Expr::Literal(Value::Null) => expr.add(f(EPlanNode::Null(Null))),
Expr::Literal(Value::Bool(b)) => expr.add(f(EPlanNode::Bool(*b))),
Expr::Literal(Value::String(s)) => expr.add(f(EPlanNode::String(s.to_string()))),
Expr::Literal(Value::Int(i)) => expr.add(f(EPlanNode::Int(*i))),
Expr::Column(Column(t, c)) => expr.add(f(EPlanNode::NamedColumn(ColumnRef(
t.to_string(),
c.to_string(),
)))),
Expr::Binding(b) => expr.add(f(EPlanNode::Binding(Binding(*b)))),
Expr::Unary(UnOp::Not, rhs) => {
let rhs_id = compile_expr(expr, rhs, f);
expr.add(f(EPlanNode::Not([rhs_id])))
}
Expr::Bin(lhs, BinOp::Eq, rhs) => {
let lhs_id = compile_expr(expr, lhs, f);
let rhs_id = compile_expr(expr, rhs, f);
expr.add(f(EPlanNode::Eq([lhs_id, rhs_id])))
}
Expr::Bin(lhs, BinOp::NEq, rhs) => {
let lhs_id = compile_expr(expr, lhs, f);
let rhs_id = compile_expr(expr, rhs, f);
expr.add(f(EPlanNode::NEq([lhs_id, rhs_id])))
}
Expr::Bin(lhs, BinOp::Gt, rhs) => {
let lhs_id = compile_expr(expr, lhs, f);
let rhs_id = compile_expr(expr, rhs, f);
expr.add(f(EPlanNode::Gt([lhs_id, rhs_id])))
}
Expr::Bin(lhs, BinOp::GtEq, rhs) => {
let lhs_id = compile_expr(expr, lhs, f);
let rhs_id = compile_expr(expr, rhs, f);
expr.add(f(EPlanNode::GtEq([lhs_id, rhs_id])))
}
Expr::Bin(lhs, BinOp::Lt, rhs) => {
let lhs_id = compile_expr(expr, lhs, f);
let rhs_id = compile_expr(expr, rhs, f);
expr.add(f(EPlanNode::Lt([lhs_id, rhs_id])))
}
Expr::Bin(lhs, BinOp::LtEq, rhs) => {
let lhs_id = compile_expr(expr, lhs, f);
let rhs_id = compile_expr(expr, rhs, f);
expr.add(f(EPlanNode::LtEq([lhs_id, rhs_id])))
}
Expr::Bin(lhs, BinOp::And, rhs) => {
let lhs_id = compile_expr(expr, lhs, f);
let rhs_id = compile_expr(expr, rhs, f);
expr.add(f(EPlanNode::And([lhs_id, rhs_id])))
}
Expr::Bin(lhs, BinOp::Or, rhs) => {
let lhs_id = compile_expr(expr, lhs, f);
let rhs_id = compile_expr(expr, rhs, f);
expr.add(f(EPlanNode::Or([lhs_id, rhs_id])))
}
Expr::Bin(lhs, BinOp::Add, rhs) => {
let lhs_id = compile_expr(expr, lhs, f);
let rhs_id = compile_expr(expr, rhs, f);
expr.add(f(EPlanNode::Add([lhs_id, rhs_id])))
}
Expr::Bin(lhs, BinOp::Sub, rhs) => {
let lhs_id = compile_expr(expr, lhs, f);
let rhs_id = compile_expr(expr, rhs, f);
expr.add(f(EPlanNode::Sub([lhs_id, rhs_id])))
}
Expr::Bin(lhs, BinOp::Mul, rhs) => {
let lhs_id = compile_expr(expr, lhs, f);
let rhs_id = compile_expr(expr, rhs, f);
expr.add(f(EPlanNode::Mul([lhs_id, rhs_id])))
}
Expr::Bin(lhs, BinOp::Div, rhs) => {
let lhs_id = compile_expr(expr, lhs, f);
let rhs_id = compile_expr(expr, rhs, f);
expr.add(f(EPlanNode::Div([lhs_id, rhs_id])))
}
Expr::Bin(lhs, BinOp::Mod, rhs) => {
let lhs_id = compile_expr(expr, lhs, f);
let rhs_id = compile_expr(expr, rhs, f);
expr.add(f(EPlanNode::Mod([lhs_id, rhs_id])))
}
Expr::Edge(_, _, _, e) => compile_expr(expr, e, f),
}
}
pub fn compile_update(db: &mut Db, update_statement: &UpdateStatement) -> RecExpr<EPlanNode> {
let mut expr = RecExpr::default();
let table = expr.add(EPlanNode::TableRef(TableRef(
db.this_tx_id,
update_statement.table.clone(),
)));
let seqscan = expr.add(EPlanNode::SeqScan(table));
let cond = compile_expr(
&mut expr,
&update_statement
.cond
.clone()
.unwrap_or(Expr::Literal(Value::Bool(true))),
|a| a,
);
let update = EPlanNode::Update([table, expr.add(EPlanNode::FilterScan([cond, seqscan]))]);
expr.add(update);
expr
}
pub fn compile_delete(db: &mut Db, delete_statement: &DeleteStatement) -> RecExpr<EPlanNode> {
let mut expr = RecExpr::default();
let table = expr.add(EPlanNode::TableRef(TableRef(
db.this_tx_id,
delete_statement.table.clone(),
)));
let seqscan = expr.add(EPlanNode::SeqScan(table));
let cond = compile_expr(
&mut expr,
&delete_statement
.cond
.clone()
.unwrap_or(Expr::Literal(Value::Bool(true))),
|a| a,
);
let delete = EPlanNode::Delete([table, expr.add(EPlanNode::FilterScan([cond, seqscan]))]);
expr.add(delete);
expr
}
pub fn compile_select(db: &mut Db, select_statement: &SelectStatement) -> RecExpr<EPlanNode> {
let mut expr = RecExpr::default();
let mut last_join = None;
for (a, _) in &select_statement.table_mappings {
let table_id = expr.add(EPlanNode::TableRef(TableRef(
select_statement.version.unwrap_or(db.this_tx_id),
a.to_string(),
)));
let scan = expr.add(EPlanNode::SeqScan(table_id));
if let Some(id) = last_join {
last_join = Some(expr.add(EPlanNode::Join([scan, id])));
} else {
last_join = Some(scan);
}
}
if let Some(cond) = &select_statement.cond {
let cond_id = compile_expr(&mut expr, cond, |a| a);
last_join = Some(expr.add(EPlanNode::FilterScan([cond_id, last_join.unwrap()])));
}
let select = EPlanNode::Select(
once(last_join.unwrap())
.chain(
select_statement
.expr
.iter()
.map(|e| compile_expr(&mut expr, e, |a| a)),
)
.collect(),
);
expr.add(select);
expr
}
pub fn decompile_expr(best: &RecExpr<EPlanNode>, id: Id) -> Expr {
match &best[id] {
EPlanNode::Binding(b) => Expr::Binding(b.0),
EPlanNode::Null(_) => Expr::Literal(Value::Null),
EPlanNode::Int(i) => Expr::Literal(Value::Int(*i)),
EPlanNode::Bool(b) => Expr::Literal(Value::Bool(*b)),
EPlanNode::String(s) => Expr::Literal(Value::String(s.to_string())),
EPlanNode::Not([rhs]) => Expr::Unary(UnOp::Not, Box::new(decompile_expr(best, *rhs))),
EPlanNode::Eq([lhs, rhs]) => Expr::Bin(
Box::new(decompile_expr(best, *lhs)),
BinOp::Eq,
Box::new(decompile_expr(best, *rhs)),
),
EPlanNode::NEq([lhs, rhs]) => Expr::Bin(
Box::new(decompile_expr(best, *lhs)),
BinOp::NEq,
Box::new(decompile_expr(best, *rhs)),
),
EPlanNode::Gt([lhs, rhs]) => Expr::Bin(
Box::new(decompile_expr(best, *lhs)),
BinOp::Gt,
Box::new(decompile_expr(best, *rhs)),
),
EPlanNode::GtEq([lhs, rhs]) => Expr::Bin(
Box::new(decompile_expr(best, *lhs)),
BinOp::GtEq,
Box::new(decompile_expr(best, *rhs)),
),
EPlanNode::Lt([lhs, rhs]) => Expr::Bin(
Box::new(decompile_expr(best, *lhs)),
BinOp::Lt,
Box::new(decompile_expr(best, *rhs)),
),
EPlanNode::LtEq([lhs, rhs]) => Expr::Bin(
Box::new(decompile_expr(best, *lhs)),
BinOp::LtEq,
Box::new(decompile_expr(best, *rhs)),
),
EPlanNode::And([lhs, rhs]) => Expr::Bin(
Box::new(decompile_expr(best, *lhs)),
BinOp::And,
Box::new(decompile_expr(best, *rhs)),
),
EPlanNode::Add([lhs, rhs]) => Expr::Bin(
Box::new(decompile_expr(best, *lhs)),
BinOp::Add,
Box::new(decompile_expr(best, *rhs)),
),
EPlanNode::Or([lhs, rhs]) => Expr::Bin(
Box::new(decompile_expr(best, *lhs)),
BinOp::Or,
Box::new(decompile_expr(best, *rhs)),
),
EPlanNode::Sub([lhs, rhs]) => Expr::Bin(
Box::new(decompile_expr(best, *lhs)),
BinOp::Sub,
Box::new(decompile_expr(best, *rhs)),
),
EPlanNode::Mul([lhs, rhs]) => Expr::Bin(
Box::new(decompile_expr(best, *lhs)),
BinOp::Mul,
Box::new(decompile_expr(best, *rhs)),
),
EPlanNode::Div([lhs, rhs]) => Expr::Bin(
Box::new(decompile_expr(best, *lhs)),
BinOp::Div,
Box::new(decompile_expr(best, *rhs)),
),
EPlanNode::Mod([lhs, rhs]) => Expr::Bin(
Box::new(decompile_expr(best, *lhs)),
BinOp::Mod,
Box::new(decompile_expr(best, *rhs)),
),
EPlanNode::NamedColumn(r) => Expr::Column(Column(r.0.clone(), r.1.clone())),
EPlanNode::SeqScan(_)
| EPlanNode::NoneScan
| EPlanNode::FilterScan(_)
| EPlanNode::TableScanEq(_)
| EPlanNode::TableScanGt(_)
| EPlanNode::TableScanGtEq(_)
| EPlanNode::Select(_)
| EPlanNode::Delete(_)
| EPlanNode::Update(_)
| EPlanNode::Join(_)
| EPlanNode::TableRef(_)
| EPlanNode::IndexRef(_)
| EPlanNode::TableScanLt(_)
| EPlanNode::TableScanLtEq(_) => todo!(),
}
}
pub fn decompile_plan(
table_mappings: &BTreeMap<String, String>,
best: &RecExpr<EPlanNode>,
id: Id,
sets: Option<Vec<Set>>,
) -> PlanNode {
match &best[id] {
EPlanNode::SeqScan(id) => match &best[*id] {
EPlanNode::TableRef(TableRef(tx_id, name)) => PlanNode::RangeScan(RangeScan::new(
table_mappings.get(name).unwrap().clone(),
name.to_string(),
*tx_id,
None,
Range {
backwards: false,
prefix: Vec::new(),
start: None,
end: None,
},
)),
node => todo!("{:?}", node),
},
EPlanNode::NoneScan => PlanNode::None,
EPlanNode::FilterScan([cond, child]) => PlanNode::Filter(Filter {
child: Box::new(decompile_plan(table_mappings, best, *child, None)),
cond: decompile_expr(best, *cond),
}),
EPlanNode::TableScanEq(items) => {
let (tx_id, alias, index, tail, head) = index_scan_decode_items(best, items);
PlanNode::RangeScan(RangeScan::new(
table_mappings.get(&alias).unwrap().to_string(),
alias,
tx_id,
index,
Range {
backwards: false,
prefix: {
let mut v = tail;
v.push(head);
v
},
start: None,
end: None,
},
))
}
EPlanNode::TableScanGtEq(items) => {
let (tx_id, alias, index, tail, head) = index_scan_decode_items(best, items);
PlanNode::RangeScan(RangeScan::new(
table_mappings.get(&alias).unwrap().to_string(),
alias,
tx_id,
index,
Range {
backwards: false,
prefix: tail,
start: Some(RangeEnd {
inclusive: true,
value: head,
}),
end: None,
},
))
}
EPlanNode::TableScanGt(items) => {
let (tx_id, alias, index, tail, head) = index_scan_decode_items(best, items);
PlanNode::RangeScan(RangeScan::new(
table_mappings.get(&alias).unwrap().to_string(),
alias,
tx_id,
index,
Range {
backwards: false,
prefix: tail,
start: Some(RangeEnd {
inclusive: false,
value: head,
}),
end: None,
},
))
}
EPlanNode::TableScanLtEq(items) => {
let (tx_id, alias, index, tail, head) = index_scan_decode_items(best, items);
PlanNode::RangeScan(RangeScan::new(
table_mappings.get(&alias).unwrap().to_string(),
alias,
tx_id,
index,
Range {
backwards: false,
prefix: tail,
start: None,
end: Some(RangeEnd {
inclusive: true,
value: head,
}),
},
))
}
EPlanNode::TableScanLt(items) => {
let (tx_id, alias, index, tail, head) = index_scan_decode_items(best, items);
PlanNode::RangeScan(RangeScan::new(
table_mappings.get(&alias).unwrap().to_string(),
alias,
tx_id,
index,
Range {
backwards: false,
prefix: tail,
start: None,
end: Some(RangeEnd {
inclusive: false,
value: head,
}),
},
))
}
EPlanNode::Join([left, right]) => PlanNode::Join(Join::new(
decompile_plan(table_mappings, best, *left, None),
decompile_plan(table_mappings, best, *right, None),
)),
EPlanNode::Select(items) => {
let mut iter = items.into_iter();
let child = decompile_plan(table_mappings, best, *iter.next().unwrap(), None);
let exprs = iter.map(|e| decompile_expr(best, *e)).collect::<Vec<_>>();
PlanNode::Select(Select::new(child, exprs, "select".to_string()))
}
EPlanNode::Delete([table_id, child_id]) => {
let table = match &best[*table_id] {
EPlanNode::TableRef(name) => name.clone(),
node => todo!("{:?}", node),
};
let child = decompile_plan(table_mappings, best, *child_id, None);
PlanNode::Delete(Delete::new(table.1, child))
}
EPlanNode::Update([table_id, child_id]) => {
let table = match &best[*table_id] {
EPlanNode::TableRef(name) => name.clone(),
node => todo!("{:?}", node),
};
let child = decompile_plan(table_mappings, best, *child_id, None);
PlanNode::Update(Update::new(table.1, sets.unwrap(), child))
}
EPlanNode::NamedColumn(_)
| EPlanNode::Binding(_)
| EPlanNode::Null(_)
| EPlanNode::Int(_)
| EPlanNode::Bool(_)
| EPlanNode::String(_)
| EPlanNode::Eq(_)
| EPlanNode::NEq(_)
| EPlanNode::Gt(_)
| EPlanNode::GtEq(_)
| EPlanNode::Lt(_)
| EPlanNode::LtEq(_)
| EPlanNode::Not(_)
| EPlanNode::And(_)
| EPlanNode::Or(_)
| EPlanNode::TableRef(_)
| EPlanNode::IndexRef(_)
| EPlanNode::Add(_)
| EPlanNode::Sub(_)
| EPlanNode::Mul(_)
| EPlanNode::Div(_)
| EPlanNode::Mod(_) => todo!(),
}
}
fn index_scan_decode_items(
best: &RecExpr<EPlanNode>,
items: &[Id],
) -> (u64, String, Option<String>, Vec<Expr>, Expr) {
let mut iter = items.into_iter();
let (tx_id, alias, index) = match &best[*iter.next().unwrap()] {
EPlanNode::TableRef(TableRef(tx_id, alias)) => (tx_id, alias.clone(), None),
EPlanNode::IndexRef(IndexRef(tx_id, alias, index)) => {
(tx_id, alias.clone(), Some(index.clone()))
}
_ => todo!(),
};
let mut riter = iter.rev();
let head = decompile_expr(best, *riter.next().unwrap());
let tail = riter
.rev()
.map(|e| decompile_expr(best, *e))
.collect::<Vec<_>>();
(*tx_id, alias, index, tail, head)
}