128 lines
3.7 KiB
Rust
128 lines
3.7 KiB
Rust
use std::path::{PathBuf, Component};
|
|
|
|
#[derive(Debug, Clone, Default)]
|
|
pub struct Directory {
|
|
pub name: String,
|
|
pub size: u64,
|
|
children: Vec<Directory>,
|
|
}
|
|
impl Directory {
|
|
pub fn new(name: String, size: u64) -> Self {
|
|
Self {
|
|
name,
|
|
size,
|
|
children: Vec::new(),
|
|
}
|
|
}
|
|
|
|
pub fn insert(&mut self, path: PathBuf, size: u64) {
|
|
let components =
|
|
path.components()
|
|
.filter_map(
|
|
|cmp| match cmp {
|
|
Component::Normal(os_str) => os_str.to_str(),
|
|
_ => None,
|
|
}
|
|
).map(ToOwned::to_owned);
|
|
|
|
let mut current = self;
|
|
for next in components {
|
|
if &next == ¤t.name {
|
|
continue;
|
|
}
|
|
|
|
let idx = if let Some(idx) = current.children.iter().position(|d| d.name == next) {
|
|
idx
|
|
} else {
|
|
current.children.push(Directory::new(next, size));
|
|
current.children.len()-1
|
|
};
|
|
current = &mut current.children[idx];
|
|
}
|
|
}
|
|
|
|
/// Recursive function to find the greaest file size
|
|
/// in this directory tree, used for formatting in the output
|
|
fn find_max_size(&self) -> u64 {
|
|
self.children.iter()
|
|
.map(|d| d.find_max_size())
|
|
.max()
|
|
.unwrap_or(self.size)
|
|
}
|
|
|
|
/// public-exposed print function, does the setup for the
|
|
/// real print function, `print2`
|
|
pub fn print(self) {
|
|
// fake print function to give the depth param
|
|
let max_size = self.find_max_size();
|
|
let tabwidth = max_size.to_string().len() + 2;
|
|
|
|
let mut stack = Vec::new();
|
|
self.print2(tabwidth as usize, &mut stack);
|
|
}
|
|
|
|
/// real print function that makes the tree
|
|
fn print2(self, tabwidth: usize, stack: &mut Vec<TreePart>) {
|
|
let indent = build_indent(stack);
|
|
println!("{:tabwidth$}{indent} {}", self.size.to_string(), self.name);
|
|
|
|
stack.push(TreePart::Edge);
|
|
let mut iter = self.children.into_iter().peekable();
|
|
while let Some(child) = iter.next() {
|
|
if iter.peek().is_none() {
|
|
let idx = stack.len()-1;
|
|
stack[idx] = TreePart::Corner;
|
|
};
|
|
child.print2(tabwidth, stack);
|
|
}
|
|
stack.pop();
|
|
}
|
|
}
|
|
impl Extend<(PathBuf, u64)> for Directory {
|
|
fn extend<T: IntoIterator<Item = (PathBuf, u64)>>(&mut self, iter: T) {
|
|
for (p, s) in iter {
|
|
self.insert(p, s);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(PartialEq, Eq, Debug, Clone, Copy)]
|
|
enum TreePart {
|
|
/// Rightmost column, not the last in the directory
|
|
Edge,
|
|
/// Not rightmost, but dir not finished yet
|
|
Line,
|
|
/// Rightmost column in the last directory
|
|
Corner,
|
|
/// Not rightmost, but the dir has finished
|
|
Blank
|
|
}
|
|
impl TreePart {
|
|
/// convert to ascii art
|
|
pub fn display(&self) -> &'static str {
|
|
match self {
|
|
Self::Edge => "├──",
|
|
Self::Line => "│ ",
|
|
Self::Corner => "└──",
|
|
Self::Blank => " ",
|
|
}
|
|
}
|
|
}
|
|
/// Construct the indent string for the given stack.
|
|
/// must be called at the top of the recursive function,
|
|
/// and does mutate the stack
|
|
fn build_indent(stack: &mut Vec<TreePart>) -> String {
|
|
let mut indent = String::new();
|
|
for (i, mut tp) in stack.iter().enumerate() {
|
|
if i < stack.len()-1 && tp == &TreePart::Edge {
|
|
tp = &TreePart::Line;
|
|
}
|
|
indent += tp.display();
|
|
}
|
|
// essentially, if the last time element on the stack was a corner,
|
|
// make it blank for all future prints
|
|
if let Some(last @ TreePart::Corner) = stack.last_mut() {
|
|
*last = TreePart::Blank;
|
|
}
|
|
indent
|
|
} |