use crate::{
ByteFormat,
snapshot::{Replay, ReplayEntries},
};
use anyhow::{Context, Result};
use owo_colors::OwoColorize;
use std::{
borrow::Cow,
cmp::Ordering,
io::{self, Read, Seek},
path::{Path, PathBuf},
};
#[derive(Clone, Eq, Ord, PartialEq, PartialOrd)]
struct KeyPart {
name: Vec<u8>,
sibling_ordinal: u64,
}
#[derive(Clone, Copy)]
struct Current {
depth: usize,
size: u128,
is_dir: bool,
}
struct Cursor<'a, R> {
entries: ReplayEntries<'a, R>,
current: Option<Current>,
root_ordinal: usize,
roots_seen: usize,
key: Vec<KeyPart>,
key_depth: usize,
path: PathBuf,
path_depth: usize,
}
#[derive(Clone, Copy)]
enum Change {
Added(u128),
Removed(u128),
Modified { sign: char, magnitude: u128 },
}
struct Location<'a> {
root_ordinal: usize,
key: &'a [KeyPart],
path: &'a Path,
depth: usize,
is_dir: bool,
}
#[derive(Default)]
struct TreeState {
root_ordinal: Option<usize>,
key: Vec<KeyPart>,
collapsed: Option<CollapsedChange>,
}
struct CollapsedChange {
root_ordinal: usize,
key: Vec<KeyPart>,
path: PathBuf,
depth: usize,
additions: Option<u128>,
removals: Option<u128>,
}
impl CollapsedChange {
fn matches(&self, location: &Location<'_>, visible_depth: usize) -> bool {
self.root_ordinal == location.root_ordinal && self.key == location.key[..=visible_depth]
}
fn record(&mut self, change: Change) -> Result<()> {
let (total, magnitude, direction) = match change {
Change::Added(magnitude)
| Change::Modified {
sign: '+',
magnitude,
} => (&mut self.additions, magnitude, "addition"),
Change::Removed(magnitude) | Change::Modified { magnitude, .. } => {
(&mut self.removals, magnitude, "removal")
}
};
*total = Some(
(*total)
.unwrap_or(0)
.checked_add(magnitude)
.with_context(|| format!("collapsed {direction} total overflows u128"))?,
);
Ok(())
}
}
#[derive(Default)]
struct Summary {
additions: Vec<SummaryEntry>,
removals: Vec<SummaryEntry>,
additions_total: u64,
removals_total: u64,
changes_total: u64,
}
struct SummaryEntry {
size: u128,
path: PathBuf,
is_dir: bool,
}
impl Summary {
fn record(&mut self, change: Change, entry: Option<SummaryEntry>, limit: usize) {
self.changes_total = self.changes_total.saturating_add(1);
let (entries, total) = match change {
Change::Added(_) => (&mut self.additions, &mut self.additions_total),
Change::Removed(_) => (&mut self.removals, &mut self.removals_total),
Change::Modified { .. } => return,
};
*total = total.saturating_add(1);
if limit == 0 {
return;
}
entries.push(entry.expect("additions and removals carry their summary entry"));
entries.sort_by(|left, right| {
right
.size
.cmp(&left.size)
.then_with(|| left.path.cmp(&right.path))
});
entries.truncate(limit);
}
}
impl<'a, R: Read> Cursor<'a, R> {
fn new(entries: ReplayEntries<'a, R>, prefix: Option<&Path>) -> Result<Self> {
let mut cursor = Self {
entries,
current: None,
root_ordinal: 0,
roots_seen: 0,
key: Vec::new(),
key_depth: 0,
path: PathBuf::new(),
path_depth: 0,
};
cursor.advance(prefix)?;
Ok(cursor)
}
fn advance(&mut self, prefix: Option<&Path>) -> Result<()> {
self.advance_raw()?;
self.skip_to_prefix(prefix)
}
fn skip_to_prefix(&mut self, prefix: Option<&Path>) -> Result<()> {
while self
.current
.is_some_and(|_| prefix.is_some_and(|prefix| !self.path.starts_with(prefix)))
{
self.advance_raw()?;
}
Ok(())
}
fn advance_raw(&mut self) -> Result<()> {
let Some(entry) = self.entries.next_entry()? else {
self.current = None;
return Ok(());
};
let depth = entry.depth;
let data = entry.data;
let native_name = entry.native_name;
let sibling_ordinal = entry.sibling_ordinal;
let name = entry.name();
if depth == 0 {
self.root_ordinal = self.roots_seen;
self.roots_seen = self
.roots_seen
.checked_add(1)
.context("snapshot contains too many roots")?;
self.key_depth = 0;
self.path.clear();
self.path.push(name.as_ref());
self.path_depth = 0;
} else {
while self.path_depth >= depth {
self.path.pop();
self.path_depth -= 1;
}
self.path.push(name.as_ref());
self.path_depth = depth;
}
if self.key.len() <= depth {
self.key.push(KeyPart {
name: Vec::new(),
sibling_ordinal,
});
}
let key = &mut self.key[depth];
key.name.clear();
key.name.extend_from_slice(native_name);
key.sibling_ordinal = sibling_ordinal;
self.key_depth = depth + 1;
self.current = Some(Current {
depth,
size: data.size,
is_dir: data.is_dir,
});
Ok(())
}
fn cmp_key<Rhs: Read>(&self, rhs: &Cursor<'_, Rhs>) -> Ordering {
self.root_ordinal
.cmp(&rhs.root_ordinal)
.then_with(|| self.key[..self.key_depth].cmp(&rhs.key[..rhs.key_depth]))
}
fn advance_past_current(&mut self, prefix: Option<&Path>) -> Result<()> {
let Some(current) = self.current else {
return Ok(());
};
self.advance_raw()?;
if current.is_dir {
while self
.current
.is_some_and(|entry| entry.depth > current.depth)
{
self.advance_raw()?;
}
}
self.skip_to_prefix(prefix)
}
fn location(&self, current: Current) -> Location<'_> {
Location {
root_ordinal: self.root_ordinal,
key: &self.key[..self.key_depth],
path: &self.path,
depth: current.depth,
is_dir: current.is_dir,
}
}
}
#[allow(clippy::too_many_arguments)]
pub fn diff_snapshots<Old: Read + Seek, New: Read + Seek>(
out: (impl io::Write, bool),
old: &mut Replay<Old>,
new: &mut Replay<New>,
byte_format: ByteFormat,
directories_only: bool,
prefix: Option<&Path>,
max_depth: Option<usize>,
summary_limit: usize,
) -> Result<()> {
let (mut out, out_is_terminal) = out;
let (send, receive) = std::sync::mpsc::channel();
let summary = std::thread::spawn(move || {
let mut summary = Summary::default();
for (change, entry) in receive {
summary.record(change, entry, summary_limit);
}
summary
});
let mut state = TreeState::default();
let walk_result = walk_changes(old, new, directories_only, prefix, |change, location| {
let summary_entry = match change {
_ if summary_limit == 0 => None,
Change::Added(size) | Change::Removed(size) => Some(SummaryEntry {
size,
path: location.path.to_owned(),
is_dir: location.is_dir,
}),
Change::Modified { .. } => None,
};
send.send((change, summary_entry))
.map_err(|_| anyhow::anyhow!("summary collector stopped"))?;
write_tree_entry(
&mut out,
&mut state,
change,
location,
byte_format,
out_is_terminal,
prefix,
max_depth,
!directories_only,
)?;
Ok(())
});
drop(send);
let summary = summary
.join()
.map_err(|_| anyhow::anyhow!("summary collector panicked"))?;
walk_result?;
flush_collapsed(&mut out, &mut state, byte_format, out_is_terminal, prefix)?;
if summary.changes_total == 0 {
return Ok(());
}
writeln!(out)?;
write_summary(&mut out, &summary, byte_format, out_is_terminal)?;
Ok(())
}
fn write_summary(
out: &mut impl io::Write,
summary: &Summary,
byte_format: ByteFormat,
out_is_terminal: bool,
) -> io::Result<()> {
let mut wrote_entries = false;
for (title, entries, total, change) in [
(
"removals",
&summary.removals,
summary.removals_total,
Change::Removed as fn(u128) -> Change,
),
(
"additions",
&summary.additions,
summary.additions_total,
Change::Added as fn(u128) -> Change,
),
] {
if entries.is_empty() {
continue;
}
if wrote_entries {
writeln!(out)?;
}
writeln!(
out,
"Largest {title} (showing {} of {total}):",
entries.len()
)?;
for entry in entries {
write_change(
out,
change(entry.size),
&entry.path,
entry.is_dir,
0,
false,
(byte_format, out_is_terminal),
)?;
}
wrote_entries = true;
}
if wrote_entries {
writeln!(out)?;
}
writeln!(out, "Changes: {}", summary.changes_total)
}
fn walk_changes<Old: Read + Seek, New: Read + Seek>(
old: &mut Replay<Old>,
new: &mut Replay<New>,
directories_only: bool,
prefix: Option<&Path>,
mut on_change: impl FnMut(Change, Location<'_>) -> Result<()>,
) -> Result<()> {
let mut old = Cursor::new(old.entries()?, prefix)?;
let mut new = Cursor::new(new.entries()?, prefix)?;
loop {
let ordering = match (old.current, new.current) {
(Some(_), Some(_)) => Some(old.cmp_key(&new)),
(Some(_), None) => Some(Ordering::Less),
(None, Some(_)) => Some(Ordering::Greater),
(None, None) => None,
};
match ordering {
Some(Ordering::Less) => {
let Some(entry) = old.current else {
return Ok(());
};
if !directories_only || entry.is_dir {
on_change(Change::Removed(entry.size), old.location(entry))?;
}
old.advance_past_current(prefix)?;
}
Some(Ordering::Greater) => {
let Some(entry) = new.current else {
return Ok(());
};
if !directories_only || entry.is_dir {
on_change(Change::Added(entry.size), new.location(entry))?;
}
new.advance_past_current(prefix)?;
}
Some(Ordering::Equal) => {
let (Some(old_entry), Some(new_entry)) = (old.current, new.current) else {
return Ok(());
};
if old_entry.is_dir == new_entry.is_dir {
if old_entry.is_dir == directories_only && old_entry.size != new_entry.size {
let (sign, magnitude) = if new_entry.size > old_entry.size {
('+', new_entry.size - old_entry.size)
} else {
('-', old_entry.size - new_entry.size)
};
on_change(
Change::Modified { sign, magnitude },
new.location(new_entry),
)?;
}
old.advance(prefix)?;
new.advance(prefix)?;
} else {
if !directories_only || old_entry.is_dir {
on_change(Change::Removed(old_entry.size), old.location(old_entry))?;
}
if !directories_only || new_entry.is_dir {
on_change(Change::Added(new_entry.size), new.location(new_entry))?;
}
old.advance_past_current(prefix)?;
new.advance_past_current(prefix)?;
}
}
None => return Ok(()),
}
}
}
#[allow(clippy::too_many_arguments)]
fn write_tree_entry(
out: &mut impl io::Write,
state: &mut TreeState,
change: Change,
location: Location<'_>,
byte_format: ByteFormat,
out_is_terminal: bool,
prefix: Option<&Path>,
max_depth: Option<usize>,
aggregate_hidden: bool,
) -> Result<()> {
let mut hierarchy = location
.path
.ancestors()
.take(location.depth + 1)
.collect::<Vec<_>>();
hierarchy.reverse();
let base_depth = prefix
.and_then(|prefix| hierarchy.iter().position(|path| *path == prefix))
.unwrap_or(0);
let visible_depth = location
.depth
.min(base_depth.saturating_add(max_depth.unwrap_or(usize::MAX)));
let change_is_visible = visible_depth == location.depth;
if aggregate_hidden && !change_is_visible {
if state
.collapsed
.as_ref()
.is_none_or(|collapsed| !collapsed.matches(&location, visible_depth))
{
flush_collapsed(out, state, byte_format, out_is_terminal, prefix)?;
state.collapsed = Some(CollapsedChange {
root_ordinal: location.root_ordinal,
key: location.key[..=visible_depth].to_vec(),
path: hierarchy[visible_depth].to_owned(),
depth: visible_depth,
additions: None,
removals: None,
});
}
state
.collapsed
.as_mut()
.expect("collapsed change was initialized")
.record(change)?;
return Ok(());
}
flush_collapsed(out, state, byte_format, out_is_terminal, prefix)?;
let common_depth = if state.root_ordinal == Some(location.root_ordinal) {
state
.key
.iter()
.zip(location.key)
.take_while(|(left, right)| left == right)
.count()
} else {
0
};
let context_count = if change_is_visible {
location.depth
} else {
visible_depth + 1
};
for (depth, path) in hierarchy
.iter()
.enumerate()
.take(context_count)
.skip(common_depth.max(base_depth))
{
write_context(
out,
display_path(path, depth == base_depth),
depth - base_depth,
!change_is_visible && depth == visible_depth,
out_is_terminal,
)?;
}
if change_is_visible {
write_change(
out,
change,
display_path(location.path, location.depth == base_depth),
location.is_dir,
location.depth - base_depth,
location.is_dir
&& matches!(change, Change::Modified { .. })
&& max_depth.is_some_and(|max_depth| location.depth - base_depth == max_depth),
(byte_format, out_is_terminal),
)?;
}
state.root_ordinal = Some(location.root_ordinal);
state.key.clear();
state.key.extend_from_slice(&location.key[..=visible_depth]);
Ok(())
}
fn flush_collapsed(
out: &mut impl io::Write,
state: &mut TreeState,
byte_format: ByteFormat,
out_is_terminal: bool,
prefix: Option<&Path>,
) -> io::Result<()> {
let Some(collapsed) = state.collapsed.take() else {
return Ok(());
};
let mut hierarchy = collapsed
.path
.ancestors()
.take(collapsed.depth + 1)
.collect::<Vec<_>>();
hierarchy.reverse();
let base_depth = prefix
.and_then(|prefix| hierarchy.iter().position(|path| *path == prefix))
.unwrap_or(0);
let common_depth = if state.root_ordinal == Some(collapsed.root_ordinal) {
state
.key
.iter()
.zip(&collapsed.key)
.take_while(|(left, right)| left == right)
.count()
} else {
0
};
for (depth, path) in hierarchy
.iter()
.enumerate()
.take(collapsed.depth)
.skip(common_depth.max(base_depth))
{
write_context(
out,
display_path(path, depth == base_depth),
depth - base_depth,
false,
out_is_terminal,
)?;
}
write_collapsed_change(
out,
display_path(&collapsed.path, collapsed.depth == base_depth),
collapsed.depth - base_depth,
collapsed.additions,
collapsed.removals,
byte_format,
out_is_terminal,
)?;
state.root_ordinal = Some(collapsed.root_ordinal);
state.key = collapsed.key;
Ok(())
}
fn write_context(
out: &mut impl io::Write,
path: &Path,
depth: usize,
has_hidden_changes: bool,
out_is_terminal: bool,
) -> io::Result<()> {
let path = path_for_output(path);
let suffix = directory_suffix(path.as_ref());
let hidden = if has_hidden_changes { " …" } else { "" };
let line = format!("{}{path}{suffix}{hidden}", " ".repeat(depth));
if out_is_terminal {
writeln!(out, "{}", line.cyan())
} else {
writeln!(out, "{line}")
}
}
#[allow(clippy::too_many_arguments)]
fn write_collapsed_change(
out: &mut impl io::Write,
path: &Path,
depth: usize,
additions: Option<u128>,
removals: Option<u128>,
byte_format: ByteFormat,
out_is_terminal: bool,
) -> io::Result<()> {
let path = path_for_output(path);
let suffix = directory_suffix(path.as_ref());
let prefix = format!("{}{path}{suffix} …", " ".repeat(depth));
if out_is_terminal {
write!(out, "{} (", prefix.cyan())?;
if let Some(size) = additions {
write!(out, "{}", format!("+{}", byte_format.display(size)).green())?;
}
if additions.is_some() && removals.is_some() {
write!(out, " ")?;
}
if let Some(size) = removals {
write!(out, "{}", format!("-{}", byte_format.display(size)).red())?;
}
writeln!(out, ")")
} else {
write!(out, "{prefix} (")?;
if let Some(size) = additions {
write!(out, "+{}", byte_format.display(size))?;
}
if additions.is_some() && removals.is_some() {
write!(out, " ")?;
}
if let Some(size) = removals {
write!(out, "-{}", byte_format.display(size))?;
}
writeln!(out, ")")
}
}
fn write_change(
out: &mut impl io::Write,
change: Change,
path: &Path,
is_dir: bool,
depth: usize,
has_hidden_changes: bool,
(byte_format, out_is_terminal): (ByteFormat, bool),
) -> io::Result<()> {
let (kind, sign, magnitude) = match change {
Change::Added(size) => ('+', "", size),
Change::Removed(size) => ('-', "", size),
Change::Modified { sign, magnitude } => {
('~', if sign == '+' { "+" } else { "-" }, magnitude)
}
};
let path = path_for_output(path);
let suffix = if is_dir {
directory_suffix(path.as_ref())
} else {
""
};
let hidden = if has_hidden_changes { " …" } else { "" };
let line = format!(
"{}{kind} {sign}{} {path}{suffix}{hidden}",
" ".repeat(depth),
byte_format.display(magnitude)
);
if !out_is_terminal {
return writeln!(out, "{line}");
}
match change {
Change::Added(_) => writeln!(out, "{}", line.green()),
Change::Removed(_) => writeln!(out, "{}", line.red()),
Change::Modified { .. } => writeln!(out, "{}", line.yellow()),
}
}
fn display_path(path: &Path, full: bool) -> &Path {
if full {
path
} else {
path.file_name().map_or(path, Path::new)
}
}
fn directory_suffix(path: &str) -> &'static str {
if path.chars().last().is_some_and(std::path::is_separator) {
""
} else {
std::path::MAIN_SEPARATOR_STR
}
}
fn path_for_output(path: &Path) -> Cow<'_, str> {
let path = path.to_string_lossy();
if path.chars().any(char::is_control) {
path.chars()
.map(|character| {
if character.is_control() {
'\u{FFFD}'
} else {
character
}
})
.collect::<String>()
.into()
} else {
path
}
}