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

context_builder/
tree.rs
