diff --git a/Cargo.lock b/Cargo.lock index cba7a70ede5..7e6717a4209 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9620,7 +9620,6 @@ dependencies = [ "smallvec", "static_assertions", "tabled", - "termtree", "test-with", "tracing", "uuid", diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index d00b811a387..746ce079a6a 100644 --- a/vortex-array/Cargo.toml +++ b/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 } diff --git a/vortex-array/src/display/extractor.rs b/vortex-array/src/display/extractor.rs index bbfc25a3666..0998872529d 100644 --- a/vortex-array/src/display/extractor.rs +++ b/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 { self.ancestor_sizes.last().cloned().flatten() } - - pub(crate) fn push(&mut self, size: Option) { - 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 for TreeContext { + fn push_parent(&mut self, parent: &ArrayRef) { + self.ancestor_sizes.push(if parent.is::() { + 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(); } } diff --git a/vortex-array/src/display/extractors/buffer.rs b/vortex-array/src/display/extractors/buffer.rs index 17861a654f7..2fb1c716f3d 100644 --- a/vortex-array/src/display/extractors/buffer.rs +++ b/vortex-array/src/display/extractors/buffer.rs @@ -17,7 +17,7 @@ pub struct BufferExtractor { pub show_percent: bool, } -impl TreeExtractor for BufferExtractor { +impl TreeExtractor for BufferExtractor { fn write_details( &self, array: &ArrayRef, diff --git a/vortex-array/src/display/extractors/encoding_summary.rs b/vortex-array/src/display/extractors/encoding_summary.rs index bc678643ecb..43c924190df 100644 --- a/vortex-array/src/display/extractors/encoding_summary.rs +++ b/vortex-array/src/display/extractors/encoding_summary.rs @@ -23,7 +23,7 @@ impl EncodingSummaryExtractor { } } -impl TreeExtractor for EncodingSummaryExtractor { +impl TreeExtractor for EncodingSummaryExtractor { fn write_header( &self, array: &ArrayRef, diff --git a/vortex-array/src/display/extractors/metadata.rs b/vortex-array/src/display/extractors/metadata.rs index a355bafe63c..d68e9c7a1d8 100644 --- a/vortex-array/src/display/extractors/metadata.rs +++ b/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 for MetadataExtractor { fn write_details( &self, array: &ArrayRef, diff --git a/vortex-array/src/display/extractors/nbytes.rs b/vortex-array/src/display/extractors/nbytes.rs index 37e100279fc..27d850444db 100644 --- a/vortex-array/src/display/extractors/nbytes.rs +++ b/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 for NbytesExtractor { fn write_header( &self, array: &ArrayRef, diff --git a/vortex-array/src/display/extractors/stats.rs b/vortex-array/src/display/extractors/stats.rs index e39dafa2340..7459daccc5c 100644 --- a/vortex-array/src/display/extractors/stats.rs +++ b/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 for StatsExtractor { fn write_header( &self, array: &ArrayRef, diff --git a/vortex-array/src/display/mod.rs b/vortex-array/src/display/mod.rs index 1c12afcf8d2..780140af2ac 100644 --- a/vortex-array/src/display/mod.rs +++ b/vortex-array/src/display/mod.rs @@ -556,7 +556,7 @@ impl ArrayRef { metadata, stats, } => { - let extractors: [(bool, Box); 5] = [ + let extractors: [(bool, Box>); 5] = [ (true, Box::new(EncodingSummaryExtractor)), (*stats, Box::new(NbytesExtractor)), (*stats, Box::new(StatsExtractor)), diff --git a/vortex-array/src/display/tree_display.rs b/vortex-array/src/display/tree_display.rs index 5d0753ba354..76217da0f8d 100644 --- a/vortex-array/src/display/tree_display.rs +++ b/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>, + extractors: Vec>>, } impl TreeDisplay { @@ -64,57 +66,60 @@ impl TreeDisplay { } /// Add an extractor to the display pipeline. - pub fn with(mut self, extractor: E) -> Self { + pub fn with + '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) -> Self { + pub fn with_boxed(mut self, extractor: Box>) -> 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::() { - 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(()) } } @@ -122,6 +127,6 @@ impl TreeDisplay { 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) } } diff --git a/vortex-array/src/expr/display.rs b/vortex-array/src/expr/display.rs index d542f7d35b3..ee292e4a82e 100644 --- a/vortex-array/src/expr/display.rs +++ b/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 Display for NodeDisplay<'_, T> { - fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - self.0.fmt_tree_node(f) +impl 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 Display for DisplayTreeExpr<'_, T> { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - pub use termtree::Tree; - fn make_tree(expr: &T) -> Tree { - 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::>(); - - Tree::new(NodeDisplay(expr).to_string()).with_leaves(child_trees) - } - - write!(f, "{}", make_tree(self.0)) + write_branch_tree(self, self.0, &mut (), f) } } diff --git a/vortex-btrblocks/tests/golden.rs b/vortex-btrblocks/tests/golden.rs index c60752b064a..4f3f4a1d7e2 100644 --- a/vortex-btrblocks/tests/golden.rs +++ b/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 for ExactNbytesExtractor { fn write_header( &self, array: &ArrayRef, diff --git a/vortex-utils/src/lib.rs b/vortex-utils/src/lib.rs index 31decd406ff..993c1c82337 100644 --- a/vortex-utils/src/lib.rs +++ b/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; diff --git a/vortex-utils/src/tree.rs b/vortex-utils/src/tree.rs new file mode 100644 index 00000000000..7382b2cb66b --- /dev/null +++ b/vortex-utils/src/tree.rs @@ -0,0 +1,335 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Shared traversal and rendering utilities for named trees. + +use std::fmt; + +/// Traversal state updated as a tree renderer enters and leaves a node's children. +/// +/// The context visible while a node is rendered describes its ancestors. The renderer calls +/// [`Self::push_parent`] before visiting the node's children and [`Self::pop_parent`] afterwards. +pub trait TreeDisplayContext { + /// Record `parent` while its children are visited. + fn push_parent(&mut self, parent: &N) { + _ = parent; + } + + /// Remove `parent` after all of its children have been visited. + fn pop_parent(&mut self, parent: &N) { + _ = parent; + } +} + +impl TreeDisplayContext for () {} + +/// Tree traversal context that records only the current depth. +#[derive(Debug, Default)] +pub struct DepthContext { + depth: usize, +} + +impl DepthContext { + /// Return the current node's depth, where the root has depth zero. + pub fn depth(&self) -> usize { + self.depth + } +} + +impl TreeDisplayContext for DepthContext { + fn push_parent(&mut self, _parent: &N) { + self.depth += 1; + } + + fn pop_parent(&mut self, _parent: &N) { + debug_assert!(self.depth > 0, "tree depth push/pop mismatch"); + self.depth -= 1; + } +} + +/// Access to a formatter together with the indentation for detail lines. +pub struct IndentedFormatter<'a, 'b> { + inner: &'a mut fmt::Formatter<'b>, + indent: &'a str, +} + +impl<'a, 'b> IndentedFormatter<'a, 'b> { + fn new(inner: &'a mut fmt::Formatter<'b>, indent: &'a str) -> Self { + Self { inner, indent } + } + + /// Return the indentation string and underlying formatter together. + pub fn parts(&mut self) -> (&str, &mut fmt::Formatter<'b>) { + (self.indent, self.inner) + } + + /// Return the current indentation string. + pub fn indent(&self) -> &str { + self.indent + } + + /// Return the underlying formatter. + pub fn formatter(&mut self) -> &mut fmt::Formatter<'b> { + self.inner + } +} + +/// Contributes one composable dimension of information to tree nodes. +pub trait TreeDisplayExtractor>: Send + Sync { + /// Write space-prefixed annotations on the node's header line. + fn write_header( + &self, + node: &N, + context: &C, + formatter: &mut fmt::Formatter<'_>, + ) -> fmt::Result { + _ = (node, context, formatter); + Ok(()) + } + + /// Write detail lines beneath the node's header. + fn write_details( + &self, + node: &N, + context: &C, + formatter: &mut IndentedFormatter<'_, '_>, + ) -> fmt::Result { + _ = (node, context, formatter); + Ok(()) + } +} + +/// Adapts a domain-specific node and traversal context to the shared tree renderers. +pub trait TreeDisplayAdapter { + /// Node type traversed by this adapter. + type Node: ?Sized; + + /// State made available while each node is rendered. + type Context: TreeDisplayContext; + + /// Write the node's display content. + /// + /// The indented renderer writes `name:` first, so implementations normally write + /// space-prefixed annotations. The branch renderer uses this as the complete node label. + fn write_node( + &self, + node: &Self::Node, + context: &Self::Context, + formatter: &mut fmt::Formatter<'_>, + ) -> fmt::Result; + + /// Write detail lines beneath an indented node header. + fn write_details( + &self, + node: &Self::Node, + context: &Self::Context, + formatter: &mut IndentedFormatter<'_, '_>, + ) -> fmt::Result { + _ = (node, context, formatter); + Ok(()) + } + + /// Visit each named child in display order. + /// + /// The final argument to `visit` must be `true` only for the last child. Renderers call the + /// visitor synchronously, so adapters may pass either stored or temporarily owned nodes. + fn visit_children( + &self, + node: &Self::Node, + visit: &mut dyn FnMut(&str, &Self::Node, bool) -> fmt::Result, + ) -> fmt::Result; +} + +/// Render a named tree using two-space indentation. +/// +/// Each node is written as `name:` followed by [`TreeDisplayAdapter::write_node`]. Detail lines +/// and child nodes are indented beneath it. +pub fn write_indented_tree( + adapter: &A, + root_name: &str, + root: &A::Node, + context: &mut A::Context, + formatter: &mut fmt::Formatter<'_>, +) -> fmt::Result { + write_indented_node(adapter, root_name, root, context, "", formatter) +} + +fn write_indented_node( + adapter: &A, + name: &str, + node: &A::Node, + context: &mut A::Context, + indent: &str, + formatter: &mut fmt::Formatter<'_>, +) -> fmt::Result { + write!(formatter, "{indent}{name}:")?; + adapter.write_node(node, context, formatter)?; + writeln!(formatter)?; + + let child_indent = format!("{indent} "); + { + let mut indented = IndentedFormatter::new(formatter, &child_indent); + adapter.write_details(node, context, &mut indented)?; + } + + context.push_parent(node); + let result = adapter.visit_children(node, &mut |child_name, child, _is_last| { + write_indented_node( + adapter, + child_name, + child, + context, + &child_indent, + formatter, + ) + }); + context.pop_parent(node); + result +} + +/// Render a tree using Unicode branch connectors. +/// +/// The root contains only [`TreeDisplayAdapter::write_node`]. Descendants are prefixed with their +/// child name and connectors such as `├──` and `└──`. Detail lines are not rendered in this style. +pub fn write_branch_tree( + adapter: &A, + root: &A::Node, + context: &mut A::Context, + formatter: &mut fmt::Formatter<'_>, +) -> fmt::Result { + write_branch_node(adapter, root, context, "", formatter) +} + +fn write_branch_node( + adapter: &A, + node: &A::Node, + context: &mut A::Context, + prefix: &str, + formatter: &mut fmt::Formatter<'_>, +) -> fmt::Result { + adapter.write_node(node, context, formatter)?; + + context.push_parent(node); + let result = adapter.visit_children(node, &mut |child_name, child, is_last| { + writeln!(formatter)?; + let connector = if is_last { "└── " } else { "├── " }; + write!(formatter, "{prefix}{connector}{child_name}: ")?; + let child_prefix = format!("{prefix}{}", if is_last { " " } else { "│ " }); + write_branch_node(adapter, child, context, &child_prefix, formatter) + }); + context.pop_parent(node); + result +} + +#[cfg(test)] +mod tests { + use std::fmt; + + use super::DepthContext; + use super::TreeDisplayAdapter; + use super::write_branch_tree; + use super::write_indented_tree; + + struct TestNode { + label: &'static str, + children: Vec<(&'static str, TestNode)>, + } + + struct TestAdapter; + + impl TreeDisplayAdapter for TestAdapter { + type Context = DepthContext; + type Node = TestNode; + + fn write_node( + &self, + node: &Self::Node, + context: &Self::Context, + formatter: &mut fmt::Formatter<'_>, + ) -> fmt::Result { + write!(formatter, "{}@{}", node.label, context.depth()) + } + + fn visit_children( + &self, + node: &Self::Node, + visit: &mut dyn FnMut(&str, &Self::Node, bool) -> fmt::Result, + ) -> fmt::Result { + for (index, (name, child)) in node.children.iter().enumerate() { + visit(name, child, index + 1 == node.children.len())?; + } + Ok(()) + } + } + + struct IndentedDisplay<'a>(&'a TestNode); + + impl fmt::Display for IndentedDisplay<'_> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + write_indented_tree( + &TestAdapter, + "root", + self.0, + &mut DepthContext::default(), + formatter, + ) + } + } + + struct BranchDisplay<'a>(&'a TestNode); + + impl fmt::Display for BranchDisplay<'_> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + write_branch_tree( + &TestAdapter, + self.0, + &mut DepthContext::default(), + formatter, + ) + } + } + + fn tree() -> TestNode { + TestNode { + label: "parent", + children: vec![ + ( + "left", + TestNode { + label: "branch", + children: vec![( + "leaf", + TestNode { + label: "first", + children: Vec::new(), + }, + )], + }, + ), + ( + "right", + TestNode { + label: "second", + children: Vec::new(), + }, + ), + ], + } + } + + #[test] + fn renders_indented_tree() { + assert_eq!( + IndentedDisplay(&tree()).to_string(), + "root:parent@0\n left:branch@1\n leaf:first@2\n right:second@1\n" + ); + } + + #[test] + fn renders_branch_tree() { + assert_eq!( + BranchDisplay(&tree()).to_string(), + "parent@0\n├── left: branch@1\n│ └── leaf: first@2\n└── right: second@1" + ); + } +}