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
13pub type FileTree = BTreeMap<String, FileNode>;
14
15/// Builds a nested BTreeMap representing the file structure.
16pub fn build_file_tree(files: &[DirEntry], base_path: &Path) -> FileTree {
17    let mut tree = BTreeMap::new();
18    for entry in files {
19        let path = entry.path().strip_prefix(base_path).unwrap_or_else(|_| entry.path());
20        let components: Vec<_> = path.components().collect();
21        
22        // Insert this path into the tree
23        insert_path(&mut tree, &components);
24    }
25    tree
26}
27
28/// Helper function to insert a path into the tree structure
29fn insert_path(tree: &mut FileTree, components: &[std::path::Component]) {
30    if components.is_empty() {
31        return;
32    }
33    
34    let name = components[0].as_os_str().to_string_lossy().to_string();
35    
36    if components.len() == 1 {
37        // This is the last component, so it's a file
38        tree.insert(name, FileNode::File);
39    } else {
40        // This is a directory component
41        // Make sure the directory exists
42        tree.entry(name.clone())
43            .or_insert_with(|| FileNode::Directory(BTreeMap::new()));
44        
45        // Recursively insert the rest of the path
46        if let Some(FileNode::Directory(next_dir)) = tree.get_mut(&name) {
47            insert_path(next_dir, &components[1..]);
48        }
49    }
50}
51
52/// Recursively prints the file tree to the console.
53pub fn print_tree(tree: &FileTree, depth: usize) {
54    for (name, node) in tree {
55        let indent = "  ".repeat(depth);
56        match node {
57            FileNode::File => {
58                println!("{}- 📄 {}", indent, name);
59            }
60            FileNode::Directory(children) => {
61                println!("{}- 📁 {}", indent, name);
62                print_tree(children, depth + 1);
63            }
64        }
65    }
66}
67
68/// Recursively writes the file tree to a file.
69pub fn write_tree_to_file(output: &mut impl Write, tree: &FileTree, depth: usize) -> io::Result<()> {
70    for (name, node) in tree {
71        let indent = "  ".repeat(depth);
72        match node {
73            FileNode::File => {
74                writeln!(output, "{}- 📄 {}", indent, name)?;
75            }
76            FileNode::Directory(children) => {
77                writeln!(output, "{}- 📁 {}", indent, name)?;
78                write_tree_to_file(output, children, depth + 1)?;
79            }
80        }
81    }
82    Ok(())
83}
84
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89    use crate::file_utils::collect_files;
90    use std::fs;
91    use tempfile::tempdir;
92
93    #[test]
94    fn test_build_file_tree_with_collected_files() {
95        // 1. Setup a temporary directory with a file structure
96        let dir = tempdir().unwrap();
97        let base_path = dir.path();
98
99        fs::create_dir(base_path.join("src")).unwrap();
100        fs::File::create(base_path.join("src/main.rs")).unwrap();
101        fs::File::create(base_path.join("README.md")).unwrap();
102        // Add a hidden file that should be ignored by default
103        fs::File::create(base_path.join(".env")).unwrap();
104
105        // 2. Collect files using the actual function
106        let files = collect_files(base_path, &[], &[]).unwrap();
107
108        // 3. Assert that the correct files were collected (hidden file is ignored)
109        assert_eq!(files.len(), 2);
110
111        // 4. Build the tree with the collected files
112        let tree = build_file_tree(&files, base_path);
113
114        // 5. Assert the tree structure is correct
115        let mut expected: FileTree = BTreeMap::new();
116        let mut src_tree = BTreeMap::new();
117        src_tree.insert("main.rs".to_string(), FileNode::File);
118        expected.insert("src".to_string(), FileNode::Directory(src_tree));
119        expected.insert("README.md".to_string(), FileNode::File);
120
121        assert_eq!(tree, expected);
122    }
123}
context_builder/tree.rs

context_builder/
tree.rs
