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

lean_ctx/tools/
ctx_tree.rs
