Share tree display traversal and rendering

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

Author: joseph-isaacs

Cargo.lock

@@ -60,7 +60,6 @@ static_assertions = { workspace = true }
 tabled = { workspace = true, optional = true, default-features = false, features = [
     "std",
 ] }
-termtree = { workspace = true }
 tracing = { workspace = true }
 uuid = { workspace = true }
 vortex-array-macros = { workspace = true }

vortex-array/Cargo.toml

@@ -3,7 +3,11 @@
 
 use std::fmt;
 
+pub use vortex_utils::tree::IndentedFormatter;
+use vortex_utils::tree::TreeDisplayContext;
+
 use crate::ArrayRef;
+use crate::arrays::Chunked;
 
 /// Context threaded through tree traversal for percentage calculations etc.
 pub struct TreeContext {
@@ -25,40 +29,19 @@ impl TreeContext {
     pub fn parent_total_size(&self) -> Option<u64> {
         self.ancestor_sizes.last().cloned().flatten()
     }
-
-    pub(crate) fn push(&mut self, size: Option<u64>) {
-        self.ancestor_sizes.push(size);
-    }
-
-    pub(crate) fn pop(&mut self) {
-        self.ancestor_sizes.pop();
-    }
 }
 
-/// Wrapper providing access to a [`fmt::Formatter`] and the current indentation string.
-pub struct IndentedFormatter<'a, 'b> {
-    inner: &'a mut fmt::Formatter<'b>,
-    indent: &'a str,
-}
-
-impl<'a, 'b> IndentedFormatter<'a, 'b> {
-    pub(crate) fn new(f: &'a mut fmt::Formatter<'b>, indent: &'a str) -> Self {
-        Self { inner: f, indent }
-    }
-
-    /// Access the indent string and underlying [`fmt::Formatter`] together.
-    pub fn parts(&mut self) -> (&str, &mut fmt::Formatter<'b>) {
-        (self.indent, self.inner)
+impl TreeDisplayContext<ArrayRef> for TreeContext {
+    fn push_parent(&mut self, parent: &ArrayRef) {
+        self.ancestor_sizes.push(if parent.is::<Chunked>() {
+            None
+        } else {
+            Some(parent.nbytes())
+        });
     }
 
-    /// The current indentation string.
-    pub fn indent(&self) -> &str {
-        self.indent
-    }
-
-    /// Access the underlying [`fmt::Formatter`].
-    pub fn formatter(&mut self) -> &mut fmt::Formatter<'b> {
-        self.inner
+    fn pop_parent(&mut self, _parent: &ArrayRef) {
+        self.ancestor_sizes.pop();
     }
 }
 

vortex-array/src/display/extractor.rs

@@ -3,8 +3,10 @@
 
 use std::fmt;
 
+use vortex_utils::tree::TreeDisplayAdapter;
+use vortex_utils::tree::write_indented_tree;
+
 use crate::ArrayRef;
-use crate::arrays::Chunked;
 use crate::display::extractor::IndentedFormatter;
 use crate::display::extractor::TreeContext;
 use crate::display::extractor::TreeExtractor;
@@ -74,54 +76,57 @@ impl TreeDisplay {
         self.extractors.push(extractor);
         self
     }
+}
+
+impl TreeDisplayAdapter for TreeDisplay {
+    type Context = TreeContext;
+    type Node = ArrayRef;
 
-    /// Recursively write a node and all its descendants directly to the formatter.
     fn write_node(
         &self,
-        name: &str,
         array: &ArrayRef,
-        ctx: &mut TreeContext,
-        indent: &str,
+        ctx: &TreeContext,
         f: &mut fmt::Formatter<'_>,
     ) -> fmt::Result {
-        // Header line: "{indent}{name}:{annotations...}\n"
-        write!(f, "{indent}{name}:")?;
         for extractor in &self.extractors {
             extractor.write_header(array, ctx, f)?;
         }
-        writeln!(f)?;
+        Ok(())
+    }
 
-        // Detail lines
-        let child_indent = format!("{indent}  ");
-        {
-            let mut indented = IndentedFormatter::new(f, &child_indent);
-            for extractor in &self.extractors {
-                extractor.write_details(array, ctx, &mut indented)?;
-            }
+    fn write_details(
+        &self,
+        array: &ArrayRef,
+        ctx: &TreeContext,
+        f: &mut IndentedFormatter<'_, '_>,
+    ) -> fmt::Result {
+        for extractor in &self.extractors {
+            extractor.write_details(array, ctx, f)?;
         }
+        Ok(())
+    }
 
-        // Push context for children: chunked arrays reset the percentage root
-        let child_size = if array.is::<Chunked>() {
-            None
-        } else {
-            Some(array.nbytes())
-        };
-        ctx.push(child_size);
-
-        // Recurse into children
-        for (child_name, child) in array.children_names().into_iter().zip(array.children()) {
-            self.write_node(&child_name, &child, ctx, &child_indent, f)?;
+    fn visit_children(
+        &self,
+        array: &ArrayRef,
+        visit: &mut dyn FnMut(&str, &ArrayRef, bool) -> fmt::Result,
+    ) -> fmt::Result {
+        let mut children = array
+            .children_names()
+            .into_iter()
+            .zip(array.children())
+            .peekable();
+        while let Some((child_name, child)) = children.next() {
+            let is_last = children.peek().is_none();
+            visit(&child_name, &child, is_last)?;
         }
-
-        ctx.pop();
-
         Ok(())
     }
 }
 
 impl fmt::Display for TreeDisplay {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let mut ctx = TreeContext::new();
-        self.write_node("root", &self.array, &mut ctx, "", f)
+        write_indented_tree(self, "root", &self.array, &mut ctx, f)
     }
 }

vortex-array/src/display/tree_display.rs

@@ -5,6 +5,9 @@ use std::fmt;
 use std::fmt::Display;
 use std::fmt::Formatter;
 
+use vortex_utils::tree::TreeDisplayAdapter;
+use vortex_utils::tree::write_branch_tree;
+
 use crate::expr::BoundExpression;
 use crate::expr::BoundKind;
 use crate::expr::Expression;
@@ -90,39 +93,38 @@ impl DisplayTreeNode for BoundExpression {
     }
 }
 
-struct NodeDisplay<'a, T>(&'a T);
+pub struct DisplayTreeExpr<'a, T: ?Sized = Expression>(pub &'a T);
 
