ctx/
tree.rs

1//! ASCII project tree rendering.
2//!
3//! Builds the `├──`/`└──` style directory tree shown at the top of
4//! generated context. Entry point: [`generate_tree`].
5
6use std::collections::BTreeMap;
7
8use crate::walker::{format_size, FileEntry};
9
10/// A node in the file tree (either a file or directory).
11#[derive(Debug)]
12enum TreeNode {
13    File {
14        size: u64,
15    },
16    Directory {
17        children: BTreeMap<String, TreeNode>,
18    },
19}
20
21impl TreeNode {
22    fn new_directory() -> Self {
23        TreeNode::Directory {
24            children: BTreeMap::new(),
25        }
26    }
27
28    fn new_file(size: u64) -> Self {
29        TreeNode::File { size }
30    }
31
32    fn get_or_create_dir(&mut self, name: &str) -> &mut TreeNode {
33        if let TreeNode::Directory { children } = self {
34            children
35                .entry(name.to_string())
36                .or_insert_with(TreeNode::new_directory)
37        } else {
38            panic!("Expected directory node");
39        }
40    }
41
42    fn insert_file(&mut self, name: &str, size: u64) {
43        if let TreeNode::Directory { children } = self {
44            children.insert(name.to_string(), TreeNode::new_file(size));
45        }
46    }
47}
48
49/// Build a tree structure from file entries.
50fn build_tree(entries: &[FileEntry]) -> TreeNode {
51    let mut root = TreeNode::new_directory();
52
53    for entry in entries {
54        let components: Vec<_> = entry.relative_path.components().collect();
55        let mut current = &mut root;
56
57        for (i, component) in components.iter().enumerate() {
58            let name = component.as_os_str().to_string_lossy().to_string();
59            let is_last = i == components.len() - 1;
60
61            if is_last {
62                current.insert_file(&name, entry.size);
63            } else {
64                current = current.get_or_create_dir(&name);
65            }
66        }
67    }
68
69    root
70}
71
72/// Render the tree to an ASCII string.
73fn render_tree(node: &TreeNode, prefix: &str, is_root: bool, show_sizes: bool) -> String {
74    let mut output = String::new();
75
76    if let TreeNode::Directory { children } = node {
77        let entries: Vec<_> = children.iter().collect();
78        let total = entries.len();
79
80        for (i, (name, child)) in entries.iter().enumerate() {
81            let is_last_entry = i == total - 1;
82
83            // Determine the connector
84            let connector = if is_root {
85                ""
86            } else if is_last_entry {
87                "└── "
88            } else {
89                "├── "
90            };
91
92            // Build the line
93            match child {
94                TreeNode::File { size } => {
95                    if show_sizes {
96                        output.push_str(&format!(
97                            "{}{}{} ({})\n",
98                            prefix,
99                            connector,
100                            name,
101                            format_size(*size)
102                        ));
103                    } else {
104                        output.push_str(&format!("{}{}{}\n", prefix, connector, name));
105                    }
106                }
107                TreeNode::Directory { .. } => {
108                    output.push_str(&format!("{}{}{}/\n", prefix, connector, name));
109
110                    // Recurse into directory
111                    let new_prefix = if is_root {
112                        prefix.to_string()
113                    } else if is_last_entry {
114                        format!("{}    ", prefix)
115                    } else {
116                        format!("{}│   ", prefix)
117                    };
118
119                    output.push_str(&render_tree(child, &new_prefix, false, show_sizes));
120                }
121            }
122        }
123    }
124
125    output
126}
127
128/// Generate an ASCII tree representation of the file entries.
129pub fn generate_tree(root_name: &str, entries: &[FileEntry], show_sizes: bool) -> String {
130    if entries.is_empty() {
131        return format!("{}/\n(empty)\n", root_name);
132    }
133
134    let tree = build_tree(entries);
135    let mut output = format!("{}/\n", root_name);
136    output.push_str(&render_tree(&tree, "", true, show_sizes));
137    output
138}
139
140#[cfg(test)]
141mod tests {
142    use super::*;
143    use std::path::PathBuf;
144
145    #[test]
146    fn test_simple_tree() {
147        let entries = vec![
148            FileEntry {
149                absolute_path: PathBuf::from("/project/src/main.rs"),
150                relative_path: PathBuf::from("src/main.rs"),
151                size: 100,
152            },
153            FileEntry {
154                absolute_path: PathBuf::from("/project/Cargo.toml"),
155                relative_path: PathBuf::from("Cargo.toml"),
156                size: 200,
157            },
158        ];
159
160        let tree = generate_tree("project", &entries, false);
161        assert!(tree.contains("project/"));
162        assert!(tree.contains("src/"));
163        assert!(tree.contains("main.rs"));
164        assert!(tree.contains("Cargo.toml"));
165    }
166
167    #[test]
168    fn test_tree_with_sizes() {
169        let entries = vec![FileEntry {
170            absolute_path: PathBuf::from("/project/file.rs"),
171            relative_path: PathBuf::from("file.rs"),
172            size: 1024,
173        }];
174
175        let tree = generate_tree("project", &entries, true);
176        assert!(tree.contains("1.0 KB"));
177    }
178}
ctx/tree.rs

ctx/
tree.rs
