dux_core/tree/
arena.rs

1use std::path::{Path, PathBuf};
2
3use serde::{Deserialize, Serialize};
4
5use super::node::{NodeId, NodeKind, TreeNode};
6
7/// Arena-allocated directory tree
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct DiskTree {
10    nodes: Vec<Option<TreeNode>>,
11    root_path: PathBuf,
12}
13
14impl DiskTree {
15    pub fn new(root_path: PathBuf) -> Self {
16        let root_name = root_path
17            .file_name()
18            .map(|s| s.to_string_lossy().to_string())
19            .unwrap_or_else(|| root_path.to_string_lossy().to_string());
20
21        let root_node = TreeNode::new(
22            NodeId::ROOT,
23            root_name,
24            NodeKind::Directory,
25            root_path.clone(),
26            None,
27            0,
28        );
29
30        Self {
31            nodes: vec![Some(root_node)],
32            root_path,
33        }
34    }
35
36    /// Add a new node and return its ID
37    pub fn add_node(
38        &mut self,
39        name: String,
40        kind: NodeKind,
41        path: PathBuf,
42        parent: NodeId,
43    ) -> NodeId {
44        let parent_depth = self.get(parent).map(|n| n.depth).unwrap_or(0);
45        let id = NodeId(self.nodes.len());
46
47        let node = TreeNode::new(id, name, kind, path, Some(parent), parent_depth + 1);
48
49        self.nodes.push(Some(node));
50        if let Some(parent_node) = self.get_mut(parent) {
51            parent_node.children.push(id);
52        }
53
54        id
55    }
56
57    /// Get a reference to a node
58    pub fn get(&self, id: NodeId) -> Option<&TreeNode> {
59        self.nodes.get(id.index()).and_then(|opt| opt.as_ref())
60    }
61
62    /// Get a mutable reference to a node
63    pub fn get_mut(&mut self, id: NodeId) -> Option<&mut TreeNode> {
64        self.nodes.get_mut(id.index()).and_then(|opt| opt.as_mut())
65    }
66
67    /// Get the root node
68    pub fn root(&self) -> &TreeNode {
69        self.nodes[0].as_ref().expect("Root node must exist")
70    }
71
72    /// Get the root path
73    pub fn root_path(&self) -> &Path {
74        &self.root_path
75    }
76
77    /// Reconstruct paths and UI state for all nodes after deserialization
78    /// Must be called after loading from cache since paths and is_expanded are not serialized
79    pub fn rebuild_paths(&mut self) {
80        // Set root path and expand root
81        if let Some(root) = self.nodes.get_mut(0).and_then(|o| o.as_mut()) {
82            root.path = self.root_path.clone();
83            root.is_expanded = true;
84        }
85
86        // Process nodes in order (parents before children due to arena structure)
87        for i in 1..self.nodes.len() {
88            if let Some(node) = &self.nodes[i] {
89                let parent_path = node
90                    .parent
91                    .and_then(|pid| self.nodes.get(pid.index()))
92                    .and_then(|o| o.as_ref())
93                    .map(|p| p.path.clone());
94
95                if let Some(pp) = parent_path {
96                    let name = self.nodes[i].as_ref().map(|n| n.name.clone());
97                    if let (Some(node), Some(name)) = (self.nodes[i].as_mut(), name) {
98                        node.path = pp.join(&name);
99                    }
100                }
101            }
102        }
103    }
104
105    /// Get total number of nodes (including tombstones)
106    pub fn len(&self) -> usize {
107        self.nodes.len()
108    }
109
110    /// Get count of live (non-tombstone) nodes
111    pub fn live_count(&self) -> usize {
112        self.nodes.iter().filter(|n| n.is_some()).count()
113    }
114
115    /// Check if tree is empty (only has root)
116    pub fn is_empty(&self) -> bool {
117        self.live_count() <= 1
118    }
119
120    /// Set size for a node
121    pub fn set_size(&mut self, id: NodeId, size: u64) {
122        if let Some(node) = self.get_mut(id) {
123            node.size = size;
124        }
125    }
126
127    /// Propagate sizes from children to parents (bottom-up)
128    pub fn aggregate_sizes(&mut self) {
129        // Process nodes in reverse order (children before parents)
130        for i in (0..self.nodes.len()).rev() {
131            let node = match &self.nodes[i] {
132                Some(n) => n,
133                None => continue, // Skip tombstones
134            };
135            if node.kind.is_directory() {
136                let children = node.children.clone();
137                let mut total_size = 0u64;
138                let mut total_files = 0u64;
139
140                for child_id in &children {
141                    if let Some(child) = self.get(*child_id) {
142                        total_size += child.size;
143                        total_files += child.file_count;
144                    }
145                }
146
147                if let Some(node) = self.get_mut(NodeId(i)) {
148                    node.size = total_size;
149                    node.file_count = total_files;
150                }
151            }
152        }
153    }
154
155    /// Sort all children by size descending
156    pub fn sort_by_size(&mut self) {
157        // First collect all the size information
158        let sizes: Vec<u64> = self
159            .nodes
160            .iter()
161            .map(|n| n.as_ref().map(|n| n.size).unwrap_or(0))
162            .collect();
163
164        // Then sort each node's children
165        for node in self.nodes.iter_mut().flatten() {
166            node.children.sort_by(|a, b| {
167                let size_a = sizes.get(a.index()).copied().unwrap_or(0);
168                let size_b = sizes.get(b.index()).copied().unwrap_or(0);
169                size_b.cmp(&size_a)
170            });
171        }
172    }
173
174    /// Toggle expanded state for a node
175    pub fn toggle_expanded(&mut self, id: NodeId) {
176        if let Some(node) = self.get_mut(id)
177            && node.kind.is_directory()
178        {
179            node.is_expanded = !node.is_expanded;
180        }
181    }
182
183    /// Set expanded state for a node
184    pub fn set_expanded(&mut self, id: NodeId, expanded: bool) {
185        if let Some(node) = self.get_mut(id)
186            && node.kind.is_directory()
187        {
188            node.is_expanded = expanded;
189        }
190    }
191
192    /// Get visible nodes in tree order (respecting expansion state)
193    pub fn visible_nodes(&self, root: NodeId) -> Vec<NodeId> {
194        let mut result = Vec::new();
195        self.collect_visible(root, &mut result);
196        result
197    }
198
199    fn collect_visible(&self, id: NodeId, result: &mut Vec<NodeId>) {
200        result.push(id);
201
202        if let Some(node) = self.get(id)
203            && node.is_expanded
204        {
205            for &child_id in &node.children {
206                self.collect_visible(child_id, result);
207            }
208        }
209    }
210
211    /// Get the path from root to a node
212    pub fn path_to_node(&self, id: NodeId) -> Vec<NodeId> {
213        let mut path = Vec::new();
214        let mut current = Some(id);
215
216        while let Some(node_id) = current {
217            path.push(node_id);
218            current = self.get(node_id).and_then(|n| n.parent);
219        }
220
221        path.reverse();
222        path
223    }
224
225    /// Get breadcrumb string for a node
226    pub fn breadcrumbs(&self, id: NodeId) -> String {
227        let path = self.path_to_node(id);
228        path.iter()
229            .filter_map(|&id| self.get(id).map(|n| n.name.as_str()))
230            .collect::<Vec<_>>()
231            .join("/")
232    }
233
234    /// Expand all ancestors of a node
235    pub fn expand_to(&mut self, id: NodeId) {
236        let path = self.path_to_node(id);
237        for node_id in path {
238            self.set_expanded(node_id, true);
239        }
240    }
241
242    /// Get total size of the tree
243    pub fn total_size(&self) -> u64 {
244        self.root().size
245    }
246
247    /// Get total file count
248    pub fn total_files(&self) -> u64 {
249        self.root().file_count
250    }
251
252    /// Iterator over all live nodes (skips tombstones)
253    pub fn iter(&self) -> impl Iterator<Item = &TreeNode> {
254        self.nodes.iter().filter_map(|opt| opt.as_ref())
255    }
256
257    /// Find a node by its path
258    pub fn find_by_path(&self, path: &Path) -> Option<NodeId> {
259        for (i, node_opt) in self.nodes.iter().enumerate() {
260            if let Some(node) = node_opt
261                && node.path == path
262            {
263                return Some(NodeId(i));
264            }
265        }
266        None
267    }
268
269    /// Collect all descendant node IDs
270    fn collect_descendants(&self, id: NodeId, result: &mut Vec<NodeId>) {
271        if let Some(node) = self.get(id) {
272            for &child_id in &node.children {
273                result.push(child_id);
274                self.collect_descendants(child_id, result);
275            }
276        }
277    }
278
279    /// Remove node and descendants, return bytes freed
280    /// Does NOT perform filesystem operations - only updates tree structure
281    pub fn remove_node(&mut self, id: NodeId) -> u64 {
282        // Never remove root
283        if id == NodeId::ROOT {
284            return 0;
285        }
286
287        // Get node info before removal
288        let (size, file_count, parent_id) = match self.get(id) {
289            Some(node) => (node.size, node.file_count, node.parent),
290            None => return 0, // Already removed
291        };
292
293        // Remove from parent's children
294        if let Some(pid) = parent_id
295            && let Some(parent) = self.get_mut(pid)
296        {
297            parent.children.retain(|&c| c != id);
298        }
299
300        // Collect all descendants to tombstone
301        let mut to_remove = vec![id];
302        self.collect_descendants(id, &mut to_remove);
303
304        // Tombstone node and all descendants
305        for nid in to_remove {
306            if let Some(slot) = self.nodes.get_mut(nid.index()) {
307                *slot = None;
308            }
309        }
310
311        // Propagate size decrease up to root
312        let mut current = parent_id;
313        while let Some(nid) = current {
314            if let Some(node) = self.get_mut(nid) {
315                node.size = node.size.saturating_sub(size);
316                node.file_count = node.file_count.saturating_sub(file_count);
317                current = node.parent;
318            } else {
319                break;
320            }
321        }
322
323        size
324    }
325}
326
327#[cfg(test)]
328mod tests {
329    use super::*;
330
331    #[test]
332    fn test_tree_creation() {
333        let tree = DiskTree::new(PathBuf::from("/test"));
334        assert_eq!(tree.len(), 1);
335        assert_eq!(tree.root().name, "test");
336    }
337
338    #[test]
339    fn test_add_nodes() {
340        let mut tree = DiskTree::new(PathBuf::from("/test"));
341
342        let file_id = tree.add_node(
343            "file.txt".to_string(),
344            NodeKind::File,
345            PathBuf::from("/test/file.txt"),
346            NodeId::ROOT,
347        );
348
349        assert_eq!(tree.len(), 2);
350        assert_eq!(tree.get(file_id).unwrap().name, "file.txt");
351        assert_eq!(tree.root().children.len(), 1);
352    }
353
354    #[test]
355    fn test_remove_node_with_size_propagation() {
356        let mut tree = DiskTree::new(PathBuf::from("/test"));
357
358        // Add a subdirectory
359        let subdir_id = tree.add_node(
360            "subdir".to_string(),
361            NodeKind::Directory,
362            PathBuf::from("/test/subdir"),
363            NodeId::ROOT,
364        );
365
366        // Add files under subdir
367        let file1_id = tree.add_node(
368            "file1.txt".to_string(),
369            NodeKind::File,
370            PathBuf::from("/test/subdir/file1.txt"),
371            subdir_id,
372        );
373        tree.set_size(file1_id, 1000);
374
375        let file2_id = tree.add_node(
376            "file2.txt".to_string(),
377            NodeKind::File,
378            PathBuf::from("/test/subdir/file2.txt"),
379            subdir_id,
380        );
381        tree.set_size(file2_id, 2000);
382
383        // Aggregate sizes
384        tree.aggregate_sizes();
385
386        // Verify initial state
387        assert_eq!(tree.root().size, 3000);
388        assert_eq!(tree.get(subdir_id).unwrap().size, 3000);
389
390        // Remove file1
391        let freed = tree.remove_node(file1_id);
392        assert_eq!(freed, 1000);
393
394        // Verify sizes propagated correctly
395        assert_eq!(tree.root().size, 2000);
396        assert_eq!(tree.get(subdir_id).unwrap().size, 2000);
397
398        // file1 should be tombstoned
399        assert!(tree.get(file1_id).is_none());
400
401        // file2 should still exist
402        assert!(tree.get(file2_id).is_some());
403    }
404
405    #[test]
406    fn test_remove_node_removes_descendants() {
407        let mut tree = DiskTree::new(PathBuf::from("/test"));
408
409        // Add a subdirectory
410        let subdir_id = tree.add_node(
411            "subdir".to_string(),
412            NodeKind::Directory,
413            PathBuf::from("/test/subdir"),
414            NodeId::ROOT,
415        );
416
417        // Add files under subdir
418        let file1_id = tree.add_node(
419            "file1.txt".to_string(),
420            NodeKind::File,
421            PathBuf::from("/test/subdir/file1.txt"),
422            subdir_id,
423        );
424        tree.set_size(file1_id, 1000);
425
426        tree.aggregate_sizes();
427
428        // Remove subdir (should also remove file1)
429        let freed = tree.remove_node(subdir_id);
430        assert_eq!(freed, 1000);
431
432        // Both should be tombstoned
433        assert!(tree.get(subdir_id).is_none());
434        assert!(tree.get(file1_id).is_none());
435
436        // Root should have no children
437        assert_eq!(tree.root().children.len(), 0);
438        assert_eq!(tree.root().size, 0);
439    }
440
441    #[test]
442    fn test_find_by_path() {
443        let mut tree = DiskTree::new(PathBuf::from("/test"));
444
445        let file_id = tree.add_node(
446            "file.txt".to_string(),
447            NodeKind::File,
448            PathBuf::from("/test/file.txt"),
449            NodeId::ROOT,
450        );
451
452        // Should find by path
453        assert_eq!(
454            tree.find_by_path(&PathBuf::from("/test/file.txt")),
455            Some(file_id)
456        );
457
458        // Should return None for non-existent path
459        assert_eq!(
460            tree.find_by_path(&PathBuf::from("/test/nonexistent.txt")),
461            None
462        );
463    }
464}
dux_core/tree/arena.rs

dux_core/tree/
arena.rs
