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

oak_visualize/tree/
mod.rs
