142 lines
4.1 KiB
Rust
142 lines
4.1 KiB
Rust
use super::*;
|
|
use std::cmp::Reverse;
|
|
use std::collections::BinaryHeap;
|
|
|
|
/// Which size drives layout and sorting.
|
|
#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)]
|
|
pub enum SizeMode {
|
|
#[default]
|
|
Allocated, // disk usage
|
|
Logical, // apparent size
|
|
}
|
|
|
|
impl SizeMode {
|
|
pub fn name(self) -> &'static str {
|
|
match self {
|
|
SizeMode::Logical => "logical",
|
|
SizeMode::Allocated => "allocated",
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Node {
|
|
/// The aggregate size in the given mode.
|
|
#[inline]
|
|
pub fn total(&self, m: SizeMode) -> i64 {
|
|
match m {
|
|
SizeMode::Logical => self.tot_size,
|
|
SizeMode::Allocated => self.tot_alloc,
|
|
}
|
|
}
|
|
/// The node's own size in the given mode.
|
|
#[inline]
|
|
pub fn own(&self, m: SizeMode) -> i64 {
|
|
match m {
|
|
SizeMode::Logical => self.size,
|
|
SizeMode::Allocated => self.alloc,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Keeps the k entries with the largest keys (ties: smaller id wins); O(n log k).
|
|
pub struct TopK {
|
|
k: usize,
|
|
// Min-heap on (key, Reverse(id)): the root is the weakest entry.
|
|
h: BinaryHeap<Reverse<(i64, Reverse<NodeId>)>>,
|
|
}
|
|
|
|
impl TopK {
|
|
pub fn new(k: usize) -> TopK {
|
|
TopK { k, h: BinaryHeap::with_capacity(k.min(1 << 16)) }
|
|
}
|
|
|
|
pub fn offer(&mut self, id: NodeId, key: i64) {
|
|
if self.k == 0 {
|
|
return;
|
|
}
|
|
let e = Reverse((key, Reverse(id)));
|
|
if self.h.len() < self.k {
|
|
self.h.push(e);
|
|
} else if let Some(mut top) = self.h.peek_mut() {
|
|
if e < *top {
|
|
*top = e;
|
|
}
|
|
}
|
|
}
|
|
|
|
/// The collected IDs, largest key first (ties by ID).
|
|
pub fn sorted(self) -> Vec<NodeId> {
|
|
// into_sorted_vec is ascending in Reverse order = descending key.
|
|
self.h.into_sorted_vec().into_iter().map(|Reverse((_, Reverse(id)))| id).collect()
|
|
}
|
|
}
|
|
|
|
impl Tree {
|
|
/// Visits id and all its descendants depth-first (iteratively). fn
|
|
/// returns false to skip a node's children.
|
|
pub fn walk(&self, id: NodeId, mut f: impl FnMut(NodeId, &Node) -> bool) {
|
|
let mut stack = vec![id];
|
|
while let Some(cur) = stack.pop() {
|
|
let n = self.node(cur);
|
|
if !f(cur, n) {
|
|
continue;
|
|
}
|
|
stack.extend(self.children(cur).map(|(c, _)| c));
|
|
}
|
|
}
|
|
|
|
/// The k nodes beneath (and including) root with the largest key,
|
|
/// considering only nodes accepted by keep.
|
|
pub fn top(&self, root: NodeId, k: usize, keep: impl Fn(&Node) -> bool, key: impl Fn(&Node) -> i64) -> Vec<NodeId> {
|
|
let mut tk = TopK::new(k);
|
|
self.walk(root, |id, n| {
|
|
if keep(n) {
|
|
tk.offer(id, key(n));
|
|
}
|
|
true
|
|
});
|
|
tk.sorted()
|
|
}
|
|
|
|
/// The children of id ordered by size descending (ties by ID).
|
|
pub fn sorted_children(&self, id: NodeId, m: SizeMode) -> Vec<NodeId> {
|
|
let mut out: Vec<(i64, NodeId)> = self.children(id).map(|(c, n)| (n.total(m), c)).collect();
|
|
out.sort_unstable_by(|a, b| b.0.cmp(&a.0).then(a.1.cmp(&b.1)));
|
|
out.into_iter().map(|(_, c)| c).collect()
|
|
}
|
|
|
|
/// Aggregates file statistics by extension beneath root.
|
|
pub fn ext_stats(&self, root: NodeId) -> Vec<ExtStat> {
|
|
let mut stats = vec![ExtStat::default(); self.ext_count()];
|
|
self.walk(root, |_, n| {
|
|
if n.kind == Kind::File && n.flags & FLAG_HARDLINK_DUP == 0 {
|
|
let s = &mut stats[n.ext as usize];
|
|
s.count += 1;
|
|
s.size += n.size;
|
|
s.alloc += n.alloc;
|
|
}
|
|
true
|
|
});
|
|
stats
|
|
.into_iter()
|
|
.enumerate()
|
|
.filter(|(_, s)| s.count > 0)
|
|
.map(|(i, mut s)| {
|
|
s.ext = self.exts[i].clone();
|
|
s.cat = category_for(&s.ext, false, false);
|
|
s
|
|
})
|
|
.collect()
|
|
}
|
|
}
|
|
|
|
/// Summarises one extension.
|
|
#[derive(Clone, Default, Debug)]
|
|
pub struct ExtStat {
|
|
pub ext: String,
|
|
pub cat: Category,
|
|
pub count: i64,
|
|
pub size: i64, // logical
|
|
pub alloc: i64,
|
|
}
|