From e8007e3325119821e1d9d21a73ff485596a12928 Mon Sep 17 00:00:00 2001
From: Denis Smirnov <sd@picodata.io>
Date: Fri, 28 Oct 2022 18:15:58 +0700
Subject: [PATCH] refactoring: tree traversal
Split different traversal iterators to separate modules.
---
sbroad-core/src/ir/tree.rs | 450 +-------------------------
sbroad-core/src/ir/tree/and.rs | 76 +++++
sbroad-core/src/ir/tree/eq_class.rs | 76 +++++
sbroad-core/src/ir/tree/expression.rs | 143 ++++++++
sbroad-core/src/ir/tree/relation.rs | 84 +++++
sbroad-core/src/ir/tree/subtree.rs | 221 +++++++++++++
6 files changed, 612 insertions(+), 438 deletions(-)
create mode 100644 sbroad-core/src/ir/tree/and.rs
create mode 100644 sbroad-core/src/ir/tree/eq_class.rs
create mode 100644 sbroad-core/src/ir/tree/expression.rs
create mode 100644 sbroad-core/src/ir/tree/relation.rs
create mode 100644 sbroad-core/src/ir/tree/subtree.rs
diff --git a/sbroad-core/src/ir/tree.rs b/sbroad-core/src/ir/tree.rs
index 656dbeeecb..ea174a34ae 100644
--- a/sbroad-core/src/ir/tree.rs
+++ b/sbroad-core/src/ir/tree.rs
@@ -1,449 +1,23 @@
//! IR tree traversal module.
+use super::{Nodes, Plan};
use std::cell::RefCell;
-use std::cmp::Ordering;
-use super::expression::Expression;
-use super::operator::{Bool, Relational};
-use super::{Node, Nodes, Plan};
-
-/// Relational node's child iterator.
-///
-/// The iterator returns the next relational node in the plan tree.
-#[derive(Debug)]
-pub struct RelationalIterator<'n> {
- current: &'n usize,
- child: RefCell<usize>,
- nodes: &'n Nodes,
-}
-
-/// Expression node's children iterator.
-///
-/// The iterator returns the next child for expression
-/// nodes. It is required to use `traversal` crate.
-#[derive(Debug)]
-pub struct ExpressionIterator<'n> {
- current: &'n usize,
- child: RefCell<usize>,
- nodes: &'n Nodes,
- make_row_leaf: bool,
-}
-
-/// Expression and relational nodes iterator.
-#[derive(Debug)]
-pub struct SubtreeIterator<'p> {
- current: &'p usize,
- child: RefCell<usize>,
- plan: &'p Plan,
+trait TreeIterator<'nodes> {
+ fn get_current(&self) -> &'nodes usize;
+ fn get_child(&self) -> &RefCell<usize>;
+ fn get_nodes(&self) -> &'nodes Nodes;
}
-/// Children iterator for "and"-ed equivalent expressions.
-///
-/// The iterator returns the next child for the chained `Bool::And`
-/// and `Bool::Eq` nodes.
-#[derive(Debug)]
-pub struct EqClassIterator<'n> {
- current: &'n usize,
- child: RefCell<usize>,
- nodes: &'n Nodes,
+trait PlanTreeIterator<'plan>: TreeIterator<'plan> {
+ fn get_plan(&self) -> &'plan Plan;
}
-/// Children iterator for "and"-ed expression chains.
-///
-/// The iterator returns the next child for the chained `Bool::And` nodes.
-#[derive(Debug)]
-pub struct AndIterator<'n> {
- current: &'n usize,
- child: RefCell<usize>,
- nodes: &'n Nodes,
-}
-
-impl<'n> Nodes {
- #[must_use]
- pub fn expr_iter(&'n self, current: &'n usize, make_row_leaf: bool) -> ExpressionIterator<'n> {
- ExpressionIterator {
- current,
- child: RefCell::new(0),
- nodes: self,
- make_row_leaf,
- }
- }
-
- #[must_use]
- pub fn eq_iter(&'n self, current: &'n usize) -> EqClassIterator<'n> {
- EqClassIterator {
- current,
- child: RefCell::new(0),
- nodes: self,
- }
- }
-
- #[must_use]
- pub fn and_iter(&'n self, current: &'n usize) -> AndIterator<'n> {
- AndIterator {
- current,
- child: RefCell::new(0),
- nodes: self,
- }
- }
-
- #[must_use]
- pub fn rel_iter(&'n self, current: &'n usize) -> RelationalIterator<'n> {
- RelationalIterator {
- current,
- child: RefCell::new(0),
- nodes: self,
- }
- }
-}
-
-impl<'p> Plan {
- #[must_use]
- pub fn subtree_iter(&'p self, current: &'p usize) -> SubtreeIterator<'p> {
- SubtreeIterator {
- current,
- child: RefCell::new(0),
- plan: self,
- }
- }
-}
-
-impl<'n> Iterator for ExpressionIterator<'n> {
- type Item = &'n usize;
-
- fn next(&mut self) -> Option<Self::Item> {
- match self.nodes.arena.get(*self.current) {
- Some(Node::Expression(
- Expression::Alias { child, .. }
- | Expression::Cast { child, .. }
- | Expression::Unary { child, .. },
- )) => {
- let child_step = *self.child.borrow();
- if child_step == 0 {
- *self.child.borrow_mut() += 1;
- return Some(child);
- }
- None
- }
- Some(Node::Expression(
- Expression::Bool { left, right, .. } | Expression::Concat { left, right },
- )) => {
- let child_step = *self.child.borrow();
- if child_step == 0 {
- *self.child.borrow_mut() += 1;
- return Some(left);
- } else if child_step == 1 {
- *self.child.borrow_mut() += 1;
- return Some(right);
- }
- None
- }
- Some(Node::Expression(Expression::Row { list, .. })) => {
- let child_step = *self.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!(
- self.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 || !self.make_row_leaf {
- match list.get(child_step) {
- None => return None,
- Some(child) => {
- *self.child.borrow_mut() += 1;
- return Some(child);
- }
- }
- }
-
- None
- }
- Some(Node::Expression(Expression::StableFunction { children, .. })) => {
- let child_step = *self.child.borrow();
- match children.get(child_step) {
- None => None,
- Some(child) => {
- *self.child.borrow_mut() += 1;
- Some(child)
- }
- }
- }
- Some(
- Node::Expression(Expression::Constant { .. } | Expression::Reference { .. })
- | Node::Relational(_)
- | Node::Parameter,
- )
- | None => None,
- }
- }
-}
-
-impl<'n> Iterator for EqClassIterator<'n> {
- type Item = &'n usize;
-
- fn next(&mut self) -> Option<Self::Item> {
- if let Some(Node::Expression(Expression::Bool {
- left, op, right, ..
- })) = self.nodes.arena.get(*self.current)
- {
- if (*op != Bool::And) && (*op != Bool::Eq) {
- return None;
- }
- let child_step = *self.child.borrow();
- if child_step == 0 {
- *self.child.borrow_mut() += 1;
- return Some(left);
- } else if child_step == 1 {
- *self.child.borrow_mut() += 1;
- return Some(right);
- }
- None
- } else {
- None
- }
- }
-}
-
-impl<'n> Iterator for AndIterator<'n> {
- type Item = &'n usize;
-
- fn next(&mut self) -> Option<Self::Item> {
- let node = self.nodes.arena.get(*self.current);
- if let Some(Node::Expression(Expression::Bool {
- left, op, right, ..
- })) = node
- {
- if *op != Bool::And {
- return None;
- }
- let child_step = *self.child.borrow();
- if child_step == 0 {
- *self.child.borrow_mut() += 1;
- return Some(left);
- } else if child_step == 1 {
- *self.child.borrow_mut() += 1;
- return Some(right);
- }
- None
- } else {
- None
- }
- }
-}
-
-impl<'n> Iterator for RelationalIterator<'n> {
- type Item = &'n usize;
-
- fn next(&mut self) -> Option<Self::Item> {
- match self.nodes.arena.get(*self.current) {
- Some(Node::Relational(
- Relational::Except { children, .. }
- | Relational::InnerJoin { children, .. }
- | Relational::Insert { children, .. }
- | Relational::Motion { children, .. }
- | Relational::Projection { children, .. }
- | Relational::ScanSubQuery { children, .. }
- | Relational::Selection { children, .. }
- | Relational::UnionAll { children, .. }
- | Relational::Values { children, .. }
- | Relational::ValuesRow { children, .. },
- )) => {
- let step = *self.child.borrow();
- if step < children.len() {
- *self.child.borrow_mut() += 1;
- return children.get(step);
- }
- None
- }
- Some(
- Node::Relational(Relational::ScanRelation { .. })
- | Node::Expression(_)
- | Node::Parameter,
- )
- | None => None,
- }
- }
-}
-
-impl<'p> Iterator for SubtreeIterator<'p> {
- type Item = &'p usize;
-
- #[allow(clippy::too_many_lines)]
- fn next(&mut self) -> Option<Self::Item> {
- if let Some(child) = self.plan.nodes.arena.get(*self.current) {
- return match child {
- Node::Parameter => None,
- Node::Expression(exp) => match exp {
- Expression::Alias { child, .. }
- | Expression::Cast { child, .. }
- | Expression::Unary { child, .. } => {
- let step = *self.child.borrow();
- *self.child.borrow_mut() += 1;
- if step == 0 {
- return Some(child);
- }
- None
- }
- Expression::Bool { left, right, .. } | Expression::Concat { left, right } => {
- let child_step = *self.child.borrow();
- if child_step == 0 {
- *self.child.borrow_mut() += 1;
- return Some(left);
- } else if child_step == 1 {
- *self.child.borrow_mut() += 1;
- return Some(right);
- }
- None
- }
- Expression::Row { list, .. }
- | Expression::StableFunction { children: list, .. } => {
- let child_step = *self.child.borrow();
- return match list.get(child_step) {
- None => None,
- Some(child) => {
- *self.child.borrow_mut() += 1;
- Some(child)
- }
- };
- }
- Expression::Constant { .. } => None,
- Expression::Reference { .. } => {
- let step = *self.child.borrow();
- if step == 0 {
- *self.child.borrow_mut() += 1;
-
- // At first we need to detect the place where the reference is used:
- // 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 parent_id = match exp.get_parent() {
- Ok(parent_id) => parent_id,
- Err(_) => return None,
- };
- if let Ok(rel_id) =
- self.plan.get_relational_from_reference_node(*self.current)
- {
- match self.plan.get_relation_node(*rel_id) {
- Ok(rel_node) if rel_node.is_subquery() => {
- // Check if the sub-query is an additional one.
- let parent = self.plan.get_relation_node(parent_id);
- let mut is_additional = false;
- if let Ok(Relational::Selection { children, .. }) = parent {
- if children.iter().skip(1).any(|&c| c == *rel_id) {
- is_additional = true;
- }
- }
- if let Ok(Relational::InnerJoin { children, .. }) = parent {
- if children.iter().skip(2).any(|&c| c == *rel_id) {
- is_additional = true;
- }
- }
- if is_additional {
- return Some(rel_id);
- }
- }
- _ => {}
- }
- }
- }
- None
- }
- },
-
- Node::Relational(r) => match r {
- Relational::InnerJoin {
- children,
- condition,
- ..
- } => {
- let step = *self.child.borrow();
-
- *self.child.borrow_mut() += 1;
- match step.cmp(&2) {
- Ordering::Less => {
- return children.get(step);
- }
- Ordering::Equal => {
- return Some(condition);
- }
- Ordering::Greater => None,
- }
- }
-
- Relational::Except { children, .. }
- | Relational::Insert { children, .. }
- | Relational::Motion { children, .. }
- | Relational::ScanSubQuery { children, .. }
- | Relational::UnionAll { children, .. } => {
- let step = *self.child.borrow();
- if step < children.len() {
- *self.child.borrow_mut() += 1;
- return children.get(step);
- }
- None
- }
- Relational::Values {
- output, children, ..
- }
- | Relational::Projection {
- output, children, ..
- } => {
- let step = *self.child.borrow();
- *self.child.borrow_mut() += 1;
- if step == 0 {
- return Some(output);
- }
- if step <= children.len() {
- return children.get(step - 1);
- }
- None
- }
- Relational::Selection {
- children, filter, ..
- } => {
- let step = *self.child.borrow();
-
- *self.child.borrow_mut() += 1;
- match step.cmp(&1) {
- Ordering::Less => {
- return children.get(step);
- }
- Ordering::Equal => {
- return Some(filter);
- }
- Ordering::Greater => None,
- }
- }
- Relational::ValuesRow { data, .. } => {
- let step = *self.child.borrow();
-
- *self.child.borrow_mut() += 1;
- if step == 0 {
- return Some(data);
- }
- None
- }
- Relational::ScanRelation { .. } => None,
- },
- };
- }
- None
- }
-}
+pub mod and;
+pub mod eq_class;
+pub mod expression;
+pub mod relation;
+pub mod subtree;
#[cfg(test)]
mod tests;
diff --git a/sbroad-core/src/ir/tree/and.rs b/sbroad-core/src/ir/tree/and.rs
new file mode 100644
index 0000000000..f38259e079
--- /dev/null
+++ b/sbroad-core/src/ir/tree/and.rs
@@ -0,0 +1,76 @@
+use std::cell::RefCell;
+
+use super::TreeIterator;
+use crate::ir::expression::Expression;
+use crate::ir::operator::Bool;
+use crate::ir::{Node, Nodes};
+
+trait AndTreeIterator<'nodes>: TreeIterator<'nodes> {}
+
+/// Children iterator for "and"-ed expression chains.
+///
+/// The iterator returns the next child for the chained `Bool::And` nodes.
+#[derive(Debug)]
+pub struct AndIterator<'n> {
+ current: &'n usize,
+ child: RefCell<usize>,
+ nodes: &'n Nodes,
+}
+
+impl<'nodes> TreeIterator<'nodes> for AndIterator<'nodes> {
+ fn get_current(&self) -> &'nodes usize {
+ self.current
+ }
+
+ fn get_child(&self) -> &RefCell<usize> {
+ &self.child
+ }
+
+ fn get_nodes(&self) -> &'nodes Nodes {
+ self.nodes
+ }
+}
+
+impl<'nodes> AndTreeIterator<'nodes> for AndIterator<'nodes> {}
+
+impl<'n> Nodes {
+ #[must_use]
+ pub fn and_iter(&'n self, current: &'n usize) -> AndIterator<'n> {
+ AndIterator {
+ current,
+ child: RefCell::new(0),
+ nodes: self,
+ }
+ }
+}
+
+impl<'n> Iterator for AndIterator<'n> {
+ type Item = &'n usize;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ and_next(self)
+ }
+}
+
+fn and_next<'nodes>(iter: &mut impl AndTreeIterator<'nodes>) -> Option<&'nodes usize> {
+ let node = iter.get_nodes().arena.get(*iter.get_current());
+ if let Some(Node::Expression(Expression::Bool {
+ left, op, right, ..
+ })) = node
+ {
+ if *op != Bool::And {
+ return None;
+ }
+ 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
+ } else {
+ None
+ }
+}
diff --git a/sbroad-core/src/ir/tree/eq_class.rs b/sbroad-core/src/ir/tree/eq_class.rs
new file mode 100644
index 0000000000..a96ad63b8c
--- /dev/null
+++ b/sbroad-core/src/ir/tree/eq_class.rs
@@ -0,0 +1,76 @@
+use std::cell::RefCell;
+
+use super::TreeIterator;
+use crate::ir::expression::Expression;
+use crate::ir::operator::Bool;
+use crate::ir::{Node, Nodes};
+
+trait EqClassTreeIterator<'nodes>: TreeIterator<'nodes> {}
+
+/// Children iterator for "and"-ed equivalent expressions.
+///
+/// The iterator returns the next child for the chained `Bool::And`
+/// and `Bool::Eq` nodes.
+#[derive(Debug)]
+pub struct EqClassIterator<'n> {
+ current: &'n usize,
+ child: RefCell<usize>,
+ nodes: &'n Nodes,
+}
+
+impl<'nodes> TreeIterator<'nodes> for EqClassIterator<'nodes> {
+ fn get_current(&self) -> &'nodes usize {
+ self.current
+ }
+
+ fn get_child(&self) -> &RefCell<usize> {
+ &self.child
+ }
+
+ fn get_nodes(&self) -> &'nodes Nodes {
+ self.nodes
+ }
+}
+
+impl<'nodes> EqClassTreeIterator<'nodes> for EqClassIterator<'nodes> {}
+
+impl<'n> Nodes {
+ #[must_use]
+ pub fn eq_iter(&'n self, current: &'n usize) -> EqClassIterator<'n> {
+ EqClassIterator {
+ current,
+ child: RefCell::new(0),
+ nodes: self,
+ }
+ }
+}
+
+impl<'n> Iterator for EqClassIterator<'n> {
+ type Item = &'n usize;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ eq_class_next(self)
+ }
+}
+
+fn eq_class_next<'nodes>(iter: &mut impl EqClassTreeIterator<'nodes>) -> Option<&'nodes usize> {
+ if let Some(Node::Expression(Expression::Bool {
+ left, op, right, ..
+ })) = iter.get_nodes().arena.get(*iter.get_current())
+ {
+ if (*op != Bool::And) && (*op != Bool::Eq) {
+ return None;
+ }
+ 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
+ } else {
+ None
+ }
+}
diff --git a/sbroad-core/src/ir/tree/expression.rs b/sbroad-core/src/ir/tree/expression.rs
new file mode 100644
index 0000000000..aeafaebc88
--- /dev/null
+++ b/sbroad-core/src/ir/tree/expression.rs
@@ -0,0 +1,143 @@
+use std::cell::RefCell;
+
+use super::TreeIterator;
+use crate::ir::expression::Expression;
+use crate::ir::{Node, Nodes};
+
+trait ExpressionTreeIterator<'nodes>: TreeIterator<'nodes> {
+ fn get_make_row_leaf(&self) -> bool;
+}
+
+/// Expression node's children iterator.
+///
+/// The iterator returns the next child for expression
+/// nodes. It is required to use `traversal` crate.
+#[derive(Debug)]
+pub struct ExpressionIterator<'n> {
+ current: &'n usize,
+ child: RefCell<usize>,
+ nodes: &'n Nodes,
+ make_row_leaf: bool,
+}
+
+impl<'n> Nodes {
+ #[must_use]
+ pub fn expr_iter(&'n self, current: &'n usize, make_row_leaf: bool) -> ExpressionIterator<'n> {
+ ExpressionIterator {
+ current,
+ child: RefCell::new(0),
+ nodes: self,
+ make_row_leaf,
+ }
+ }
+}
+
+impl<'nodes> TreeIterator<'nodes> for ExpressionIterator<'nodes> {
+ fn get_current(&self) -> &'nodes usize {
+ self.current
+ }
+
+ fn get_child(&self) -> &RefCell<usize> {
+ &self.child
+ }
+
+ fn get_nodes(&self) -> &'nodes Nodes {
+ self.nodes
+ }
+}
+
+impl<'nodes> ExpressionTreeIterator<'nodes> for ExpressionIterator<'nodes> {
+ fn get_make_row_leaf(&self) -> bool {
+ self.make_row_leaf
+ }
+}
+
+impl<'n> Iterator for ExpressionIterator<'n> {
+ type Item = &'n usize;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ expression_next(self)
+ }
+}
+
+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::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::Bool { 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::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)
+ }
+ }
+ }
+ Some(
+ Node::Expression(Expression::Constant { .. } | Expression::Reference { .. })
+ | Node::Relational(_)
+ | Node::Parameter,
+ )
+ | None => None,
+ }
+}
diff --git a/sbroad-core/src/ir/tree/relation.rs b/sbroad-core/src/ir/tree/relation.rs
new file mode 100644
index 0000000000..8bb0ff454c
--- /dev/null
+++ b/sbroad-core/src/ir/tree/relation.rs
@@ -0,0 +1,84 @@
+use std::cell::RefCell;
+
+use super::TreeIterator;
+use crate::ir::operator::Relational;
+use crate::ir::{Node, Nodes};
+
+trait RelationalTreeIterator<'nodes>: TreeIterator<'nodes> {}
+
+/// Relational node's child iterator.
+///
+/// The iterator returns the next relational node in the plan tree.
+#[derive(Debug)]
+pub struct RelationalIterator<'n> {
+ current: &'n usize,
+ child: RefCell<usize>,
+ nodes: &'n Nodes,
+}
+
+impl<'n> Nodes {
+ #[must_use]
+ pub fn rel_iter(&'n self, current: &'n usize) -> RelationalIterator<'n> {
+ RelationalIterator {
+ current,
+ child: RefCell::new(0),
+ nodes: self,
+ }
+ }
+}
+
+impl<'nodes> TreeIterator<'nodes> for RelationalIterator<'nodes> {
+ fn get_current(&self) -> &'nodes usize {
+ self.current
+ }
+
+ fn get_child(&self) -> &RefCell<usize> {
+ &self.child
+ }
+
+ fn get_nodes(&self) -> &'nodes Nodes {
+ self.nodes
+ }
+}
+
+impl<'nodes> RelationalTreeIterator<'nodes> for RelationalIterator<'nodes> {}
+
+impl<'n> Iterator for RelationalIterator<'n> {
+ type Item = &'n usize;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ relational_next(self)
+ }
+}
+
+fn relational_next<'nodes>(
+ iter: &mut impl RelationalTreeIterator<'nodes>,
+) -> Option<&'nodes usize> {
+ match iter.get_nodes().arena.get(*iter.get_current()) {
+ Some(Node::Relational(
+ Relational::Except { children, .. }
+ | Relational::InnerJoin { children, .. }
+ | Relational::Insert { children, .. }
+ | Relational::Motion { children, .. }
+ | Relational::Projection { children, .. }
+ | Relational::ScanSubQuery { children, .. }
+ | Relational::Selection { children, .. }
+ | Relational::UnionAll { children, .. }
+ | Relational::Values { children, .. }
+ | Relational::ValuesRow { children, .. },
+ )) => {
+ let step = *iter.get_child().borrow();
+ if step < children.len() {
+ *iter.get_child().borrow_mut() += 1;
+ return children.get(step);
+ }
+ None
+ }
+ Some(
+ Node::Relational(Relational::ScanRelation { .. })
+ | Node::Expression(_)
+ | Node::Parameter,
+ )
+ | None => None,
+ }
+}
diff --git a/sbroad-core/src/ir/tree/subtree.rs b/sbroad-core/src/ir/tree/subtree.rs
new file mode 100644
index 0000000000..0868168be5
--- /dev/null
+++ b/sbroad-core/src/ir/tree/subtree.rs
@@ -0,0 +1,221 @@
+use std::cell::RefCell;
+use std::cmp::Ordering;
+
+use super::{PlanTreeIterator, TreeIterator};
+use crate::ir::expression::Expression;
+use crate::ir::operator::Relational;
+use crate::ir::{Node, Nodes, Plan};
+
+trait SubtreePlanIterator<'plan>: PlanTreeIterator<'plan> {}
+
+/// Expression and relational nodes iterator.
+#[derive(Debug)]
+pub struct SubtreeIterator<'plan> {
+ current: &'plan usize,
+ child: RefCell<usize>,
+ plan: &'plan Plan,
+}
+
+impl<'nodes> TreeIterator<'nodes> for SubtreeIterator<'nodes> {
+ fn get_current(&self) -> &'nodes usize {
+ self.current
+ }
+
+ fn get_child(&self) -> &RefCell<usize> {
+ &self.child
+ }
+
+ fn get_nodes(&self) -> &'nodes Nodes {
+ &self.plan.nodes
+ }
+}
+
+impl<'plan> PlanTreeIterator<'plan> for SubtreeIterator<'plan> {
+ fn get_plan(&self) -> &'plan Plan {
+ self.plan
+ }
+}
+
+impl<'plan> SubtreePlanIterator<'plan> for SubtreeIterator<'plan> {}
+
+impl<'plan> Iterator for SubtreeIterator<'plan> {
+ type Item = &'plan usize;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ subtree_next(self)
+ }
+}
+
+impl<'plan> Plan {
+ #[must_use]
+ pub fn subtree_iter(&'plan self, current: &'plan usize) -> SubtreeIterator<'plan> {
+ SubtreeIterator {
+ current,
+ child: RefCell::new(0),
+ plan: self,
+ }
+ }
+}
+
+#[allow(clippy::too_many_lines)]
+fn subtree_next<'plan>(iter: &mut impl SubtreePlanIterator<'plan>) -> Option<&'plan usize> {
+ if let Some(child) = iter.get_nodes().arena.get(*iter.get_current()) {
+ return match child {
+ Node::Parameter => None,
+ Node::Expression(exp) => match exp {
+ Expression::Alias { 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::Bool { 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::Row { list, .. }
+ | Expression::StableFunction { children: list, .. } => {
+ let child_step = *iter.get_child().borrow();
+ return match list.get(child_step) {
+ None => None,
+ Some(child) => {
+ *iter.get_child().borrow_mut() += 1;
+ Some(child)
+ }
+ };
+ }
+ Expression::Constant { .. } => None,
+ Expression::Reference { .. } => {
+ let step = *iter.get_child().borrow();
+ if step == 0 {
+ *iter.get_child().borrow_mut() += 1;
+
+ // At first we need to detect the place where the reference is used:
+ // 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 parent_id = match exp.get_parent() {
+ Ok(parent_id) => parent_id,
+ Err(_) => return None,
+ };
+ if let Ok(rel_id) = iter
+ .get_plan()
+ .get_relational_from_reference_node(*iter.get_current())
+ {
+ match iter.get_plan().get_relation_node(*rel_id) {
+ Ok(rel_node) if rel_node.is_subquery() => {
+ // Check if the sub-query is an additional one.
+ let parent = iter.get_plan().get_relation_node(parent_id);
+ let mut is_additional = false;
+ if let Ok(Relational::Selection { children, .. }) = parent {
+ if children.iter().skip(1).any(|&c| c == *rel_id) {
+ is_additional = true;
+ }
+ }
+ if let Ok(Relational::InnerJoin { children, .. }) = parent {
+ if children.iter().skip(2).any(|&c| c == *rel_id) {
+ is_additional = true;
+ }
+ }
+ if is_additional {
+ return Some(rel_id);
+ }
+ }
+ _ => {}
+ }
+ }
+ }
+ None
+ }
+ },
+
+ Node::Relational(r) => match r {
+ Relational::InnerJoin {
+ children,
+ condition,
+ ..
+ } => {
+ let step = *iter.get_child().borrow();
+
+ *iter.get_child().borrow_mut() += 1;
+ match step.cmp(&2) {
+ Ordering::Less => {
+ return children.get(step);
+ }
+ Ordering::Equal => {
+ return Some(condition);
+ }
+ Ordering::Greater => None,
+ }
+ }
+
+ Relational::Except { children, .. }
+ | Relational::Insert { children, .. }
+ | Relational::Motion { children, .. }
+ | Relational::ScanSubQuery { children, .. }
+ | Relational::UnionAll { children, .. } => {
+ let step = *iter.get_child().borrow();
+ if step < children.len() {
+ *iter.get_child().borrow_mut() += 1;
+ return children.get(step);
+ }
+ None
+ }
+ Relational::Values {
+ output, children, ..
+ }
+ | Relational::Projection {
+ output, children, ..
+ } => {
+ let step = *iter.get_child().borrow();
+ *iter.get_child().borrow_mut() += 1;
+ if step == 0 {
+ return Some(output);
+ }
+ if step <= children.len() {
+ return children.get(step - 1);
+ }
+ None
+ }
+ Relational::Selection {
+ children, filter, ..
+ } => {
+ let step = *iter.get_child().borrow();
+
+ *iter.get_child().borrow_mut() += 1;
+ match step.cmp(&1) {
+ Ordering::Less => {
+ return children.get(step);
+ }
+ Ordering::Equal => {
+ return Some(filter);
+ }
+ Ordering::Greater => None,
+ }
+ }
+ Relational::ValuesRow { data, .. } => {
+ let step = *iter.get_child().borrow();
+
+ *iter.get_child().borrow_mut() += 1;
+ if step == 0 {
+ return Some(data);
+ }
+ None
+ }
+ Relational::ScanRelation { .. } => None,
+ },
+ };
+ }
+ None
+}
--
GitLab