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}
context_builder/tree.rs

context_builder/
tree.rs
