Share tree display traversal and rendering

vortex-array/Cargo.toml

@@ -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/src/display/extractor.rs

@@ -1,9 +1,12 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use std::fmt;
+pub use vortex_utils::tree::IndentedFormatter;
+use vortex_utils::tree::TreeDisplayContext;
+pub use vortex_utils::tree::TreeDisplayExtractor as TreeExtractor;
 
 use crate::ArrayRef;
+use crate::arrays::Chunked;
 
 /// Context threaded through tree traversal for percentage calculations etc.
 pub struct TreeContext {
@@ -25,73 +28,18 @@ 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
-    }
-}
-
-/// Trait for contributing display information to tree nodes.
-///
-/// Each extractor represents one "dimension" of display (e.g., nbytes, stats, metadata, buffers).
-/// Extractors are composable: you can combine any number of them via [`TreeDisplay::with`].
-///
-/// [`TreeDisplay::with`]: super::TreeDisplay::with
-pub trait TreeExtractor: Send + Sync {
-    /// Write header annotations (space-prefixed) to the formatter.
-    fn write_header(
-        &self,
-        array: &ArrayRef,
-        ctx: &TreeContext,
-        f: &mut fmt::Formatter<'_>,
-    ) -> fmt::Result {
-        let _ = (array, ctx, f);
-        Ok(())
-    }
-
-    /// Write detail lines below the header.
-    ///
-    /// Content written through `f` is automatically indented. Use
-    /// [`f.formatter()`](IndentedFormatter::formatter) to access the underlying
-    /// [`fmt::Formatter`] for formatting flags.
-    fn write_details(
-        &self,
-        array: &ArrayRef,
-        ctx: &TreeContext,
-        f: &mut IndentedFormatter<'_, '_>,
-    ) -> fmt::Result {
-        let _ = (array, ctx, f);
-        Ok(())
+    fn pop_parent(&mut self, _parent: &ArrayRef) {
+        self.ancestor_sizes.pop();
     }
 }

vortex-array/src/display/extractors/buffer.rs

@@ -17,7 +17,7 @@ pub struct BufferExtractor {
     pub show_percent: bool,
 }
 
-impl TreeExtractor for BufferExtractor {
+impl TreeExtractor<ArrayRef, TreeContext> for BufferExtractor {
     fn write_details(
         &self,
         array: &ArrayRef,

vortex-array/src/display/extractors/encoding_summary.rs

@@ -23,7 +23,7 @@ impl EncodingSummaryExtractor {
     }
 }
 
-impl TreeExtractor for EncodingSummaryExtractor {
+impl TreeExtractor<ArrayRef, TreeContext> for EncodingSummaryExtractor {
     fn write_header(
         &self,
         array: &ArrayRef,

vortex-array/src/display/extractors/metadata.rs

@@ -11,7 +11,7 @@ use crate::display::extractor::TreeExtractor;
 /// Extractor that adds a `metadata: ...` detail line.
 pub struct MetadataExtractor;
 
-impl TreeExtractor for MetadataExtractor {
+impl TreeExtractor<ArrayRef, TreeContext> for MetadataExtractor {
     fn write_details(
         &self,
         array: &ArrayRef,

vortex-array/src/display/extractors/nbytes.rs

@@ -13,7 +13,7 @@ use crate::display::extractor::TreeExtractor;
 /// Extractor that adds `nbytes=X (Y%)` to the header line.
 pub struct NbytesExtractor;
 
-impl TreeExtractor for NbytesExtractor {
+impl TreeExtractor<ArrayRef, TreeContext> for NbytesExtractor {
     fn write_header(
         &self,
         array: &ArrayRef,

vortex-array/src/display/extractors/stats.rs

@@ -116,7 +116,7 @@ impl fmt::Display for StatsDisplay<'_> {
 /// Extractor that adds stats annotations (e.g. `[nulls=3, min=5]`) to the header line.
 pub struct StatsExtractor;
 
-impl TreeExtractor for StatsExtractor {
+impl TreeExtractor<ArrayRef, TreeContext> for StatsExtractor {
     fn write_header(
         &self,
         array: &ArrayRef,

vortex-array/src/display/mod.rs

@@ -556,7 +556,7 @@ impl ArrayRef {
                 metadata,
                 stats,
             } => {
-                let extractors: [(bool, Box<dyn TreeExtractor>); 5] = [
+                let extractors: [(bool, Box<dyn TreeExtractor<ArrayRef, TreeContext>>); 5] = [
                     (true, Box::new(EncodingSummaryExtractor)),
                     (*stats, Box::new(NbytesExtractor)),
                     (*stats, Box::new(StatsExtractor)),

vortex-array/src/display/tree_display.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;
@@ -37,7 +39,7 @@ use crate::display::extractors::StatsExtractor;
 /// ```
 pub struct TreeDisplay {
     array: ArrayRef,
-    extractors: Vec<Box<dyn TreeExtractor>>,
+    extractors: Vec<Box<dyn TreeExtractor<ArrayRef, TreeContext>>>,
 }
 
 impl TreeDisplay {
@@ -64,64 +66,67 @@ impl TreeDisplay {
     }
 
     /// Add an extractor to the display pipeline.
-    pub fn with<E: TreeExtractor + 'static>(mut self, extractor: E) -> Self {
+    pub fn with<E: TreeExtractor<ArrayRef, TreeContext> + 'static>(mut self, extractor: E) -> Self {
         self.extractors.push(Box::new(extractor));
         self
     }
 
     /// Add a pre-boxed extractor to the display pipeline.
-    pub fn with_boxed(mut self, extractor: Box<dyn TreeExtractor>) -> Self {
+    pub fn with_boxed(mut self, extractor: Box<dyn TreeExtractor<ArrayRef, TreeContext>>) -> Self {
         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/expr/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-btrblocks/tests/golden.rs

@@ -65,7 +65,7 @@ const N: usize = 16_384;
 /// [`NbytesExtractor`]: vortex_array::display::NbytesExtractor
 struct ExactNbytesExtractor;
 
-impl TreeExtractor for ExactNbytesExtractor {
+impl TreeExtractor<ArrayRef, TreeContext> for ExactNbytesExtractor {
     fn write_header(
         &self,
         array: &ArrayRef,

vortex-utils/src/lib.rs

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