Skip to content
Snippets Groups Projects
Commit 7fbcf5aa authored by EmirVildanov's avatar EmirVildanov Committed by Denis Smirnov
Browse files

feat: refactor expressions iterators

parent 0a56d83c
No related branches found
No related tags found
1 merge request!1414sbroad import
//! IR tree traversal module.
use super::{Nodes, Plan};
use crate::ir::expression::Expression;
use std::cell::RefCell;
trait TreeIterator<'nodes> {
fn get_current(&self) -> usize;
fn get_child(&self) -> &RefCell<usize>;
fn get_nodes(&self) -> &'nodes Nodes;
fn handle_trim(&mut self, expr: &'nodes Expression) -> Option<&'nodes usize> {
let Expression::Trim {
pattern, target, ..
} = expr
else {
panic!("Trim expected")
};
let child_step = *self.get_child().borrow();
match child_step {
0 => {
*self.get_child().borrow_mut() += 1;
match pattern {
Some(_) => pattern.as_ref(),
None => Some(target),
}
}
1 => {
*self.get_child().borrow_mut() += 1;
match pattern {
Some(_) => Some(target),
None => None,
}
}
_ => None,
}
}
fn handle_left_right_children(&mut self, expr: &'nodes Expression) -> Option<&'nodes usize> {
let (Expression::Bool { left, right, .. }
| Expression::Arithmetic { left, right, .. }
| Expression::Concat { left, right, .. }) = expr
else {
panic!("Expected expression with left and right children")
};
let child_step = *self.get_child().borrow();
if child_step == 0 {
*self.get_child().borrow_mut() += 1;
return Some(left);
} else if child_step == 1 {
*self.get_child().borrow_mut() += 1;
return Some(right);
}
None
}
fn handle_single_child(&mut self, expr: &'nodes Expression) -> Option<&'nodes usize> {
let (Expression::Alias { child, .. }
| Expression::ExprInParentheses { child }
| Expression::Cast { child, .. }
| Expression::Unary { child, .. }) = expr
else {
panic!("Expected expression with single child")
};
let step = *self.get_child().borrow();
*self.get_child().borrow_mut() += 1;
if step == 0 {
return Some(child);
}
None
}
fn handle_case_iter(&mut self, expr: &'nodes Expression) -> Option<&'nodes usize> {
let Expression::Case {
search_expr,
when_blocks,
else_expr,
} = expr
else {
panic!("Case expression expected");
};
let mut child_step = *self.get_child().borrow();
*self.get_child().borrow_mut() += 1;
if let Some(search_expr) = search_expr {
if child_step == 0 {
return Some(search_expr);
}
child_step -= 1;
}
let when_blocks_index = child_step / 2;
let index_reminder = child_step % 2;
return if when_blocks_index < when_blocks.len() {
let (cond_expr, res_expr) = when_blocks
.get(when_blocks_index)
.expect("When block must have been found.");
return match index_reminder {
0 => Some(cond_expr),
1 => Some(res_expr),
_ => unreachable!("Impossible reminder"),
};
} else if when_blocks_index == when_blocks.len() && index_reminder == 0 {
if let Some(else_expr) = else_expr {
Some(else_expr)
} else {
None
}
} else {
None
};
}
}
trait PlanTreeIterator<'plan>: TreeIterator<'plan> {
......
......@@ -101,148 +101,78 @@ impl<'n> Iterator for AggregateIterator<'n> {
fn expression_next<'nodes>(
iter: &mut impl ExpressionTreeIterator<'nodes>,
) -> Option<&'nodes usize> {
match iter.get_nodes().arena.get(iter.get_current()) {
Some(Node::Expression(
Expression::Alias { child, .. }
| Expression::ExprInParentheses { child, .. }
| Expression::Cast { child, .. }
| Expression::Unary { child, .. },
)) => {
let child_step = *iter.get_child().borrow();
if child_step == 0 {
*iter.get_child().borrow_mut() += 1;
return Some(child);
}
None
}
Some(Node::Expression(Expression::Case {
search_expr,
when_blocks,
else_expr,
})) => {
let mut child_step = *iter.get_child().borrow();
*iter.get_child().borrow_mut() += 1;
if let Some(search_expr) = search_expr {
if child_step == 0 {
return Some(search_expr);
}
child_step -= 1;
}
let when_blocks_index = child_step / 2;
let index_reminder = child_step % 2;
return if when_blocks_index < when_blocks.len() {
let (cond_expr, res_expr) = when_blocks
.get(when_blocks_index)
.expect("When block must have been found.");
return match index_reminder {
0 => Some(cond_expr),
1 => Some(res_expr),
_ => unreachable!("Impossible reminder"),
};
} else if when_blocks_index == when_blocks.len() && index_reminder == 0 {
if let Some(else_expr) = else_expr {
Some(else_expr)
} else {
None
}
} else {
None
};
}
Some(Node::Expression(
Expression::Bool { left, right, .. }
| Expression::Arithmetic { left, right, .. }
| Expression::Concat { left, right },
)) => {
let child_step = *iter.get_child().borrow();
if child_step == 0 {
*iter.get_child().borrow_mut() += 1;
return Some(left);
} else if child_step == 1 {
*iter.get_child().borrow_mut() += 1;
return Some(right);
}
None
}
Some(Node::Expression(Expression::Trim {
pattern, target, ..
})) => {
let child_step = *iter.get_child().borrow();
match child_step {
0 => {
*iter.get_child().borrow_mut() += 1;
match pattern {
Some(_) => pattern.as_ref(),
None => Some(target),
let node = iter.get_nodes().arena.get(iter.get_current());
match node {
Some(node) => {
match node {
Node::Expression(expr) => {
match expr {
Expression::Alias { .. }
| Expression::ExprInParentheses { .. }
| Expression::Cast { .. }
| Expression::Unary { .. } => iter.handle_single_child(expr),
Expression::Bool { .. }
| Expression::Arithmetic { .. }
| Expression::Concat { .. } => iter.handle_left_right_children(expr),
Expression::Row { list, .. } => {
let child_step = *iter.get_child().borrow();
let mut is_leaf = false;
// Check on the first step, if the row contains only leaf nodes.
if child_step == 0 {
is_leaf = true;
for col in list {
if !matches!(
iter.get_nodes().arena.get(*col),
Some(Node::Expression(
Expression::Reference { .. }
| Expression::Constant { .. }
))
) {
is_leaf = false;
break;
}
}
}
// If the row contains only leaf nodes (or we don't want to go deeper
// into the row tree for some reasons), skip traversal.
if !is_leaf || !iter.get_make_row_leaf() {
match list.get(child_step) {
None => return None,
Some(child) => {
*iter.get_child().borrow_mut() += 1;
return Some(child);
}
}
}
None
}
Expression::StableFunction { children, .. } => {
let child_step = *iter.get_child().borrow();
match children.get(child_step) {
None => None,
Some(child) => {
*iter.get_child().borrow_mut() += 1;
Some(child)
}
}
}
Expression::Trim { .. } => iter.handle_trim(expr),
Expression::Case { .. } => iter.handle_case_iter(expr),
Expression::Constant { .. }
| Expression::Reference { .. }
| Expression::CountAsterisk => None,
}
}
1 => {
*iter.get_child().borrow_mut() += 1;
match pattern {
Some(_) => Some(target),
None => None,
}
}
_ => None,
}
}
Some(Node::Expression(Expression::Row { list, .. })) => {
let child_step = *iter.get_child().borrow();
let mut is_leaf = false;
// Check on the first step, if the row contains only leaf nodes.
if child_step == 0 {
is_leaf = true;
for col in list {
if !matches!(
iter.get_nodes().arena.get(*col),
Some(Node::Expression(
Expression::Reference { .. } | Expression::Constant { .. }
))
) {
is_leaf = false;
break;
}
}
}
// If the row contains only leaf nodes (or we don't want to go deeper
// into the row tree for some reasons), skip traversal.
if !is_leaf || !iter.get_make_row_leaf() {
match list.get(child_step) {
None => return None,
Some(child) => {
*iter.get_child().borrow_mut() += 1;
return Some(child);
}
}
}
None
}
Some(Node::Expression(Expression::StableFunction { children, .. })) => {
let child_step = *iter.get_child().borrow();
match children.get(child_step) {
None => None,
Some(child) => {
*iter.get_child().borrow_mut() += 1;
Some(child)
}
Node::Acl(_)
| Node::Block(_)
| Node::Ddl(_)
| Node::Relational(_)
| Node::Parameter => None,
}
}
Some(
Node::Expression(
Expression::Constant { .. }
| Expression::Reference { .. }
| Expression::CountAsterisk,
)
| Node::Relational(_)
| Node::Parameter
| Node::Ddl(_)
| Node::Acl(_)
| Node::Block(_),
)
| None => None,
None => None,
}
}
......@@ -196,88 +196,16 @@ fn subtree_next<'plan>(
if let Some(child) = iter.get_nodes().arena.get(iter.get_current()) {
return match child {
Node::Parameter | Node::Ddl(..) | Node::Acl(..) | Node::Block(..) => None,
Node::Expression(exp) => match exp {
Expression::Alias { child, .. }
| Expression::ExprInParentheses { child }
| Expression::Cast { child, .. }
| Expression::Unary { child, .. } => {
let step = *iter.get_child().borrow();
*iter.get_child().borrow_mut() += 1;
if step == 0 {
return Some(child);
}
None
}
Expression::Case {
search_expr,
when_blocks,
else_expr,
} => {
let mut child_step = *iter.get_child().borrow();
*iter.get_child().borrow_mut() += 1;
if let Some(search_expr) = search_expr {
if child_step == 0 {
return Some(search_expr);
}
child_step -= 1;
}
let when_blocks_index = child_step / 2;
let index_reminder = child_step % 2;
return if when_blocks_index < when_blocks.len() {
let (cond_expr, res_expr) = when_blocks
.get(when_blocks_index)
.expect("When block must have been found.");
return match index_reminder {
0 => Some(cond_expr),
1 => Some(res_expr),
_ => unreachable!("Impossible reminder"),
};
} else if when_blocks_index == when_blocks.len() && index_reminder == 0 {
if let Some(else_expr) = else_expr {
Some(else_expr)
} else {
None
}
} else {
None
};
}
Expression::Bool { left, right, .. }
| Expression::Arithmetic { left, right, .. }
| Expression::Concat { left, right } => {
let child_step = *iter.get_child().borrow();
if child_step == 0 {
*iter.get_child().borrow_mut() += 1;
return Some(left);
} else if child_step == 1 {
*iter.get_child().borrow_mut() += 1;
return Some(right);
}
None
}
Expression::Trim {
pattern, target, ..
} => {
let child_step = *iter.get_child().borrow();
match child_step {
0 => {
*iter.get_child().borrow_mut() += 1;
match pattern {
Some(_) => pattern.as_ref(),
None => Some(target),
}
}
1 => {
*iter.get_child().borrow_mut() += 1;
match pattern {
Some(_) => Some(target),
None => None,
}
}
_ => None,
}
}
Node::Expression(expr) => match expr {
Expression::Alias { .. }
| Expression::ExprInParentheses { .. }
| Expression::Cast { .. }
| Expression::Unary { .. } => iter.handle_single_child(expr),
Expression::Case { .. } => iter.handle_case_iter(expr),
Expression::Bool { .. }
| Expression::Arithmetic { .. }
| Expression::Concat { .. } => iter.handle_left_right_children(expr),
Expression::Trim { .. } => iter.handle_trim(expr),
Expression::Row { list, .. }
| Expression::StableFunction { children: list, .. } => {
let child_step = *iter.get_child().borrow();
......@@ -299,7 +227,7 @@ fn subtree_next<'plan>(
// for selection filter or a join condition, we need to check whether
// the reference points to an **additional** sub-query and then traverse
// into it. Otherwise, stop traversal.
let Ok(parent_id) = exp.get_parent() else {
let Ok(parent_id) = expr.get_parent() else {
return None;
};
if let Ok(rel_id) = iter
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment