context_builder/
tree.rs

1use ignore::DirEntry;
2use std::collections::BTreeMap;
3use std::io::{self, Write};
4use std::path::Path;
5
6/// A nested map to represent the file tree structure.
7#[derive(Debug, Clone, PartialEq)]
8pub enum FileNode {
9    File,
10    Directory(BTreeMap<String, FileNode>),
11}
12
13/// Type alias for the file tree structure.
14pub type FileTree = BTreeMap<String, FileNode>;
15
16/// Builds a nested BTreeMap representing the file structure.
17pub fn build_file_tree(files: &[DirEntry], base_path: &Path) -> FileTree {
18    let mut tree = BTreeMap::new();
19    for entry in files {
20        let path = entry
21            .path()
22            .strip_prefix(base_path)
23            .unwrap_or_else(|_| entry.path());
24        let components: Vec<_> = path.components().collect();
25
26        // Insert this path into the tree
27        insert_path(&mut tree, &components);
28    }
29    tree
30}
31
32/// Helper function to insert a path into the tree structure
33fn insert_path(tree: &mut FileTree, components: &[std::path::Component]) {
34    if components.is_empty() {
35        return;
36    }
37
38    let name = components[0].as_os_str().to_string_lossy().to_string();
39
40    if components.len() == 1 {
41        // This is the last component, so it's a file
42        tree.insert(name, FileNode::File);
43    } else {
44        // This is a directory component
45        // Make sure the directory exists
46        tree.entry(name.clone())
47            .or_insert_with(|| FileNode::Directory(BTreeMap::new()));
48
49        // Recursively insert the rest of the path
50        if let Some(FileNode::Directory(next_dir)) = tree.get_mut(&name) {
51            insert_path(next_dir, &components[1..]);
52        }
53    }
54}
55
56/// Recursively prints the file tree to the console.
57pub fn print_tree(tree: &FileTree, depth: usize) {
58    for (name, node) in tree {
59        let indent = "  ".repeat(depth);
60        match node {
61            FileNode::File => {
62                println!("{}- 📄 {}", indent, name);
63            }
64            FileNode::Directory(children) => {
65                println!("{}- 📁 {}", indent, name);
66                print_tree(children, depth + 1);
67            }
68        }
69    }
70}
71
72/// Recursively writes the file tree to a file.
73pub fn write_tree_to_file(
74    output: &mut impl Write,
75    tree: &FileTree,
76    depth: usize,
77) -> io::Result<()> {
78    for (name, node) in tree {
79        let indent = "  ".repeat(depth);
80        match node {
81            FileNode::File => {
82                writeln!(output, "{}- 📄 {}", indent, name)?;
83            }
84            FileNode::Directory(children) => {
85                writeln!(output, "{}- 📁 {}", indent, name)?;
86                write_tree_to_file(output, children, depth + 1)?;
87            }
88        }
89    }
90    Ok(())
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96    use crate::file_utils::collect_files;
97    use std::fs;
98    use tempfile::tempdir;
99
100    #[test]
101    fn test_build_file_tree_with_collected_files() {
102        // 1. Set up a temporary directory with a file structure
103        let dir = tempdir().unwrap();
104        let base_path = dir.path();
105
106        fs::create_dir(base_path.join("src")).unwrap();
107        fs::File::create(base_path.join("src/main.rs")).unwrap();
108        fs::File::create(base_path.join("README.md")).unwrap();
109        // Add a hidden file that should be ignored by default
110        fs::File::create(base_path.join(".env")).unwrap();
111
112        // 2. Collect files using the actual function
113        let files = collect_files(base_path, &[], &[]).unwrap();
114
115        // 3. Assert that the correct files were collected (a hidden file is ignored)
116        assert_eq!(files.len(), 2);
117
118        // 4. Build the tree with the collected files
119        let tree = build_file_tree(&files, base_path);
120
121        // 5. Assert the tree structure is correct
122        let mut expected: FileTree = BTreeMap::new();
123        let mut src_tree = BTreeMap::new();
124        src_tree.insert("main.rs".to_string(), FileNode::File);
125        expected.insert("src".to_string(), FileNode::Directory(src_tree));
126        expected.insert("README.md".to_string(), FileNode::File);
127
128        assert_eq!(tree, expected);
129    }
130
131    #[test]
132    fn test_build_file_tree_empty() {
133        let dir = tempdir().unwrap();
134        let base_path = dir.path();
135
136        let files = collect_files(base_path, &[], &[]).unwrap();
137        let tree = build_file_tree(&files, base_path);
138
139        assert!(tree.is_empty());
140    }
141
142    #[test]
143    fn test_build_file_tree_single_file() {
144        let dir = tempdir().unwrap();
145        let base_path = dir.path();
146
147        fs::File::create(base_path.join("single.txt")).unwrap();
148
149        let files = collect_files(base_path, &[], &[]).unwrap();
150        let tree = build_file_tree(&files, base_path);
151
152        let mut expected: FileTree = BTreeMap::new();
153        expected.insert("single.txt".to_string(), FileNode::File);
154
155        assert_eq!(tree, expected);
156    }
157
158    #[test]
159    fn test_build_file_tree_nested_directories() {
160        let dir = tempdir().unwrap();
161        let base_path = dir.path();
162
163        fs::create_dir_all(base_path.join("a/b/c")).unwrap();
164        fs::File::create(base_path.join("a/b/c/deep.txt")).unwrap();
165        fs::File::create(base_path.join("a/shallow.txt")).unwrap();
166
167        let files = collect_files(base_path, &[], &[]).unwrap();
168        let tree = build_file_tree(&files, base_path);
169
170        // Build expected structure
171        let mut c_tree = BTreeMap::new();
172        c_tree.insert("deep.txt".to_string(), FileNode::File);
173
174        let mut b_tree = BTreeMap::new();
175        b_tree.insert("c".to_string(), FileNode::Directory(c_tree));
176
177        let mut a_tree = BTreeMap::new();
178        a_tree.insert("b".to_string(), FileNode::Directory(b_tree));
179        a_tree.insert("shallow.txt".to_string(), FileNode::File);
180
181        let mut expected: FileTree = BTreeMap::new();
182        expected.insert("a".to_string(), FileNode::Directory(a_tree));
183
184        assert_eq!(tree, expected);
185    }
186
187    #[test]
188    fn test_build_file_tree_unicode_filenames() {
189        let dir = tempdir().unwrap();
190        let base_path = dir.path();
191
192        fs::create_dir(base_path.join("测试目录")).unwrap();
193        fs::File::create(base_path.join("测试目录/文件.txt")).unwrap();
194        fs::File::create(base_path.join("🦀.rs")).unwrap();
195
196        let files = collect_files(base_path, &[], &[]).unwrap();
197        let tree = build_file_tree(&files, base_path);
198
199        let mut test_dir = BTreeMap::new();
200        test_dir.insert("文件.txt".to_string(), FileNode::File);
201
202        let mut expected: FileTree = BTreeMap::new();
203        expected.insert("测试目录".to_string(), FileNode::Directory(test_dir));
204        expected.insert("🦀.rs".to_string(), FileNode::File);
205
206        assert_eq!(tree, expected);
207    }
208
209    #[test]
210    fn test_insert_path_empty_components() {
211        let mut tree = BTreeMap::new();
212        insert_path(&mut tree, &[]);
213        assert!(tree.is_empty());
214    }
215
216    #[test]
217    fn test_write_tree_to_file() {
218        let mut tree = BTreeMap::new();
219        tree.insert("file1.txt".to_string(), FileNode::File);
220
221        let mut subdir = BTreeMap::new();
222        subdir.insert("file2.md".to_string(), FileNode::File);
223        tree.insert("src".to_string(), FileNode::Directory(subdir));
224
225        let mut output = Vec::new();
226        write_tree_to_file(&mut output, &tree, 0).unwrap();
227
228        let result = String::from_utf8(output).unwrap();
229        assert!(result.contains("- 📄 file1.txt"));
230        assert!(result.contains("- 📁 src"));
231        assert!(result.contains("  - 📄 file2.md"));
232    }
233
234    #[test]
235    fn test_write_tree_to_file_with_depth() {
236        let mut tree = BTreeMap::new();
237        tree.insert("nested.txt".to_string(), FileNode::File);
238
239        let mut output = Vec::new();
240        write_tree_to_file(&mut output, &tree, 2).unwrap();
241
242        let result = String::from_utf8(output).unwrap();
243        assert!(result.contains("    - 📄 nested.txt")); // 2 levels of indentation
244    }
245
246    #[test]
247    fn test_write_tree_to_file_empty_tree() {
248        let tree = BTreeMap::new();
249        let mut output = Vec::new();
250        write_tree_to_file(&mut output, &tree, 0).unwrap();
251
252        let result = String::from_utf8(output).unwrap();
253        assert!(result.is_empty());
254    }
255
256    #[test]
257    fn test_file_node_equality() {
258        let file1 = FileNode::File;
259        let file2 = FileNode::File;
260        assert_eq!(file1, file2);
261
262        let mut dir1 = BTreeMap::new();
263        dir1.insert("test.txt".to_string(), FileNode::File);
264        let node1 = FileNode::Directory(dir1.clone());
265        let node2 = FileNode::Directory(dir1);
266        assert_eq!(node1, node2);
267
268        // Different directories should not be equal
269        let mut dir2 = BTreeMap::new();
270        dir2.insert("other.txt".to_string(), FileNode::File);
271        let node3 = FileNode::Directory(dir2);
272        assert_ne!(node1, node3);
273
274        // File and directory should not be equal
275        assert_ne!(file1, node1);
276    }
277
278    #[test]
279    fn test_build_file_tree_absolute_path_fallback() {
280        // Test the fallback case when strip_prefix fails by using different base paths
281        let dir = tempdir().unwrap();
282        let base_path = dir.path();
283        let other_dir = tempdir().unwrap();
284        let other_base = other_dir.path();
285
286        // Create a file in the first directory
287        fs::File::create(base_path.join("test.txt")).unwrap();
288
289        // Create a DirEntry from the first directory but use a different base_path
290        let files = collect_files(base_path, &[], &[]).unwrap();
291
292        // This should trigger the unwrap_or_else case since other_base is unrelated to the file path
293        let tree = build_file_tree(&files, other_base);
294
295        // The tree should still contain the file, but with its full path
296        assert!(!tree.is_empty());
297    }
298
299    #[test]
300    fn test_build_file_tree_multiple_files_same_directory() {
301        let dir = tempdir().unwrap();
302        let base_path = dir.path();
303
304        fs::create_dir(base_path.join("docs")).unwrap();
305        fs::File::create(base_path.join("docs/readme.md")).unwrap();
306        fs::File::create(base_path.join("docs/guide.md")).unwrap();
307        fs::File::create(base_path.join("docs/api.md")).unwrap();
308
309        let files = collect_files(base_path, &[], &[]).unwrap();
310        let tree = build_file_tree(&files, base_path);
311
312        let mut docs_tree = BTreeMap::new();
313        docs_tree.insert("api.md".to_string(), FileNode::File);
314        docs_tree.insert("guide.md".to_string(), FileNode::File);
315        docs_tree.insert("readme.md".to_string(), FileNode::File);
316
317        let mut expected: FileTree = BTreeMap::new();
318        expected.insert("docs".to_string(), FileNode::Directory(docs_tree));
319
320        assert_eq!(tree, expected);
321    }
322}
context_builder/tree.rs

context_builder/
tree.rs
