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

lean_ctx/tools/
ctx_tree.rs
