winx_code_agent/utils/
display_tree.rs

1//! Directory-tree renderer, ported from wcgw's `display_tree.DirectoryTree`.
2//!
3//! The repo context is shown as a partially-expanded tree: only the ranked
4//! "interesting" files (and their parent directories) are expanded; everything
5//! else in a directory collapses into a single `...` line so the LLM sees
6//! structure without being flooded.
7
8use std::collections::HashSet;
9use std::fmt::Write as _;
10use std::fs;
11use std::path::{Path, PathBuf};
12
13pub struct DirectoryTree {
14    root: PathBuf,
15    expanded_files: HashSet<PathBuf>,
16    expanded_dirs: HashSet<PathBuf>,
17}
18
19impl DirectoryTree {
20    pub fn new(root: &Path) -> Self {
21        Self {
22            root: root.to_path_buf(),
23            expanded_files: HashSet::new(),
24            expanded_dirs: HashSet::new(),
25        }
26    }
27
28    /// Expand a file (given relative to the root) and all of its parent
29    /// directories, so the renderer walks down to it.
30    pub fn expand(&mut self, rel_path: &str) {
31        let abs_path = self.root.join(rel_path);
32        if !abs_path.is_file() || !abs_path.starts_with(&self.root) {
33            return;
34        }
35        self.expanded_files.insert(abs_path.clone());
36
37        let mut current = abs_path.parent().map(Path::to_path_buf);
38        while let Some(dir) = current {
39            if dir != self.root && !dir.starts_with(&self.root) {
40                break;
41            }
42            self.expanded_dirs.insert(dir.clone());
43            if dir == self.root {
44                break;
45            }
46            current = dir.parent().map(Path::to_path_buf);
47        }
48    }
49
50    /// Directory contents sorted directories-first, then case-insensitively by name.
51    fn list_directory(dir_path: &Path) -> Vec<PathBuf> {
52        let Ok(read_dir) = fs::read_dir(dir_path) else {
53            return Vec::new();
54        };
55        let mut contents: Vec<PathBuf> =
56            read_dir.filter_map(|entry| entry.ok().map(|e| e.path())).collect();
57        contents.sort_by(|a, b| {
58            let a_is_dir = a.is_dir();
59            let b_is_dir = b.is_dir();
60            // `not is_dir` first key → directories (false) come before files (true).
61            (!a_is_dir, file_name_lower(a)).cmp(&(!b_is_dir, file_name_lower(b)))
62        });
63        contents
64    }
65
66    fn count_hidden(dir_path: &Path, shown: &[PathBuf]) -> (usize, usize) {
67        let shown_set: HashSet<&PathBuf> = shown.iter().collect();
68        let mut hidden_files = 0;
69        let mut hidden_dirs = 0;
70        for item in Self::list_directory(dir_path) {
71            if shown_set.contains(&item) {
72                continue;
73            }
74            if item.is_dir() {
75                hidden_dirs += 1;
76            } else {
77                hidden_files += 1;
78            }
79        }
80        (hidden_files, hidden_dirs)
81    }
82
83    pub fn display(&self) -> String {
84        let mut out = String::new();
85        self.display_recursive(&self.root, 0, 0, &mut out);
86        out
87    }
88
89    fn display_recursive(&self, current: &Path, indent: usize, depth: usize, out: &mut String) {
90        if current == self.root {
91            let _ = writeln!(out, "{}/", current.display());
92        } else {
93            let name = file_name_str(current);
94            let _ = writeln!(out, "{:indent$}{}/", "", name, indent = indent);
95        }
96
97        // Past the top level, only descend into directories we explicitly expanded.
98        if depth > 0 && !self.expanded_dirs.contains(current) {
99            return;
100        }
101
102        let mut shown = Vec::new();
103        for item in Self::list_directory(current) {
104            let should_show =
105                self.expanded_files.contains(&item) || self.expanded_dirs.contains(&item);
106            if !should_show {
107                continue;
108            }
109            shown.push(item.clone());
110            if item.is_dir() {
111                self.display_recursive(&item, indent + 2, depth + 1, out);
112            } else {
113                let _ = writeln!(out, "{:width$}{}", "", file_name_str(&item), width = indent + 2);
114            }
115        }
116
117        let (hidden_files, hidden_dirs) = Self::count_hidden(current, &shown);
118        if hidden_files > 0 || hidden_dirs > 0 {
119            let _ = writeln!(out, "{:width$}...", "", width = indent + 2);
120        }
121    }
122}
123
124fn file_name_str(path: &Path) -> String {
125    path.file_name().map(|n| n.to_string_lossy().into_owned()).unwrap_or_default()
126}
127
128fn file_name_lower(path: &Path) -> String {
129    file_name_str(path).to_lowercase()
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135    use crate::errors::Result;
136    use tempfile::TempDir;
137
138    #[test]
139    fn renders_expanded_files_and_collapses_rest() -> Result<()> {
140        let temp = TempDir::new()?;
141        let root = temp.path();
142        fs::create_dir(root.join("src"))?;
143        fs::write(root.join("src/main.rs"), "fn main() {}\n")?;
144        fs::write(root.join("src/hidden.rs"), "\n")?;
145        fs::write(root.join("README.md"), "x\n")?;
146
147        let mut tree = DirectoryTree::new(root);
148        tree.expand("src/main.rs");
149        let display = tree.display();
150
151        assert!(display.contains("src/"));
152        assert!(display.contains("main.rs"));
153        // hidden.rs and README.md were never expanded → collapsed into "..."
154        assert!(display.contains("..."));
155        assert!(!display.contains("hidden.rs"));
156        Ok(())
157    }
158}
winx_code_agent/utils/display_tree.rs

winx_code_agent/utils/
display_tree.rs