-impl<T: DisplayTreeNode> Display for NodeDisplay<'_, T> {
-    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
-        self.0.fmt_tree_node(f)
+impl<T: DisplayTreeNode> TreeDisplayAdapter for DisplayTreeExpr<'_, T> {
+    type Context = ();
+    type Node = T;
+
+    fn write_node(
+        &self,
+        node: &Self::Node,
+        _context: &Self::Context,
+        formatter: &mut Formatter<'_>,
+    ) -> fmt::Result {
+        node.fmt_tree_node(formatter)
     }
-}
 
-pub struct DisplayTreeExpr<'a, T: ?Sized = Expression>(pub &'a T);
+    fn visit_children(
+        &self,
+        node: &Self::Node,
+        visit: &mut dyn FnMut(&str, &Self::Node, bool) -> fmt::Result,
+    ) -> fmt::Result {
+        let children = node.tree_children();
+        for (index, child) in children.iter().enumerate() {
+            let child_name = node.tree_child_name(index);
+            visit(child_name.as_ref(), child, index + 1 == children.len())?;
+        }
+        Ok(())
+    }
+}
 
 impl<T: DisplayTreeNode> Display for DisplayTreeExpr<'_, T> {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
-        pub use termtree::Tree;
-        fn make_tree<T: DisplayTreeNode>(expr: &T) -> Tree<String> {
-            let child_trees = expr
-                .tree_children()
-                .iter()
-                .enumerate()
-                .map(|(index, child)| {
-                    let child_tree = make_tree(child);
-                    Tree::new(format!(
-                        "{}: {}",
-                        expr.tree_child_name(index),
-                        child_tree.root
-                    ))
-                    .with_leaves(child_tree.leaves)
-                })
-                .collect::<Vec<_>>();
-
-            Tree::new(NodeDisplay(expr).to_string()).with_leaves(child_trees)
-        }
-
-        write!(f, "{}", make_tree(self.0))
+        write_branch_tree(self, self.0, &mut (), f)
     }
 }
 

vortex-array/src/expr/display.rs

@@ -11,3 +11,4 @@ pub mod debug_with;
 pub mod dyn_traits;
 pub mod iter;
 pub mod parallelism;
+pub mod tree;