fix(core): improve dead-loop avoidance in OptimizeInputs (v2) (#206)

Currently, OptimizeInputsTask avoids a dead-loop with a sequential
traversal of a potentially long list of tasks, with a dynamic cast for
each of them (see #29). This alternative uses a hashset of previously
visited expressions for the same purpose, similarly to explored_group.

This is the second attempt, after #196.

Author: jopereira

optd-core/src/cascades/optimizer.rs

@@ -43,6 +43,7 @@ pub struct CascadesOptimizer<T: RelNodeTyp> {
     memo: Memo<T>,
     pub(super) tasks: VecDeque<Box<dyn Task<T>>>,
     explored_group: HashSet<GroupId>,
+    explored_expr: HashSet<ExprId>,
     fired_rules: HashMap<ExprId, HashSet<RuleId>>,
     rules: Arc<[Arc<RuleWrapper<T, Self>>]>,
     disabled_rules: HashSet<usize>,
@@ -105,6 +106,7 @@ impl<T: RelNodeTyp> CascadesOptimizer<T> {
             memo,
             tasks,
             explored_group: HashSet::new(),
+            explored_expr: HashSet::new(),
             fired_rules: HashMap::new(),
             rules: rules.into(),
             cost: cost.into(),
@@ -202,11 +204,13 @@ impl<T: RelNodeTyp> CascadesOptimizer<T> {
         self.memo = Memo::new(self.property_builders.clone());
         self.fired_rules.clear();
         self.explored_group.clear();
+        self.explored_expr.clear();
     }
 
     /// Clear the winner so that the optimizer can continue to explore the group.
     pub fn step_clear_winner(&mut self) {
         self.memo.clear_winner();
+        self.explored_expr.clear();
     }
 
     /// Optimize a `RelNode`.
@@ -347,6 +351,18 @@ impl<T: RelNodeTyp> CascadesOptimizer<T> {
         self.explored_group.insert(group_id);
     }
 
+    pub(super) fn is_expr_explored(&self, expr_id: ExprId) -> bool {
+        self.explored_expr.contains(&expr_id)
+    }
+
+    pub(super) fn mark_expr_explored(&mut self, expr_id: ExprId) {
+        self.explored_expr.insert(expr_id);
+    }
+
+    pub(super) fn unmark_expr_explored(&mut self, expr_id: ExprId) {
+        self.explored_expr.remove(&expr_id);
+    }
+
     pub(super) fn is_rule_fired(&self, group_expr_id: ExprId, rule_id: RuleId) -> bool {
         self.fired_rules
             .get(&group_expr_id)

optd-core/src/cascades/tasks.rs

@@ -18,6 +18,5 @@ pub use optimize_inputs::OptimizeInputsTask;
 
 pub trait Task<T: RelNodeTyp>: 'static + Send + Sync {
     fn execute(&self, optimizer: &mut CascadesOptimizer<T>) -> Result<Vec<Box<dyn Task<T>>>>;
-    fn as_any(&self) -> &dyn std::any::Any;
     fn describe(&self) -> String;
 }

optd-core/src/cascades/tasks/apply_rule.rs

@@ -172,10 +172,6 @@ fn match_and_pick<T: RelNodeTyp>(
 }
 
 impl<T: RelNodeTyp> Task<T> for ApplyRuleTask {
-    fn as_any(&self) -> &dyn std::any::Any {
-        self
-    }
-
     fn execute(&self, optimizer: &mut CascadesOptimizer<T>) -> Result<Vec<Box<dyn Task<T>>>> {
         if optimizer.is_rule_fired(self.expr_id, self.rule_id) {
             return Ok(vec![]);
@@ -215,6 +211,7 @@ impl<T: RelNodeTyp> Task<T> for ApplyRuleTask {
                             .push(Box::new(OptimizeInputsTask::new(expr_id, true))
                                 as Box<dyn Task<T>>);
                     }
+                    optimizer.unmark_expr_explored(expr_id);
                     trace!(event = "apply_rule", expr_id = %self.expr_id, rule_id = %self.rule_id, new_expr_id = %expr_id);
                 } else {
                     trace!(event = "apply_rule", expr_id = %self.expr_id, rule_id = %self.rule_id, "triggered group merge");

optd-core/src/cascades/tasks/explore_group.rs

@@ -22,10 +22,6 @@ impl ExploreGroupTask {
 }
 
 impl<T: RelNodeTyp> Task<T> for ExploreGroupTask {
-    fn as_any(&self) -> &dyn std::any::Any {
-        self
-    }
-
     fn execute(&self, optimizer: &mut CascadesOptimizer<T>) -> Result<Vec<Box<dyn Task<T>>>> {
         trace!(event = "task_begin", task = "explore_group", group_id = %self.group_id);
         let mut tasks = vec![];

optd-core/src/cascades/tasks/optimize_expression.rs

@@ -36,10 +36,6 @@ fn top_matches<T: RelNodeTyp>(
 }
 
 impl<T: RelNodeTyp> Task<T> for OptimizeExpressionTask {
-    fn as_any(&self) -> &dyn std::any::Any {
-        self
-    }
-
     fn execute(&self, optimizer: &mut CascadesOptimizer<T>) -> Result<Vec<Box<dyn Task<T>>>> {
         let expr = optimizer.get_expr_memoed(self.expr_id);
         trace!(event = "task_begin", task = "optimize_expr", expr_id = %self.expr_id, expr = %expr);

optd-core/src/cascades/tasks/optimize_group.rs

@@ -23,10 +23,6 @@ impl OptimizeGroupTask {
 }
 
 impl<T: RelNodeTyp> Task<T> for OptimizeGroupTask {
-    fn as_any(&self) -> &dyn std::any::Any {
-        self
-    }
-
     fn execute(&self, optimizer: &mut CascadesOptimizer<T>) -> Result<Vec<Box<dyn Task<T>>>> {
         trace!(event = "task_begin", task = "optimize_group", group_id = %self.group_id);
         let group_info = optimizer.get_group_info(self.group_id);

optd-core/src/cascades/tasks/optimize_inputs.rs

@@ -117,22 +117,15 @@ impl OptimizeInputsTask {
 }
 
 impl<T: RelNodeTyp> Task<T> for OptimizeInputsTask {
-    fn as_any(&self) -> &dyn std::any::Any {
-        self
-    }
-
     fn execute(&self, optimizer: &mut CascadesOptimizer<T>) -> Result<Vec<Box<dyn Task<T>>>> {
-        if optimizer.tasks.iter().any(|t| {
-            if let Some(task) = t.as_any().downcast_ref::<Self>() {
+        if self.continue_from.is_none() {
+            if optimizer.is_expr_explored(self.expr_id) {
                 // skip optimize_inputs to avoid dead-loop: consider join commute being fired twice that produces
                 // two projections, therefore having groups like projection1 -> projection2 -> join = projection1.
-                task.expr_id == self.expr_id
-            } else {
-                false
+                trace!(event = "task_skip", task = "optimize_inputs", expr_id = %self.expr_id);
+                return Ok(vec![]);
             }
-        }) {
-            trace!(event = "task_skip", task = "optimize_inputs", expr_id = %self.expr_id);
-            return Ok(vec![]);
+            optimizer.mark_expr_explored(self.expr_id);
         }
         trace!(event = "task_begin", task = "optimize_inputs", expr_id = %self.expr_id, continue_from = ?self.continue_from);
         let expr = optimizer.get_expr_memoed(self.expr_id);