in work

Author: waralexrom

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/aggregate_multiplied_subquery.rs

@@ -99,6 +99,14 @@ impl NodeInputs for AggregateMultipliedSubqueryInput {
                 .chain(self.dimension_subqueries.iter()),
         )
     }
+    
+    fn iter_mut(&mut self) -> Box<dyn Iterator<Item = &mut PlanNode> + '_> {
+        Box::new(
+            std::iter::once(&mut self.keys_subquery)
+                .chain(std::iter::once(&mut self.source))
+                .chain(self.dimension_subqueries.iter_mut()),
+        )
+    }
 }
 
 impl PrettyPrint for AggregateMultipliedSubquery {

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/join.rs

@@ -117,6 +117,14 @@ impl NodeInputs for LogicalJoinInput {
                 .chain(self.dimension_subqueries.iter()),
         )
     }
+    
+    fn iter_mut(&mut self) -> Box<dyn Iterator<Item = &mut PlanNode> + '_> {
+        Box::new(
+            std::iter::once(&mut self.root)
+                .chain(self.joins.iter_mut())
+                .chain(self.dimension_subqueries.iter_mut()),
+        )
+    }
 }
 
 impl PrettyPrint for LogicalJoin {

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/keys_subquery.rs

@@ -60,6 +60,10 @@ impl NodeInputs for KeysSubQueryInputs {
     fn iter(&self) -> Box<dyn Iterator<Item = &PlanNode> + '_> {
         Box::new(std::iter::once(&self.pk_cube).chain(std::iter::once(&self.source)))
     }
+    
+    fn iter_mut(&mut self) -> Box<dyn Iterator<Item = &mut PlanNode> + '_> {
+        Box::new(std::iter::once(&mut self.pk_cube).chain(std::iter::once(&mut self.source)))
+    }
 }
 
 impl PrettyPrint for KeysSubQuery {

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/logical_node.rs

@@ -4,6 +4,7 @@ use std::rc::Rc;
 
 pub trait NodeInputs {
     fn iter(&self) -> Box<dyn Iterator<Item = &PlanNode> + '_>;
+    fn iter_mut(&mut self) -> Box<dyn Iterator<Item = &mut PlanNode> + '_>;
 }
 
 pub struct EmptyNodeInput {}
@@ -17,6 +18,10 @@ impl NodeInputs for EmptyNodeInput {
     fn iter(&self) -> Box<dyn Iterator<Item = &PlanNode> + '_> {
         Box::new(std::iter::empty())
     }
+    
+    fn iter_mut(&mut self) -> Box<dyn Iterator<Item = &mut PlanNode> + '_> {
+        Box::new(std::iter::empty())
+    }
 }
 
 pub struct SingleNodeInput {
@@ -37,6 +42,10 @@ impl NodeInputs for SingleNodeInput {
     fn iter(&self) -> Box<dyn Iterator<Item = &PlanNode> + '_> {
         Box::new(std::iter::once(&self.item))
     }
+    
+    fn iter_mut(&mut self) -> Box<dyn Iterator<Item = &mut PlanNode> + '_> {
+        Box::new(std::iter::once(&mut self.item))
+    }
 }
 
 pub struct OptionNodeInput {
@@ -61,6 +70,14 @@ impl NodeInputs for OptionNodeInput {
             Box::new(std::iter::empty())
         }
     }
+    
+    fn iter_mut(&mut self) -> Box<dyn Iterator<Item = &mut PlanNode> + '_> {
+        if let Some(item) = &mut self.item {
+            Box::new(std::iter::once(item))
+        } else {
+            Box::new(std::iter::empty())
+        }
+    }
 }
 
 pub struct VecNodeInput {
@@ -81,6 +98,10 @@ impl NodeInputs for VecNodeInput {
     fn iter(&self) -> Box<dyn Iterator<Item = &PlanNode> + '_> {
         Box::new(self.items.iter())
     }
+    
+    fn iter_mut(&mut self) -> Box<dyn Iterator<Item = &mut PlanNode> + '_> {
+        Box::new(self.items.iter_mut())
+    }
 }
 
 pub trait LogicalNode {

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/query.rs

@@ -115,6 +115,13 @@ impl NodeInputs for QueryInput {
                 .chain(self.multistage_members.iter())
         )
     }
+    
+    fn iter_mut(&mut self) -> Box<dyn Iterator<Item = &mut PlanNode> + '_> {
+        Box::new(
+            std::iter::once(&mut self.source)
+                .chain(self.multistage_members.iter_mut())
+        )
+    }
 }
 
 impl PrettyPrint for Query {

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/resolve_multiplied_measures.rs

@@ -100,6 +100,14 @@ impl NodeInputs for ResolveMultipliedMeasuresInput {
                 .chain(self.aggregate_multiplied_subqueries.iter()),
         )
     }
+    
+    fn iter_mut(&mut self) -> Box<dyn Iterator<Item = &mut PlanNode> + '_> {
+        Box::new(
+            self.regular_measure_subqueries
+                .iter_mut()
+                .chain(self.aggregate_multiplied_subqueries.iter_mut()),
+        )
+    }
 }
 
 impl PrettyPrint for ResolveMultipliedMeasures {