Skip to main content

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 _mode_guard = crate::core::savings_footer::ModeGuard::new("tree");
43    let raw_tokens = count_tokens(&raw_output);
44    let compact_tokens = count_tokens(&compact_output);
45    let savings = protocol::format_savings(raw_tokens, compact_tokens);
46
47    (format!("{compact_output}\n{savings}"), raw_tokens)
48}
49
50fn generate_compact_tree(
51    root: &Path,
52    max_depth: usize,
53    show_hidden: bool,
54    respect_gitignore: bool,
55) -> String {
56    let mut lines = Vec::new();
57
58    struct Entry {
59        depth: usize,
60        name: String,
61        is_dir: bool,
62        path: std::path::PathBuf,
63    }
64    let mut entries: Vec<Entry> = Vec::new();
65
66    let walker = WalkBuilder::new(root)
67        .hidden(!show_hidden)
68        .git_ignore(respect_gitignore)
69        .git_global(respect_gitignore)
70        .git_exclude(respect_gitignore)
71        .max_depth(Some(max_depth))
72        .sort_by_file_name(std::cmp::Ord::cmp)
73        .build();
74
75    for entry in walker.filter_map(std::result::Result::ok) {
76        if entry.depth() == 0 {
77            continue;
78        }
79        entries.push(Entry {
80            depth: entry.depth(),
81            name: entry.file_name().to_string_lossy().to_string(),
82            is_dir: entry.file_type().is_some_and(|ft| ft.is_dir()),
83            path: entry.path().to_path_buf(),
84        });
85    }
86
87    let mut dir_file_counts: std::collections::HashMap<&std::path::Path, usize> =
88        std::collections::HashMap::new();
89    for e in &entries {
90        if !e.is_dir {
91            if let Some(parent) = e.path.parent() {
92                *dir_file_counts.entry(parent).or_default() += 1;
93            }
94        }
95    }
96
97    for e in &entries {
98        let indent = "  ".repeat(e.depth.saturating_sub(1));
99        if e.is_dir {
100            let count = dir_file_counts.get(e.path.as_path()).copied().unwrap_or(0);
101            lines.push(format!("{indent}{}/ ({count})", e.name));
102        } else {
103            lines.push(format!("{indent}{}", e.name));
104        }
105    }
106
107    lines.join("\n")
108}
109
110fn generate_raw_tree(
111    root: &Path,
112    depth: usize,
113    show_hidden: bool,
114    respect_gitignore: bool,
115) -> String {
116    let mut lines = Vec::new();
117
118    let walker = WalkBuilder::new(root)
119        .hidden(!show_hidden)
120        .git_ignore(respect_gitignore)
121        .git_global(respect_gitignore)
122        .git_exclude(respect_gitignore)
123        .max_depth(Some(depth))
124        .sort_by_file_name(std::cmp::Ord::cmp)
125        .build();
126
127    for entry in walker.filter_map(std::result::Result::ok) {
128        if entry.depth() == 0 {
129            continue;
130        }
131        let rel = entry
132            .path()
133            .strip_prefix(root)
134            .unwrap_or(entry.path())
135            .to_string_lossy();
136        lines.push(rel.to_string());
137    }
138
139    lines.join("\n")
140}
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145
146    /// Builds a deterministic source-tree fixture so the assertions do not
147    /// depend on the live repository size or platform path separators (the live
148    /// repo coupling previously made this test tip over its token threshold on
149    /// Windows as the codebase grew).
150    fn make_fixture() -> tempfile::TempDir {
151        let dir = tempfile::tempdir().unwrap();
152        let root = dir.path();
153        let files = [
154            "Cargo.toml",
155            "README.md",
156            "src/main.rs",
157            "src/lib.rs",
158            "src/core/mod.rs",
159            "src/core/engine.rs",
160            "src/core/util.rs",
161            "src/tools/mod.rs",
162            "src/tools/reader.rs",
163            "tests/integration.rs",
164            "tests/smoke.rs",
165        ];
166        for rel in files {
167            let p = root.join(rel);
168            std::fs::create_dir_all(p.parent().unwrap()).unwrap();
169            std::fs::write(&p, "// fixture\n").unwrap();
170        }
171        dir
172    }
173
174    #[test]
175    fn tree_savings_are_reasonable() {
176        let dir = make_fixture();
177        let (output, original) = handle(&dir.path().to_string_lossy(), 3, false, true);
178        let compact_tokens = count_tokens(&output);
179
180        eprintln!("=== ctx_tree savings test ===");
181        eprintln!("  original (raw) tokens: {original}");
182        eprintln!("  compact tokens:        {compact_tokens}");
183        eprintln!(
184            "  savings:               {}",
185            original.saturating_sub(compact_tokens)
186        );
187
188        assert!(original > 0, "raw tree should have some tokens");
189        assert!(
190            original < 2000,
191            "raw tree for the fixture should be small, got {original}"
192        );
193        if original > compact_tokens {
194            let ratio = (original - compact_tokens) as f64 / original as f64;
195            eprintln!("  savings ratio:         {:.1}%", ratio * 100.0);
196            assert!(
197                ratio < 0.90,
198                "savings ratio should be < 90% for same-depth comparison, got {:.1}%",
199                ratio * 100.0
200            );
201        }
202    }
203}