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_dir() {
12        return (format!("ERROR: {path} is not a directory"), 0);
13    }
14
15    let raw_output = generate_raw_tree(root, depth, show_hidden);
16    let compact_output = generate_compact_tree(root, depth, show_hidden);
17
18    let raw_tokens = count_tokens(&raw_output);
19    let compact_tokens = count_tokens(&compact_output);
20    let savings = protocol::format_savings(raw_tokens, compact_tokens);
21
22    (format!("{compact_output}\n{savings}"), raw_tokens)
23}
24
25fn generate_compact_tree(root: &Path, max_depth: usize, show_hidden: bool) -> String {
26    let mut lines = Vec::new();
27    let mut entries: Vec<(usize, String, bool, usize)> = Vec::new();
28
29    let walker = WalkBuilder::new(root)
30        .hidden(!show_hidden)
31        .git_ignore(true)
32        .git_global(true)
33        .git_exclude(true)
34        .max_depth(Some(max_depth))
35        .sort_by_file_name(std::cmp::Ord::cmp)
36        .build();
37
38    for entry in walker.filter_map(std::result::Result::ok) {
39        if entry.depth() == 0 {
40            continue;
41        }
42
43        let name = entry.file_name().to_string_lossy().to_string();
44
45        let depth = entry.depth();
46        let is_dir = entry.file_type().is_some_and(|ft| ft.is_dir());
47
48        let file_count = if is_dir {
49            count_files_in_dir(entry.path())
50        } else {
51            0
52        };
53
54        entries.push((depth, name, is_dir, file_count));
55    }
56
57    for (depth, name, is_dir, file_count) in &entries {
58        let indent = "  ".repeat(depth.saturating_sub(1));
59        if *is_dir {
60            lines.push(format!("{indent}{name}/ ({file_count})"));
61        } else {
62            lines.push(format!("{indent}{name}"));
63        }
64    }
65
66    lines.join("\n")
67}
68
69fn generate_raw_tree(root: &Path, depth: usize, show_hidden: bool) -> String {
70    let mut lines = Vec::new();
71
72    let walker = WalkBuilder::new(root)
73        .hidden(!show_hidden)
74        .git_ignore(true)
75        .git_global(true)
76        .git_exclude(true)
77        .max_depth(Some(depth))
78        .sort_by_file_name(std::cmp::Ord::cmp)
79        .build();
80
81    for entry in walker.filter_map(std::result::Result::ok) {
82        if entry.depth() == 0 {
83            continue;
84        }
85        let rel = entry
86            .path()
87            .strip_prefix(root)
88            .unwrap_or(entry.path())
89            .to_string_lossy();
90        lines.push(rel.to_string());
91    }
92
93    lines.join("\n")
94}
95
96fn count_files_in_dir(dir: &Path) -> usize {
97    WalkBuilder::new(dir)
98        .hidden(false)
99        .git_ignore(true)
100        .max_depth(Some(5))
101        .build()
102        .filter_map(std::result::Result::ok)
103        .filter(|e| e.file_type().is_some_and(|ft| ft.is_file()))
104        .count()
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110
111    #[test]
112    fn tree_savings_are_reasonable() {
113        let dir = env!("CARGO_MANIFEST_DIR");
114        let (output, original) = handle(dir, 3, false);
115        let compact_tokens = count_tokens(&output);
116
117        eprintln!("=== ctx_tree savings test ===");
118        eprintln!("  original (raw) tokens: {original}");
119        eprintln!("  compact tokens:        {compact_tokens}");
120        eprintln!(
121            "  savings:               {}",
122            original.saturating_sub(compact_tokens)
123        );
124
125        assert!(
126            original < 5000,
127            "raw tree at depth 3 should be < 5000 tokens, got {original}"
128        );
129        assert!(original > 0, "raw tree should have some tokens");
130        if original > compact_tokens {
131            let ratio = (original - compact_tokens) as f64 / original as f64;
132            eprintln!("  savings ratio:         {:.1}%", ratio * 100.0);
133            assert!(
134                ratio < 0.90,
135                "savings ratio should be < 90% for same-depth comparison, got {:.1}%",
136                ratio * 100.0
137            );
138        }
139    }
140}
lean_ctx/tools/ctx_tree.rs

lean_ctx/tools/
ctx_tree.rs
