feat: prepare commit for motions

src/ir.rs

@@ -34,7 +34,7 @@ pub mod value;
 ///
 /// Enum was chosen as we don't want to mess with dynamic
 /// dispatching and its performance penalties.
-#[derive(Serialize, Deserialize, PartialEq, Debug)]
+#[derive(Serialize, Deserialize, PartialEq, Debug, Clone)]
 pub enum Node {
     Expression(Expression),
     Relational(Relational),

src/ir/distribution.rs

@@ -30,7 +30,7 @@ impl Key {
 }
 
 /// Tuple distribution in the cluster.
-#[derive(Serialize, Deserialize, PartialEq, Debug)]
+#[derive(Serialize, Deserialize, PartialEq, Debug, Clone)]
 pub enum Distribution {
     /// A tuple can be located on any data node
     Any,

src/ir/expression.rs

@@ -24,7 +24,7 @@ use std::collections::HashSet;
 /// and should not be changed. Thanks to this fact we always know the
 /// name of any column in the tuple that should simplify AST
 /// deserialization.
-#[derive(Serialize, Deserialize, PartialEq, Debug)]
+#[derive(Serialize, Deserialize, PartialEq, Debug, Clone)]
 pub enum Expression {
     /// Expression name.
     ///

src/ir/operator.rs

@@ -77,7 +77,7 @@ impl Display for Bool {
 ///
 /// Transforms input tuple(s) into the output one using
 /// relation algebra logic.
-#[derive(Serialize, Deserialize, PartialEq, Debug)]
+#[derive(Serialize, Deserialize, PartialEq, Debug, Clone)]
 pub enum Relational {
     /// Inner Join
     InnerJoin {

src/ir/transformation.rs

@@ -2,4 +2,5 @@
 //!
 //! Contains rule-based transformations.
 
+pub mod distribution_enforcement;
 pub mod equality_propagation;

src/ir/transformation/distribution_enforcement.rs

+use crate::errors::QueryPlannerError;
+use crate::ir::distribution::{Distribution, Key};
+use crate::ir::expression::Expression;
+use crate::ir::operator::{Bool, Relational};
+use crate::ir::{Node, Plan};
+use std::collections::{HashMap, HashSet};
+use traversal::DftPost;
+
+impl Plan {
+    /// Get a list of relational nodes in a DFS post order.
+    ///
+    /// # Errors
+    /// - plan doesn't contain the top node
+    fn get_relational_nodes_dfs_post(&self) -> Result<Vec<usize>, QueryPlannerError> {
+        let top = self.get_top()?;
+        let mut nodes: Vec<usize> = Vec::new();
+
+        let post_tree = DftPost::new(&top, |node| self.nodes.rel_iter(node));
+        for (_, node) in post_tree {
+            nodes.push(*node);
+        }
+        Ok(nodes)
+    }
+
+    ///
+    /// # Errors
+    /// - top node is invalid
+    fn get_expr_nodes_dfs_post(&self, top: usize) -> Vec<usize> {
+        let mut nodes: Vec<usize> = Vec::new();
+
+        let post_tree = DftPost::new(&top, |node| self.nodes.expr_iter(node));
+        for (_, node) in post_tree {
+            nodes.push(*node);
+        }
+        nodes
+    }
+
+    fn resolve_join_conflicts(
+        &mut self,
+        rel: usize,
+        expr: usize,
+        map: &HashMap<usize, usize>,
+    ) -> Result<(), QueryPlannerError> {
+        // TODO: resolve joins
+        Ok(())
+    }
+
+    fn resolve_subquery_conflicts(
+        &mut self,
+        rel: usize,
+        expr: usize,
+        map: &HashMap<usize, usize>,
+    ) -> Result<(), QueryPlannerError> {
+        // TODO: resolve sub queries
+        Ok(())
+    }
+
+    /// Add motion nodes to the plan tree.
+    ///
+    /// # Errors
+    /// - failed to get relational nodes (plan is invalid?)
+    /// - failed to resolve distribution conflicts
+    /// - failed to set distribution
+    pub fn add_motions(&mut self) -> Result<(), QueryPlannerError> {
+        let map = self.relational_id_map();
+
+        let nodes = self.get_relational_nodes_dfs_post()?;
+        for id in &nodes {
+            let rel_op: Relational = match self.get_node(*id)? {
+                Node::Expression(_) => return Err(QueryPlannerError::InvalidNode),
+                Node::Relational(rel) => rel.clone(),
+            };
+            match rel_op {
+                // At the moment our grammar and IR constructor
+                // don't support projection with sub queries.
+                Relational::Projection { output, .. }
+                | Relational::ScanRelation { output, .. }
+                | Relational::ScanSubQuery { output, .. }
+                | Relational::UnionAll { output, .. } => {
+                    self.set_distribution(output, &map)?;
+                }
+                Relational::Motion { .. } => {
+                    // We can apply this transformation only once,
+                    // i.e. to the plan without any motion nodes.
+                    return Err(QueryPlannerError::InvalidNode);
+                }
+                Relational::Selection { output, filter, .. } => {
+                    self.resolve_subquery_conflicts(*id, filter, &map)?;
+                    self.set_distribution(output, &map)?;
+                }
+                Relational::InnerJoin {
+                    output, condition, ..
+                } => {
+                    self.resolve_join_conflicts(*id, condition, &map)?;
+                    self.set_distribution(output, &map)?;
+                }
+            }
+        }
+        Ok(())
+    }
+}