Use shared tree display rendering for plans

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

vortex-layout/src/plan/display.rs

@@ -3,60 +3,14 @@
 
 use std::fmt;
 
+pub use vortex_utils::tree::DepthContext as PlanTreeContext;
+pub use vortex_utils::tree::IndentedFormatter as PlanIndentedFormatter;
+use vortex_utils::tree::TreeDisplayAdapter;
+use vortex_utils::tree::write_indented_tree;
+
 use super::DynPlan;
 use super::Eval;
 
-/// Context threaded through a plan tree traversal.
-pub struct PlanTreeContext {
-    depth: usize,
-}
-
-impl PlanTreeContext {
-    fn new() -> Self {
-        Self { depth: 0 }
-    }
-
-    /// Returns the current node's depth, where the root has depth zero.
-    pub fn depth(&self) -> usize {
-        self.depth
-    }
-
-    fn push(&mut self) {
-        self.depth += 1;
-    }
-
-    fn pop(&mut self) {
-        self.depth -= 1;
-    }
-}
-
-/// Wrapper providing access to a formatter and the current indentation string.
-pub struct PlanIndentedFormatter<'a, 'b> {
-    inner: &'a mut fmt::Formatter<'b>,
-    indent: &'a str,
-}
-
-impl<'a, 'b> PlanIndentedFormatter<'a, 'b> {
-    fn new(inner: &'a mut fmt::Formatter<'b>, indent: &'a str) -> Self {
-        Self { inner, indent }
-    }
-
-    /// Returns the indentation string and underlying formatter together.
-    pub fn parts(&mut self) -> (&str, &mut fmt::Formatter<'b>) {
-        (self.indent, self.inner)
-    }
-
-    /// Returns the current indentation string.
-    pub fn indent(&self) -> &str {
-        self.indent
-    }
-
-    /// Returns the underlying formatter.
-    pub fn formatter(&mut self) -> &mut fmt::Formatter<'b> {
-        self.inner
-    }
-}
-
 /// Contributes one composable dimension of information to plan tree nodes.
 pub trait PlanTreeExtractor: Send + Sync {
     /// Writes space-prefixed annotations on the node's header line.
@@ -163,49 +117,62 @@ impl<'a> PlanTreeDisplay<'a> {
         self.extractors.push(extractor);
         self
     }
+}
+
+impl TreeDisplayAdapter for PlanTreeDisplay<'_> {
+    type Context = PlanTreeContext;
+    type Node = dyn DynPlan;
 
     fn write_node(
         &self,
-        name: &str,
         plan: &dyn DynPlan,
-        context: &mut PlanTreeContext,
-        indent: &str,
+        context: &PlanTreeContext,
         formatter: &mut fmt::Formatter<'_>,
     ) -> fmt::Result {
-        write!(formatter, "{indent}{name}:")?;
         for extractor in &self.extractors {
             extractor.write_header(plan, context, formatter)?;
         }
-        writeln!(formatter)?;
-
-        let child_indent = format!("{indent}  ");
-        {
-            let mut indented = PlanIndentedFormatter::new(formatter, &child_indent);
-            for extractor in &self.extractors {
-                extractor.write_details(plan, context, &mut indented)?;
-            }
+        Ok(())
+    }
+
+    fn write_details(
+        &self,
+        plan: &dyn DynPlan,
+        context: &PlanTreeContext,
+        formatter: &mut PlanIndentedFormatter<'_, '_>,
+    ) -> fmt::Result {
+        for extractor in &self.extractors {
+            extractor.write_details(plan, context, formatter)?;
         }
+        Ok(())
+    }
 
-        context.push();
-        for (index, child) in plan.children().iter().enumerate() {
+    fn visit_children(
+        &self,
+        plan: &dyn DynPlan,
+        visit: &mut dyn FnMut(&str, &dyn DynPlan, bool) -> fmt::Result,
+    ) -> fmt::Result {
+        let children = plan.children();
+        for (index, child) in children.iter().enumerate() {
             let child_name = plan.child_name(index);
-            self.write_node(
+            visit(
                 child_name.as_ref(),
                 child.as_ref(),
-                context,
-                &child_indent,
-                formatter,
+                index + 1 == children.len(),
             )?;
         }
-        context.pop();
-
         Ok(())
     }
 }
 
 impl fmt::Display for PlanTreeDisplay<'_> {
     fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
-        let mut context = PlanTreeContext::new();
-        self.write_node("root", self.plan, &mut context, "", formatter)
+        write_indented_tree(
+            self,
+            "root",
+            self.plan,
+            &mut PlanTreeContext::default(),
+            formatter,
+        )
     }
 }