From 2aac95702d257c6f1b665ddd544dddc355063930 Mon Sep 17 00:00:00 2001
From: Desmond Cheong <desmondcheongzx@gmail.com>
Date: Mon, 6 Jan 2025 14:14:43 -0800
Subject: [PATCH] chore: Clean up join order iteration (#3638)

Missed a comment to
https://github.com/Eventual-Inc/Daft/pull/3616#discussion_r1893614089:
```
In [src/daft-logical-plan/src/optimization/rules/reorder_joins/join_graph.rs](https://github.com/Eventual-Inc/Daft/pull/3616#discussion_r1893614089):

>  };

-#[derive(Debug)]
-struct JoinNode {
+// TODO(desmond): In the future these trees should keep track of current cost estimates.
+#[derive(Clone, Debug)]
+pub(super) enum JoinOrderTree {
+    Relation(usize),                                          // (id).
+    Join(Box<JoinOrderTree>, Box<JoinOrderTree>, Vec<usize>), // (subtree, subtree, nodes involved).
I dont think you need to keep an explicit stack. I think you should be able to something like

std::iter::chain(left.into_iter(), right.into_iter())
```

This PR removes the stack and simply chains the iterators.
---
 .../rules/reorder_joins/join_graph.rs         | 26 +++----------------
 1 file changed, 4 insertions(+), 22 deletions(-)

diff --git a/src/daft-logical-plan/src/optimization/rules/reorder_joins/join_graph.rs b/src/daft-logical-plan/src/optimization/rules/reorder_joins/join_graph.rs
index e7cd907657..cd2829664a 100644
--- a/src/daft-logical-plan/src/optimization/rules/reorder_joins/join_graph.rs
+++ b/src/daft-logical-plan/src/optimization/rules/reorder_joins/join_graph.rs
@@ -38,29 +38,11 @@ impl JoinOrderTree {
         }
     }
 
-    pub(super) fn iter(&self) -> JoinOrderTreeIterator {
-        JoinOrderTreeIterator { stack: vec![self] }
-    }
-}
-
-pub(super) struct JoinOrderTreeIterator<'a> {
-    stack: Vec<&'a JoinOrderTree>,
-}
-
-impl<'a> Iterator for JoinOrderTreeIterator<'a> {
-    type Item = usize;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        while let Some(node) = self.stack.pop() {
-            match node {
-                JoinOrderTree::Relation(id) => return Some(*id),
-                JoinOrderTree::Join(left, right) => {
-                    self.stack.push(left);
-                    self.stack.push(right);
-                }
-            }
+    fn iter(&self) -> Box<dyn Iterator<Item = usize> + '_> {
+        match self {
+            JoinOrderTree::Relation(id) => Box::new(std::iter::once(*id)),
+            JoinOrderTree::Join(left, right) => Box::new(left.iter().chain(right.iter())),
         }
-        None
     }
 }