ctx/
tree.rs

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

ctx/
tree.rs
