From abf3e4ab50ded9c77054e45e13ec161c25be3a13 Mon Sep 17 00:00:00 2001
From: David Campbell <dcampbell24@gmail.com>
Date: Sat, 26 Oct 2024 14:56:01 -0400
Subject: [PATCH] Display what is removed from mdbook clean.

This is based off of [cargo's][1] clean command. cargo is licensed
under MIT or Apache-2.0.

[1]: https://github.com/rust-lang/cargo
---
 src/cmd/clean.rs | 89 +++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 84 insertions(+), 5 deletions(-)

diff --git a/src/cmd/clean.rs b/src/cmd/clean.rs
index 48b4147ca9..ec77537eaf 100644
--- a/src/cmd/clean.rs
+++ b/src/cmd/clean.rs
@@ -2,8 +2,9 @@ use super::command_prelude::*;
 use crate::get_book_dir;
 use anyhow::Context;
 use mdbook::MDBook;
-use std::fs;
+use std::mem::take;
 use std::path::PathBuf;
+use std::{fmt, fs};
 
 // Create clap subcommand arguments
 pub fn make_subcommand() -> Command {
@@ -23,10 +24,88 @@ pub fn execute(args: &ArgMatches) -> mdbook::errors::Result<()> {
         None => book.root.join(&book.config.build.build_dir),
     };
 
-    if dir_to_remove.exists() {
-        fs::remove_dir_all(&dir_to_remove)
-            .with_context(|| "Unable to remove the build directory")?;
-    }
+    let removed = Clean::new(&dir_to_remove)?;
+    println!("{removed}");
 
     Ok(())
 }
+
+/// Formats a number of bytes into a human readable SI-prefixed size.
+/// Returns a tuple of `(quantity, units)`.
+pub fn human_readable_bytes(bytes: u64) -> (f32, &'static str) {
+    static UNITS: [&str; 7] = ["B", "KiB", "MiB", "GiB", "TiB", "PiB", "EiB"];
+    let bytes = bytes as f32;
+    let i = ((bytes.log2() / 10.0) as usize).min(UNITS.len() - 1);
+    (bytes / 1024_f32.powi(i as i32), UNITS[i])
+}
+
+#[derive(Debug)]
+pub struct Clean {
+    num_files_removed: u64,
+    num_dirs_removed: u64,
+    total_bytes_removed: u64,
+}
+
+impl Clean {
+    fn new(dir: &PathBuf) -> mdbook::errors::Result<Clean> {
+        let mut files = vec![dir.clone()];
+        let mut children = Vec::new();
+        let mut num_files_removed = 0;
+        let mut num_dirs_removed = 0;
+        let mut total_bytes_removed = 0;
+
+        if dir.exists() {
+            while !files.is_empty() {
+                for file in files {
+                    if let Ok(meta) = file.metadata() {
+                        // Note: This can over-count bytes removed for hard-linked
+                        // files. It also under-counts since it only counts the exact
+                        // byte sizes and not the block sizes.
+                        total_bytes_removed += meta.len();
+                    }
+                    if file.is_file() {
+                        num_files_removed += 1;
+                    } else if file.is_dir() {
+                        num_dirs_removed += 1;
+                        for entry in fs::read_dir(file)? {
+                            children.push(entry?.path());
+                        }
+                    }
+                }
+                files = take(&mut children);
+            }
+            fs::remove_dir_all(&dir).with_context(|| "Unable to remove the build directory")?;
+        }
+
+        Ok(Clean {
+            num_files_removed,
+            num_dirs_removed,
+            total_bytes_removed,
+        })
+    }
+}
+
+impl fmt::Display for Clean {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "Removed ")?;
+        match (self.num_files_removed, self.num_dirs_removed) {
+            (0, 0) => write!(f, "0 files")?,
+            (0, 1) => write!(f, "1 directory")?,
+            (0, 2..) => write!(f, "{} directories", self.num_dirs_removed)?,
+            (1, _) => write!(f, "1 file")?,
+            (2.., _) => write!(f, "{} files", self.num_files_removed)?,
+        }
+
+        if self.total_bytes_removed == 0 {
+            Ok(())
+        } else {
+            // Don't show a fractional number of bytes.
+            if self.total_bytes_removed < 1024 {
+                write!(f, ", {}B total", self.total_bytes_removed)
+            } else {
+                let (bytes, unit) = human_readable_bytes(self.total_bytes_removed);
+                write!(f, ", {bytes:.2}{unit} total")
+            }
+        }
+    }
+}