lean_ctx/tools/
ctx_tree.rs

1use std::path::Path;
2
3use ignore::WalkBuilder;
4
5use crate::core::protocol;
6use crate::core::tokens::count_tokens;
7
8/// Generates a compact directory tree listing with file counts, respecting gitignore.
9pub fn handle(path: &str, depth: usize, show_hidden: bool) -> (String, usize) {
10    let root = Path::new(path);
11    if root.is_file() {
12        let parent = root
13            .parent()
14            .map_or(path.to_string(), |p| p.display().to_string());
15        return (
16            format!(
17                "ERROR: '{path}' is a file, not a directory. Use path=\"{parent}\" for the containing directory."
18            ),
19            0,
20        );
21    }
22    if !root.is_dir() {
23        return (
24            format!("ERROR: {path} does not exist or is not a directory"),
25            0,
26        );
27    }
28
29    let raw_output = generate_raw_tree(root, depth, show_hidden);
30    let compact_output = generate_compact_tree(root, depth, show_hidden);
31
32    if compact_output.trim().is_empty() {
33        return (format!("{path}/ (empty directory, depth={depth})"), 0);
34    }
35
36    let raw_tokens = count_tokens(&raw_output);
37    let compact_tokens = count_tokens(&compact_output);
38    let savings = protocol::format_savings(raw_tokens, compact_tokens);
39
40    (format!("{compact_output}\n{savings}"), raw_tokens)
41}
42
43fn generate_compact_tree(root: &Path, max_depth: usize, show_hidden: bool) -> String {
44    let mut lines = Vec::new();
45
46    struct Entry {
47        depth: usize,
48        name: String,
49        is_dir: bool,
50        path: std::path::PathBuf,
51    }
52    let mut entries: Vec<Entry> = Vec::new();
53
54    let walker = WalkBuilder::new(root)
55        .hidden(!show_hidden)
56        .git_ignore(true)
57        .git_global(true)
58        .git_exclude(true)
59        .max_depth(Some(max_depth))
60        .sort_by_file_name(std::cmp::Ord::cmp)
61        .build();
62
63    for entry in walker.filter_map(std::result::Result::ok) {
64        if entry.depth() == 0 {
65            continue;
66        }
67        entries.push(Entry {
68            depth: entry.depth(),
69            name: entry.file_name().to_string_lossy().to_string(),
70            is_dir: entry.file_type().is_some_and(|ft| ft.is_dir()),
71            path: entry.path().to_path_buf(),
72        });
73    }
74
75    let mut dir_file_counts: std::collections::HashMap<&std::path::Path, usize> =
76        std::collections::HashMap::new();
77    for e in &entries {
78        if !e.is_dir {
79            if let Some(parent) = e.path.parent() {
80                *dir_file_counts.entry(parent).or_default() += 1;
81            }
82        }
83    }
84
85    for e in &entries {
86        let indent = "  ".repeat(e.depth.saturating_sub(1));
87        if e.is_dir {
88            let count = dir_file_counts.get(e.path.as_path()).copied().unwrap_or(0);
89            lines.push(format!("{indent}{}/ ({count})", e.name));
90        } else {
91            lines.push(format!("{indent}{}", e.name));
92        }
93    }
94
95    lines.join("\n")
96}
97
98fn generate_raw_tree(root: &Path, depth: usize, show_hidden: bool) -> String {
99    let mut lines = Vec::new();
100
101    let walker = WalkBuilder::new(root)
102        .hidden(!show_hidden)
103        .git_ignore(true)
104        .git_global(true)
105        .git_exclude(true)
106        .max_depth(Some(depth))
107        .sort_by_file_name(std::cmp::Ord::cmp)
108        .build();
109
110    for entry in walker.filter_map(std::result::Result::ok) {
111        if entry.depth() == 0 {
112            continue;
113        }
114        let rel = entry
115            .path()
116            .strip_prefix(root)
117            .unwrap_or(entry.path())
118            .to_string_lossy();
119        lines.push(rel.to_string());
120    }
121
122    lines.join("\n")
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    #[test]
130    fn tree_savings_are_reasonable() {
131        let dir = env!("CARGO_MANIFEST_DIR");
132        let (output, original) = handle(dir, 3, false);
133        let compact_tokens = count_tokens(&output);
134
135        eprintln!("=== ctx_tree savings test ===");
136        eprintln!("  original (raw) tokens: {original}");
137        eprintln!("  compact tokens:        {compact_tokens}");
138        eprintln!(
139            "  savings:               {}",
140            original.saturating_sub(compact_tokens)
141        );
142
143        assert!(
144            original < 5000,
145            "raw tree at depth 3 should be < 5000 tokens, got {original}"
146        );
147        assert!(original > 0, "raw tree should have some tokens");
148        if original > compact_tokens {
149            let ratio = (original - compact_tokens) as f64 / original as f64;
150            eprintln!("  savings ratio:         {:.1}%", ratio * 100.0);
151            assert!(
152                ratio < 0.90,
153                "savings ratio should be < 90% for same-depth comparison, got {:.1}%",
154                ratio * 100.0
155            );
156        }
157    }
158}
lean_ctx/tools/ctx_tree.rs

lean_ctx/tools/
ctx_tree.rs
