Display bandwidth in different unit families (#328)

* Allow selecting unit families

- Supported units are {binary,SI}-{bytes,bits}

* Fix typo

* Better error msg in the unreachable case

* - I can't believe I did this. Frankly, terrible.

* Add unit test

* Add peta&pebi units to be absolutely future-proof

* Minor code style improvement

Cargo.toml

@@ -41,6 +41,7 @@ thiserror = "1.0.50"
 tokio = { version = "1.33", features = ["rt", "sync"] }
 trust-dns-resolver = "0.23.2"
 unicode-width = "0.1.11"
+strum = { version = "0.25.0", features = ["derive"] }
 
 [target.'cfg(target_os = "linux")'.dependencies]
 procfs = "0.16.0"

src/cli.rs

@@ -4,6 +4,8 @@ use clap::{Args, Parser};
 use clap_verbosity_flag::{InfoLevel, Verbosity};
 use derivative::Derivative;
 
+use crate::display::BandwidthUnitFamily;
+
 #[derive(Clone, Debug, Derivative, Parser)]
 #[derivative(Default)]
 #[command(name = "bandwhich", version)]
@@ -54,6 +56,10 @@ pub struct RenderOpts {
     /// Show remote addresses table only
     pub addresses: bool,
 
+    #[arg(short, long, value_enum, default_value_t)]
+    /// Choose a specific family of units
+    pub unit_family: BandwidthUnitFamily,
+
     #[arg(short, long)]
     /// Show total (cumulative) usages
     pub total_utilization: bool,

src/display/components/display_bandwidth.rs

@@ -1,24 +1,126 @@
 use std::fmt;
 
+use clap::ValueEnum;
+use strum::EnumIter;
+
+#[derive(Copy, Clone, Debug)]
 pub struct DisplayBandwidth {
     pub bandwidth: f64,
+    pub unit_family: BandwidthUnitFamily,
 }
 
 impl fmt::Display for DisplayBandwidth {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        // see https://github.com/rust-lang/rust/issues/41620
-        let (div, suffix) = if self.bandwidth >= 1e12 {
-            (1_099_511_627_776.0, "TiB")
-        } else if self.bandwidth >= 1e9 {
-            (1_073_741_824.0, "GiB")
-        } else if self.bandwidth >= 1e6 {
-            (1_048_576.0, "MiB")
-        } else if self.bandwidth >= 1e3 {
-            (1024.0, "KiB")
-        } else {
-            (1.0, "B")
+        let (div, suffix) = self.unit_family.get_unit_for(self.bandwidth);
+        write!(f, "{:.2}{suffix}", self.bandwidth / div)
+    }
+}
+
+#[derive(Copy, Clone, Debug, Default, Eq, PartialEq, ValueEnum, EnumIter)]
+pub enum BandwidthUnitFamily {
+    #[default]
+    /// bytes, in powers of 2^10
+    BinBytes,
+    /// bits, in powers of 2^10
+    BinBits,
+    /// bytes, in powers of 10^3
+    SiBytes,
+    /// bits, in powers of 10^3
+    SiBits,
+}
+impl BandwidthUnitFamily {
+    #[inline]
+    /// Returns an array of tuples, corresponding to the steps of this unit family.
+    ///
+    /// Each step contains a divisor, an upper bound, and a unit suffix.
+    fn steps(&self) -> [(f64, f64, &'static str); 6] {
+        /// The fraction of the next unit the value has to meet to step up.
+        const STEP_UP_FRAC: f64 = 0.95;
+        /// Binary base: 2^10.
+        const BB: f64 = 1024.0;
+
+        use BandwidthUnitFamily as F;
+        // probably could macro this stuff, but I'm too lazy
+        match self {
+            F::BinBytes => [
+                (1.0, BB * STEP_UP_FRAC, "B"),
+                (BB, BB.powi(2) * STEP_UP_FRAC, "KiB"),
+                (BB.powi(2), BB.powi(3) * STEP_UP_FRAC, "MiB"),
+                (BB.powi(3), BB.powi(4) * STEP_UP_FRAC, "GiB"),
+                (BB.powi(4), BB.powi(5) * STEP_UP_FRAC, "TiB"),
+                (BB.powi(5), f64::MAX, "PiB"),
+            ],
+            F::BinBits => [
+                (1.0 / 8.0, BB / 8.0 * STEP_UP_FRAC, "b"),
+                (BB / 8.0, BB.powi(2) / 8.0 * STEP_UP_FRAC, "Kib"),
+                (BB.powi(2) / 8.0, BB.powi(3) / 8.0 * STEP_UP_FRAC, "Mib"),
+                (BB.powi(3) / 8.0, BB.powi(4) / 8.0 * STEP_UP_FRAC, "Gib"),
+                (BB.powi(4) / 8.0, BB.powi(5) / 8.0 * STEP_UP_FRAC, "Tib"),
+                (BB.powi(5) / 8.0, f64::MAX, "Pib"),
+            ],
+            F::SiBytes => [
+                (1.0, 1e3 * STEP_UP_FRAC, "B"),
+                (1e3, 1e6 * STEP_UP_FRAC, "kB"),
+                (1e6, 1e9 * STEP_UP_FRAC, "MB"),
+                (1e9, 1e12 * STEP_UP_FRAC, "GB"),
+                (1e12, 1e15 * STEP_UP_FRAC, "TB"),
+                (1e15, f64::MAX, "PB"),
+            ],
+            F::SiBits => [
+                (1.0 / 8.0, 1e3 / 8.0 * STEP_UP_FRAC, "b"),
+                (1e3 / 8.0, 1e6 / 8.0 * STEP_UP_FRAC, "kb"),
+                (1e6 / 8.0, 1e9 / 8.0 * STEP_UP_FRAC, "Mb"),
+                (1e9 / 8.0, 1e12 / 8.0 * STEP_UP_FRAC, "Gb"),
+                (1e12 / 8.0, 1e15 / 8.0 * STEP_UP_FRAC, "Tb"),
+                (1e15 / 8.0, f64::MAX, "Pb"),
+            ],
+        }
+    }
+
+    /// Select a unit for a given value, returning its divisor and suffix.
+    fn get_unit_for(&self, bytes: f64) -> (f64, &'static str) {
+        let Some((div, _, suffix)) = self
+            .steps()
+            .into_iter()
+            .find(|&(_, bound, _)| bound >= bytes)
+        else {
+            panic!("Cannot select an appropriate unit for {bytes:.2}B.")
         };
 
-        write!(f, "{:.2}{suffix}", self.bandwidth / div)
+        (div, suffix)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::fmt::Write;
+
+    use insta::assert_snapshot;
+    use itertools::Itertools;
+    use strum::IntoEnumIterator;
+
+    use crate::display::{BandwidthUnitFamily, DisplayBandwidth};
+
+    #[test]
+    fn bandwidth_formatting() {
+        let test_bandwidths_formatted = BandwidthUnitFamily::iter()
+            .cartesian_product(
+                // I feel like this is a decent selection of values
+                (-6..60)
+                    .map(|exp| 2f64.powi(exp))
+                    .chain((-5..45).map(|exp| 2.5f64.powi(exp)))
+                    .chain((-4..38).map(|exp| 3f64.powi(exp)))
+                    .chain((-3..26).map(|exp| 5f64.powi(exp))),
+            )
+            .map(|(unit_family, bandwidth)| DisplayBandwidth {
+                bandwidth,
+                unit_family,
+            })
+            .fold(String::new(), |mut buf, b| {
+                let _ = writeln!(buf, "{b:?}: {b}");
+                buf
+            });
+
+        assert_snapshot!(test_bandwidths_formatted);
     }
 }

src/display/components/header_details.rs

@@ -78,11 +78,14 @@ impl<'a> HeaderDetails<'a> {
         } else {
             "Rate"
         };
+        let unit_family = self.state.unit_family;
         let up = DisplayBandwidth {
             bandwidth: self.state.total_bytes_uploaded as f64,
+            unit_family,
         };
         let down = DisplayBandwidth {
             bandwidth: self.state.total_bytes_downloaded as f64,
+            unit_family,
         };
         let paused = if self.paused { " [PAUSED]" } else { "" };
         format!(" Total {t} (Up / Down): {up} / {down}{paused}")