oak_visualize/tree/
mod.rs

1#![doc = "Tree layout algorithms for AST visualization"]
2
3use crate::{
4    geometry::{Point, Rect, Size},
5    layout::{Edge, Layout},
6};
7use serde::{Deserialize, Serialize};
8use std::{
9    collections::HashMap,
10    sync::atomic::{AtomicUsize, Ordering},
11};
12
13static NODE_ID_COUNTER: AtomicUsize = AtomicUsize::new(0);
14
15fn next_node_id() -> String {
16    format!("node_{}", NODE_ID_COUNTER.fetch_add(1, Ordering::SeqCst))
17}
18
19/// Tree node for visualization.
20#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct TreeNode {
22    /// Unique identifier for the node.
23    pub id: String,
24    /// Display label for the node.
25    pub label: String,
26    /// Type of the node (e.g., "node", "leaf", "function", "struct").
27    pub node_type: String,
28    /// Child nodes of this node.
29    pub children: Vec<TreeNode>,
30    /// Additional metadata associated with the node.
31    pub attributes: HashMap<String, String>,
32    /// Optional explicit size for the node.
33    pub size: Option<Size>,
34}
35
36impl TreeNode {
37    /// Creates a new tree node.
38    pub fn new(id: String, label: String, node_type: String) -> Self {
39        Self { id, label, node_type, children: Vec::new(), attributes: HashMap::new(), size: None }
40    }
41
42    /// Adds a child to the node and returns the modified node.
43    pub fn with_child(mut self, child: TreeNode) -> Self {
44        self.children.push(child);
45        self
46    }
47
48    /// Sets the children of the node and returns the modified node.
49    pub fn with_children(mut self, children: Vec<TreeNode>) -> Self {
50        self.children = children;
51        self
52    }
53
54    /// Adds an attribute to the node and returns the modified node.
55    pub fn with_attribute(mut self, key: String, value: String) -> Self {
56        self.attributes.insert(key, value);
57        self
58    }
59
60    /// Sets the size of the node and returns the modified node.
61    pub fn with_size(mut self, size: Size) -> Self {
62        self.size = Some(size);
63        self
64    }
65
66    /// Returns true if the node has no children.
67    pub fn is_leaf(&self) -> bool {
68        self.children.is_empty()
69    }
70
71    /// Calculates the maximum depth of the tree rooted at this node.
72    pub fn depth(&self) -> usize {
73        if self.children.is_empty() { 1 } else { 1 + self.children.iter().map(|child| child.depth()).max().unwrap_or(0) }
74    }
75
76    /// Calculates the total number of nodes in the tree rooted at this node.
77    pub fn node_count(&self) -> usize {
78        1 + self.children.iter().map(|child| child.node_count()).sum::<usize>()
79    }
80
81    /// Calculates the total number of leaf nodes in the tree rooted at this node.
82    pub fn leaf_count(&self) -> usize {
83        if self.children.is_empty() { 1 } else { self.children.iter().map(|child| child.leaf_count()).sum() }
84    }
85}
86
87impl crate::Visualize for TreeNode {
88    fn visualize(&self) -> crate::Result<String> {
89        TreeLayout::new().visualize(self)
90    }
91}
92
93// Bridge for oak-core types
94impl<'a, L: oak_core::Language> From<&oak_core::GreenTree<'a, L>> for TreeNode {
95    fn from(green: &oak_core::GreenTree<'a, L>) -> Self {
96        match green {
97            oak_core::GreenTree::Node(node) => {
98                let mut tree_node = TreeNode::new(next_node_id(), format!("{:?}", node.kind), "node".to_string());
99                for child in node.children() {
100                    tree_node.children.push(TreeNode::from(child));
101                }
102                tree_node
103            }
104            oak_core::GreenTree::Leaf(leaf) => TreeNode::new(next_node_id(), format!("{:?}", leaf.kind), "leaf".to_string()),
105        }
106    }
107}
108
109impl<'a, L: oak_core::Language> From<&oak_core::RedTree<'a, L>> for TreeNode {
110    fn from(red: &oak_core::RedTree<'a, L>) -> Self {
111        match red {
112            oak_core::RedTree::Node(node) => {
113                let mut tree_node = TreeNode::new(next_node_id(), format!("{:?}", node.green.kind), "node".to_string());
114                for child in node.children() {
115                    tree_node.children.push(TreeNode::from(child));
116                }
117                tree_node
118            }
119            oak_core::RedTree::Leaf(leaf) => TreeNode::new(next_node_id(), format!("{:?}", leaf.kind), "leaf".to_string()),
120        }
121    }
122}
123
124impl<'a, L: oak_core::Language> From<oak_core::GreenTree<'a, L>> for TreeNode {
125    fn from(green: oak_core::GreenTree<'a, L>) -> Self {
126        TreeNode::from(&green)
127    }
128}
129
130impl<'a, L: oak_core::Language> From<oak_core::RedTree<'a, L>> for TreeNode {
131    fn from(red: oak_core::RedTree<'a, L>) -> Self {
132        TreeNode::from(&red)
133    }
134}
135
136impl<'a, L: oak_core::Language> crate::Visualize for oak_core::GreenTree<'a, L> {
137    fn visualize(&self) -> crate::Result<String> {
138        TreeNode::from(self).visualize()
139    }
140}
141
142impl<'a, L: oak_core::Language> crate::Visualize for oak_core::RedTree<'a, L> {
143    fn visualize(&self) -> crate::Result<String> {
144        TreeNode::from(self).visualize()
145    }
146}
147
148/// Tree layout engine.
149pub struct TreeLayout {
150    algorithm: TreeLayoutAlgorithm,
151    config: TreeLayoutConfig,
152}
153
154impl Default for TreeLayout {
155    fn default() -> Self {
156        Self { algorithm: TreeLayoutAlgorithm::Layered, config: TreeLayoutConfig::default() }
157    }
158}
159
160impl TreeLayout {
161    /// Creates a new tree layout engine with default settings.
162    pub fn new() -> Self {
163        Self::default()
164    }
165
166    /// Sets the layout algorithm and returns the modified engine.
167    pub fn with_algorithm(mut self, algorithm: TreeLayoutAlgorithm) -> Self {
168        self.algorithm = algorithm;
169        self
170    }
171
172    /// Sets the layout configuration and returns the modified engine.
173    pub fn with_config(mut self, config: TreeLayoutConfig) -> Self {
174        self.config = config;
175        self
176    }
177
178    /// Visualizes the tree as an SVG string.
179    pub fn visualize(&self, tree: &TreeNode) -> crate::Result<String> {
180        let layout = self.layout_tree(tree)?;
181        crate::render::SvgRenderer::new().render_layout(&layout)
182    }
183
184    /// Calculates the layout for the given tree.
185    pub fn layout_tree(&self, tree: &TreeNode) -> crate::Result<Layout> {
186        match self.algorithm {
187            TreeLayoutAlgorithm::Layered => self.layered_layout(tree),
188            TreeLayoutAlgorithm::Radial => self.radial_layout(tree),
189            TreeLayoutAlgorithm::Compact => self.compact_layout(tree),
190            TreeLayoutAlgorithm::Balloon => self.balloon_layout(tree),
191        }
192    }
193
194    /// Layered tree layout (traditional top-down or left-right)
195    fn layered_layout(&self, tree: &TreeNode) -> crate::Result<Layout> {
196        let mut layout = Layout::new();
197        let mut node_positions = HashMap::new();
198        let mut level_widths = HashMap::new();
199
200        // Calculate level widths
201        self.calculate_level_widths(tree, 0, &mut level_widths);
202
203        // Position nodes
204        let mut current_positions = HashMap::new();
205        self.position_layered_node(tree, 0, 0.0, &level_widths, &mut current_positions, &mut node_positions);
206
207        // Add nodes to layout
208        for (id, (position, size, label, node_type)) in node_positions {
209            let rect = Rect::new(position, size);
210            let nt = match node_type.as_str() {
211                "function" => crate::layout::NodeType::Function,
212                "struct" => crate::layout::NodeType::Struct,
213                "module" => crate::layout::NodeType::Module,
214                _ => crate::layout::NodeType::Default,
215            };
216            layout.add_node_with_metadata(id, label, rect, nt);
217        }
218
219        // Add edges
220        self.add_tree_edges(tree, &mut layout);
221
222        Ok(layout)
223    }
224
225    /// Radial tree layout (root at center, children in circles)
226    fn radial_layout(&self, tree: &TreeNode) -> crate::Result<Layout> {
227        let mut layout = Layout::new();
228        let mut node_positions = HashMap::new();
229
230        // Position root at center
231        let root_size = self.get_node_size(tree);
232        let root_pos = Point::new(0.0, 0.0);
233        node_positions.insert(tree.id.clone(), (root_pos, root_size, tree.label.clone(), tree.node_type.clone()));
234
235        // Position children in concentric circles
236        if !tree.children.is_empty() {
237            let radius = self.config.level_distance;
238            self.position_radial_children(&tree.children, root_pos, radius, 0.0, 2.0 * std::f64::consts::PI, 1, &mut node_positions);
239        }
240
241        // Add nodes to layout
242        for (id, (position, size, label, node_type)) in node_positions {
243            let rect = Rect::new(Point::new(position.x - size.width / 2.0, position.y - size.height / 2.0), size);
244            let nt = match node_type.as_str() {
245                "function" => crate::layout::NodeType::Function,
246                "struct" => crate::layout::NodeType::Struct,
247                "module" => crate::layout::NodeType::Module,
248                _ => crate::layout::NodeType::Default,
249            };
250            layout.add_node_with_metadata(id, label, rect, nt);
251        }
252
253        // Add edges
254        self.add_tree_edges(tree, &mut layout);
255
256        Ok(layout)
257    }
258
259    /// Compact tree layout (minimize space usage)
260    fn compact_layout(&self, tree: &TreeNode) -> crate::Result<Layout> {
261        let mut layout = Layout::new();
262        let mut node_positions = HashMap::new();
263
264        // Use a modified layered approach with tighter packing
265        let positioned_tree = self.position_compact_node(tree, 0.0, 0.0, 0);
266        self.extract_positions(&positioned_tree, &mut node_positions);
267
268        // Add nodes to layout
269        for (id, (position, size, label, node_type)) in node_positions {
270            let rect = Rect::new(position, size);
271            let nt = match node_type.as_str() {
272                "function" => crate::layout::NodeType::Function,
273                "struct" => crate::layout::NodeType::Struct,
274                "module" => crate::layout::NodeType::Module,
275                _ => crate::layout::NodeType::Default,
276            };
277            layout.add_node_with_metadata(id, label, rect, nt);
278        }
279
280        // Add edges
281        self.add_tree_edges(tree, &mut layout);
282
283        Ok(layout)
284    }
285
286    /// Balloon tree layout (children arranged in balloon-like clusters)
287    fn balloon_layout(&self, tree: &TreeNode) -> crate::Result<Layout> {
288        let mut layout = Layout::new();
289        let mut node_positions = HashMap::new();
290
291        // Similar to radial but with balloon-like clustering
292        let root_size = self.get_node_size(tree);
293        let root_pos = Point::new(0.0, 0.0);
294        node_positions.insert(tree.id.clone(), (root_pos, root_size, tree.label.clone(), tree.node_type.clone()));
295
296        // Position children in balloon clusters
297        if !tree.children.is_empty() {
298            let cluster_radius = self.config.level_distance;
299            self.position_balloon_children(&tree.children, root_pos, cluster_radius, &mut node_positions);
300        }
301
302        // Add nodes to layout
303        for (id, (position, size, label, node_type)) in node_positions {
304            let rect = Rect::new(Point::new(position.x - size.width / 2.0, position.y - size.height / 2.0), size);
305            let nt = match node_type.as_str() {
306                "function" => crate::layout::NodeType::Function,
307                "struct" => crate::layout::NodeType::Struct,
308                "module" => crate::layout::NodeType::Module,
309                _ => crate::layout::NodeType::Default,
310            };
311            layout.add_node_with_metadata(id, label, rect, nt);
312        }
313
314        // Add edges
315        self.add_tree_edges(tree, &mut layout);
316
317        Ok(layout)
318    }
319
320    // Helper methods
321    fn get_node_size(&self, node: &TreeNode) -> Size {
322        node.size.unwrap_or(Size::new(self.config.node_width, self.config.node_height))
323    }
324
325    fn calculate_level_widths(&self, node: &TreeNode, level: usize, level_widths: &mut HashMap<usize, f64>) {
326        let node_size = self.get_node_size(node);
327        let current_width = level_widths.get(&level).unwrap_or(&0.0);
328        level_widths.insert(level, current_width + node_size.width + self.config.sibling_distance);
329
330        for child in &node.children {
331            self.calculate_level_widths(child, level + 1, level_widths);
332        }
333    }
334
335    fn position_layered_node(&self, node: &TreeNode, level: usize, _parent_x: f64, level_widths: &HashMap<usize, f64>, current_positions: &mut HashMap<usize, f64>, node_positions: &mut HashMap<String, (Point, Size, String, String)>) {
336        let node_size = self.get_node_size(node);
337        let level_width = level_widths.get(&level).unwrap_or(&0.0);
338        let default_x = -level_width / 2.0;
339        let current_x = current_positions.get(&level).unwrap_or(&default_x);
340
341        let x = if level == 0 {
342            0.0 // Root at center
343        }
344        else {
345            *current_x + node_size.width / 2.0
346        };
347
348        let y = level as f64 * self.config.level_distance;
349
350        node_positions.insert(node.id.clone(), (Point::new(x, y), node_size, node.label.clone(), node.node_type.clone()));
351        current_positions.insert(level, current_x + node_size.width + self.config.sibling_distance);
352
353        // Position children
354        for child in &node.children {
355            self.position_layered_node(child, level + 1, x, level_widths, current_positions, node_positions);
356        }
357    }
358
359    fn position_radial_children(&self, children: &[TreeNode], center: Point, radius: f64, start_angle: f64, angle_span: f64, level: usize, node_positions: &mut HashMap<String, (Point, Size, String, String)>) {
360        if children.is_empty() {
361            return;
362        }
363
364        let angle_step = angle_span / children.len() as f64;
365
366        for (i, child) in children.iter().enumerate() {
367            let angle = start_angle + i as f64 * angle_step + angle_step / 2.0;
368            let child_pos = Point::new(center.x + radius * angle.cos(), center.y + radius * angle.sin());
369
370            let child_size = self.get_node_size(child);
371            node_positions.insert(child.id.clone(), (child_pos, child_size, child.label.clone(), child.node_type.clone()));
372
373            // Recursively position grandchildren
374            if !child.children.is_empty() {
375                let child_radius = radius + self.config.level_distance;
376                let child_angle_span = angle_step * 0.8; // Reduce angle span for children
377                self.position_radial_children(&child.children, child_pos, child_radius, angle - child_angle_span / 2.0, child_angle_span, level + 1, node_positions);
378            }
379        }
380    }
381
382    fn position_compact_node(&self, node: &TreeNode, x: f64, y: f64, level: usize) -> PositionedTreeNode {
383        let size = self.get_node_size(node);
384        let mut positioned_children = Vec::new();
385        let mut child_x = x;
386
387        for child in &node.children {
388            let positioned_child = self.position_compact_node(child, child_x, y + self.config.level_distance, level + 1);
389            child_x += positioned_child.subtree_width + self.config.sibling_distance;
390            positioned_children.push(positioned_child);
391        }
392
393        let subtree_width = if positioned_children.is_empty() { size.width } else { positioned_children.iter().map(|c| c.subtree_width).sum::<f64>() + (positioned_children.len() - 1) as f64 * self.config.sibling_distance };
394
395        // Center the node over its children
396        let node_x = if positioned_children.is_empty() { x } else { x + subtree_width / 2.0 - size.width / 2.0 };
397
398        PositionedTreeNode { id: node.id.clone(), label: node.label.clone(), node_type: node.node_type.clone(), position: Point::new(node_x, y), size, subtree_width, children: positioned_children }
399    }
400
401    fn position_balloon_children(&self, children: &[TreeNode], center: Point, radius: f64, node_positions: &mut HashMap<String, (Point, Size, String, String)>) {
402        if children.is_empty() {
403            return;
404        }
405
406        // Arrange children in a circle around the parent
407        let angle_step = 2.0 * std::f64::consts::PI / children.len() as f64;
408
409        for (i, child) in children.iter().enumerate() {
410            let angle = i as f64 * angle_step;
411            let child_pos = Point::new(center.x + radius * angle.cos(), center.y + radius * angle.sin());
412
413            let child_size = self.get_node_size(child);
414            node_positions.insert(child.id.clone(), (child_pos, child_size, child.label.clone(), child.node_type.clone()));
415
416            // Recursively position grandchildren in smaller balloons
417            if !child.children.is_empty() {
418                let child_radius = radius * 0.6; // Smaller radius for child balloons
419                self.position_balloon_children(&child.children, child_pos, child_radius, node_positions);
420            }
421        }
422    }
423
424    fn extract_positions(&self, positioned_node: &PositionedTreeNode, positions: &mut HashMap<String, (Point, Size, String, String)>) {
425        positions.insert(positioned_node.id.clone(), (positioned_node.position, positioned_node.size, positioned_node.label.clone(), positioned_node.node_type.clone()));
426
427        for child in &positioned_node.children {
428            self.extract_positions(child, positions);
429        }
430    }
431
432    /// Recursively add edges from the given node to all its descendants
433    fn add_tree_edges(&self, node: &TreeNode, layout: &mut Layout) {
434        for child in &node.children {
435            let edge = Edge::new(node.id.clone(), child.id.clone());
436            layout.add_edge(edge);
437
438            // Recursively add edges for children
439            self.add_tree_edges(child, layout);
440        }
441    }
442}
443
444/// Tree layout configuration.
445#[derive(Debug, Clone)]
446pub struct TreeLayoutConfig {
447    /// Default width for nodes.
448    pub node_width: f64,
449    /// Default height for nodes.
450    pub node_height: f64,
451    /// Distance between levels of the tree.
452    pub level_distance: f64,
453    /// Distance between sibling nodes.
454    pub sibling_distance: f64,
455    /// Distance between subtrees.
456    pub subtree_distance: f64,
457}
458
459impl Default for TreeLayoutConfig {
460    fn default() -> Self {
461        Self { node_width: 100.0, node_height: 40.0, level_distance: 80.0, sibling_distance: 20.0, subtree_distance: 40.0 }
462    }
463}
464
465/// Tree layout algorithms.
466#[derive(Debug, Clone, Copy, PartialEq, Eq)]
467pub enum TreeLayoutAlgorithm {
468    /// Traditional hierarchical layout (top-down or left-right).
469    Layered,
470    /// Radial/circular layout with the root at the center.
471    Radial,
472    /// Space-efficient layout that minimizes tree width.
473    Compact,
474    /// Balloon-style clustering layout.
475    Balloon,
476}
477
478/// Positioned tree node (internal representation)
479#[derive(Debug, Clone)]
480struct PositionedTreeNode {
481    id: String,
482    label: String,
483    node_type: String,
484    position: Point,
485    size: Size,
486    subtree_width: f64,
487    children: Vec<PositionedTreeNode>,
488}
489
490/// Tree renderer for generating visual output.
491pub struct TreeRenderer {
492    config: TreeRenderConfig,
493}
494
495impl TreeRenderer {
496    /// Creates a new tree renderer with default settings.
497    pub fn new() -> Self {
498        Self { config: TreeRenderConfig::default() }
499    }
500
501    /// Sets the rendering configuration and returns the modified renderer.
502    pub fn with_config(mut self, config: TreeRenderConfig) -> Self {
503        self.config = config;
504        self
505    }
506
507    // SVG rendering functionality would be implemented here
508    // Currently disabled due to missing svg dependency
509    // #[cfg(feature = "svg")]
510    // pub fn render_svg(&self, layout: &Layout, tree: &TreeNode) -> crate::Result<String> {
511    //     use svg::{
512    //         Document,
513    //         node::element::{Group, Line, Rectangle, Text},
514    //     }
515    //
516    //     let mut document = Document::new().set("viewBox", format!("0 0 {} {}", layout.bounds.width(), layout.bounds.height()));
517    //
518    //     // Render nodes
519    //     for (node_id, rect) in &layout.nodes {
520    //         let node_group = Group::new().set("id", format!("node-{}", node_id));
521    //
522    //         // Node rectangle
523    //         let node_rect = Rectangle::new()
524    //             .set("x", rect.x())
525    //             .set("y", rect.y())
526    //             .set("width", rect.width())
527    //             .set("height", rect.height())
528    //             .set("fill", &self.config.node_fill_color)
529    //             .set("stroke", &self.config.node_stroke_color)
530    //             .set("stroke-width", self.config.node_stroke_width);
531    //
532    //         // Node label
533    //         let label = self.get_node_label(tree, node_id);
534    //         let text = Text::new()
535    //             .set("x", rect.center().x)
536    //             .set("y", rect.center().y)
537    //             .set("text-anchor", "middle")
538    //             .set("dominant-baseline", "middle")
539    //             .set("font-family", &self.config.font_family)
540    //             .set("font-size", self.config.font_size)
541    //             .set("fill", &self.config.text_color)
542    //             .add(svg::node::Text::new(label));
543    //
544    //         document = document.add(node_group.add(node_rect).add(text));
545    //     }
546    //
547    //     // Render edges
548    //     for edge in &layout.edges {
549    //         if let (Some(from_rect), Some(to_rect)) = (layout.nodes.get(&edge.from), layout.nodes.get(&edge.to)) {
550    //             let from_point = from_rect.center();
551    //             let to_point = to_rect.center();
552    //
553    //             let line = Line::new()
554    //                 .set("x1", from_point.x)
555    //                 .set("y1", from_point.y)
556    //                 .set("x2", to_point.x)
557    //                 .set("y2", to_point.y)
558    //                 .set("stroke", &self.config.edge_color)
559    //                 .set("stroke-width", self.config.edge_width);
560    //
561    //             document = document.add(line);
562    //         }
563    //     }
564    //
565    //     Ok(document.to_string())
566    // }
567
568    #[allow(dead_code)]
569    fn get_node_label(&self, tree: &TreeNode, node_id: &str) -> String {
570        self.find_node_label(tree, node_id).unwrap_or_else(|| node_id.to_string())
571    }
572
573    #[allow(dead_code)]
574    fn find_node_label(&self, node: &TreeNode, target_id: &str) -> Option<String> {
575        if node.id == target_id {
576            return Some(node.label.clone());
577        }
578
579        node.children.iter().find_map(|child| self.find_node_label(child, target_id))
580    }
581}
582
583impl Default for TreeRenderer {
584    fn default() -> Self {
585        Self::new()
586    }
587}
588
589/// Tree rendering configuration.
590#[derive(Debug, Clone)]
591pub struct TreeRenderConfig {
592    /// Fill color for node rectangles (hex string).
593    pub node_fill_color: String,
594    /// Stroke color for node rectangles (hex string).
595    pub node_stroke_color: String,
596    /// Width of the node rectangle stroke.
597    pub node_stroke_width: f64,
598    /// Color for tree edges (hex string).
599    pub edge_color: String,
600    /// Width of the tree edges.
601    pub edge_width: f64,
602    /// Color for node text (hex string).
603    pub text_color: String,
604    /// Font family for node labels.
605    pub font_family: String,
606    /// Font size for node labels.
607    pub font_size: f64,
608}
609
610impl Default for TreeRenderConfig {
611    fn default() -> Self {
612        Self {
613            node_fill_color: "#ffffff".to_string(),
614            node_stroke_color: "#000000".to_string(),
615            node_stroke_width: 1.0,
616            edge_color: "#666666".to_string(),
617            edge_width: 1.0,
618            text_color: "#000000".to_string(),
619            font_family: "Arial, sans-serif".to_string(),
620            font_size: 12.0,
621        }
622    }
623}
oak_visualize/tree/mod.rs

oak_visualize/tree/
mod.rs
