use std::{
collections::BTreeMap,
fmt::{Display, Write},
str::FromStr,
};
use analysis::ConstantFolding;
use anyhow::{bail, Result};
use compile::{compile_delete, compile_select, compile_update, decompile_plan};
use egg::{define_language, AstDepth, Extractor, Id, Runner};
use rewrite::rewrite_rules;
use crate::{
ast::{DeleteStatement, SelectStatement, UpdateStatement},
executor::PlanNode,
Db,
};
mod analysis;
mod compile;
mod rewrite;
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone)]
struct Null;
impl Display for Null {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("NULL")
}
}
impl FromStr for Null {
type Err = ();
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s {
"NULL" => Ok(Null),
_ => Err(()),
}
}
}
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone)]
struct Binding(pub u64);
impl Display for Binding {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("?")?;
f.write_fmt(format_args!("{}", self.0))
}
}
impl FromStr for Binding {
type Err = ();
fn from_str(_s: &str) -> Result<Self, Self::Err> {
Err(())
}
}
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone)]
pub struct ColumnRef(pub String, pub String);
impl FromStr for ColumnRef {
type Err = ();
fn from_str(s: &str) -> Result<Self, Self::Err> {
let (a, b) = s.split_once('.').ok_or(())?;
Ok(Self(a.to_string(), b.to_string()))
}
}
impl Display for ColumnRef {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)?;
f.write_char('.')?;
f.write_str(&self.1)
}
}
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone)]
pub struct TableRef(pub u64, pub String);
impl FromStr for TableRef {
type Err = ();
fn from_str(s: &str) -> Result<Self, Self::Err> {
let [t, a] = s.split('|').collect::<Vec<_>>()[..] else {
return Err(());
};
Ok(Self(t.parse().or(Err(()))?, a.to_string()))
}
}
impl Display for TableRef {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_fmt(format_args!("{}", &self.0))?;
f.write_char('|')?;
f.write_str(&self.1)
}
}
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone)]
pub struct IndexRef(pub u64, pub String, pub String);
impl FromStr for IndexRef {
type Err = ();
fn from_str(s: &str) -> Result<Self, Self::Err> {
let [t, a, b] = s.split('|').collect::<Vec<_>>()[..] else {
return Err(());
};
Ok(Self(t.parse().or(Err(()))?, a.to_string(), b.to_string()))
}
}
impl Display for IndexRef {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_fmt(format_args!("{}", &self.0))?;
f.write_char('|')?;
f.write_str(&self.1)?;
f.write_char('|')?;
f.write_str(&self.2)
}
}
define_language! {
pub enum EPlanNode {
"seqscan" = SeqScan(Id),
"nonescan" = NoneScan,
"filter" = FilterScan([Id; 2]),
"table-scan-eq" = TableScanEq(Vec<Id>),
"table-scan-gt" = TableScanGt(Vec<Id>),
"table-scan-gteq" = TableScanGtEq(Vec<Id>),
"table-scan-lt" = TableScanLt(Vec<Id>),
"table-scan-lteq" = TableScanLtEq(Vec<Id>),
"select" = Select(Vec<Id>),
"delete" = Delete([Id; 2]),
"update" = Update([Id; 2]),
"join" = Join([Id; 2]),
NamedColumn(ColumnRef),
Binding(Binding),
Null(Null),
Int(i64),
Bool(bool),
String(String),
"=" = Eq([Id; 2]),
"!=" = NEq([Id; 2]),
">" = Gt([Id; 2]),
">=" = GtEq([Id; 2]),
"<" = Lt([Id; 2]),
"<=" = LtEq([Id; 2]),
"and" = And([Id; 2]),
"or" = Or([Id; 2]),
"not" = Not([Id; 1]),
"+" = Add([Id; 2]),
"-" = Sub([Id; 2]),
"*" = Mul([Id; 2]),
"/" = Div([Id; 2]),
"%" = Mod([Id; 2]),
TableRef(TableRef),
IndexRef(IndexRef),
}
}
pub fn optimize_select(db: &mut Db, select_stmnt: &SelectStatement) -> Result<PlanNode> {
let mut tables = BTreeMap::new();
for (a, t) in &select_stmnt.table_mappings {
if tables.insert(a.to_string(), t.to_string()).is_some() {
bail!("duplicate table alias {}", a)
}
}
let e = compile_select(db, select_stmnt);
let runner = Runner::<EPlanNode, ConstantFolding, ()>::default()
.with_expr(&e)
.run(&rewrite_rules(
db,
select_stmnt.version.unwrap_or(db.this_tx_id),
&tables,
)?);
let extractor = Extractor::new(&runner.egraph, AstDepth);
let (_best_cost, best) = extractor.find_best(runner.roots[0]);
#[cfg(debug_assertions)]
println!("{best}");
Ok(decompile_plan(
&select_stmnt.table_mappings.clone().into_iter().collect(),
&best,
Id::from(best.as_ref().len() - 1),
None,
))
}
pub fn optimize_delete(db: &mut Db, delete_stmnt: &DeleteStatement) -> Result<PlanNode> {
let tables = BTreeMap::from([(delete_stmnt.table.clone(), delete_stmnt.table.clone())]);
let e = compile_delete(db, delete_stmnt);
let runner = Runner::<EPlanNode, ConstantFolding, ()>::default()
.with_expr(&e)
.run(&rewrite_rules(db, db.this_tx_id, &tables)?);
let extractor = Extractor::new(&runner.egraph, AstDepth);
let (_best_cost, best) = extractor.find_best(runner.roots[0]);
#[cfg(debug_assertions)]
println!("{best}");
Ok(decompile_plan(
&tables,
&best,
Id::from(best.as_ref().len() - 1),
None,
))
}
pub fn optimize_update(db: &mut Db, update_stmnt: &UpdateStatement) -> Result<PlanNode> {
let tables = BTreeMap::from([(update_stmnt.table.clone(), update_stmnt.table.clone())]);
let e = compile_update(db, update_stmnt);
let runner = Runner::<EPlanNode, ConstantFolding, ()>::default()
.with_expr(&e)
.run(&rewrite_rules(db, db.this_tx_id, &tables)?);
let extractor = Extractor::new(&runner.egraph, AstDepth);
let (_best_cost, best) = extractor.find_best(runner.roots[0]);
#[cfg(debug_assertions)]
println!("{best}");
Ok(decompile_plan(
&tables,
&best,
Id::from(best.as_ref().len() - 1),
Some(update_stmnt.sets.clone()),
))
}