allow humand-friendly durations and sizes in CLI

1 files changed, 248 insertions, 0 deletions

diff --git a/src/units.rs b/src/units.rs
new file mode 100644
index 0000000..e4576bc
--- /dev/null
+++ b/src/units.rs
@@ -0,0 +1,248 @@
+//! Durations and sizes as operators write them: in options, configuration files and
+//! API requests.
+//!
+//! * A duration is a bare number of seconds (`30`, `0.5`), or components with units in
+//!   descending order: `ms`, `s`, `m`, `h`, `d` (`500ms`, `500s`, `1h30m`, `1d2h3s`,
+//!   `1.5h`; spaces between components are allowed). Resolution is a millisecond: a
+//!   value with a finer part is an error, not rounded.
+//! * A size is a bare number of bytes, or a number with a unit: `B`, `KiB`, `MiB`, `GiB`
+//!   (`64KiB`, `1.5MiB`). A value that isn't a whole number of bytes is an error.
+//!
+//! Output (stats, reports) stays numeric, with the unit in the field's name.
+
+use std::time::Duration;
+
+const DURATION_UNITS: [(&str, u64); 5] = [
+    ("d", 86_400_000),
+    ("h", 3_600_000),
+    ("m", 60_000),
+    ("s", 1_000),
+    ("ms", 1),
+];
+
+const SIZE_UNITS: [(&str, u64); 4] = [
+    ("GiB", 1 << 30),
+    ("MiB", 1 << 20),
+    ("KiB", 1 << 10),
+    ("B", 1),
+];
+
+/// `number` (digits, optionally a fraction) times `scale`, if that is a whole number
+/// that fits.
+fn scaled(number: &str, scale: u64) -> Option<u64> {
+    let (whole, fraction) = number.split_once('.').unwrap_or((number, ""));
+    if whole.is_empty() && fraction.is_empty() {
+        return None;
+    }
+    if !whole.bytes().all(|b| b.is_ascii_digit()) || !fraction.bytes().all(|b| b.is_ascii_digit()) {
+        return None;
+    }
+    let whole: u64 = if whole.is_empty() {
+        0
+    } else {
+        whole.parse().ok()?
+    };
+    let mut value = whole.checked_mul(scale)?;
+    if !fraction.is_empty() {
+        let digits = u32::try_from(fraction.len()).ok()?;
+        let denominator = 10u128.checked_pow(digits)?;
+        let numerator = fraction.parse::<u128>().ok()? * u128::from(scale);
+        if numerator % denominator != 0 {
+            return None;
+        }
+        value = value.checked_add(u64::try_from(numerator / denominator).ok()?)?;
+    }
+    Some(value)
+}
+
+/// A duration as written by an operator (see the module doc).
+pub(crate) fn parse_duration(text: &str) -> Result<Duration, String> {
+    let text = text.trim();
+    let invalid = || {
+        format!(
+            "{text:?} is not a duration (seconds, or e.g. 500ms, 30s, 1h30m, 1d2h3s; at most millisecond precision)"
+        )
+    };
+    if text.is_empty() {
+        return Err(invalid());
+    }
+    if let Some(millis) = scaled(text, 1_000) {
+        return Ok(Duration::from_millis(millis));
+    }
+    let mut rest = text;
+    let mut total: u64 = 0;
+    // Units must descend: the index of the next one allowed.
+    let mut next_unit = 0;
+    while !rest.is_empty() {
+        let number_len = rest
+            .find(|c: char| !(c.is_ascii_digit() || c == '.'))
+            .ok_or_else(invalid)?;
+        let (number, after) = rest.split_at(number_len);
+        let unit_len = after
+            .find(|c: char| !c.is_ascii_alphabetic())
+            .unwrap_or(after.len());
+        let (unit, after) = after.split_at(unit_len);
+        let index = DURATION_UNITS
+            .iter()
+            .position(|(name, _)| *name == unit)
+            .filter(|&index| index >= next_unit)
+            .ok_or_else(invalid)?;
+        let millis = scaled(number, DURATION_UNITS[index].1).ok_or_else(invalid)?;
+        total = total.checked_add(millis).ok_or_else(invalid)?;
+        next_unit = index + 1;
+        rest = after.trim_start();
+    }
+    Ok(Duration::from_millis(total))
+}
+
+/// A duration the way [`parse_duration`] reads it: `1d2h3s`, `500ms`, `0s`.
+pub(crate) fn format_duration(duration: Duration) -> String {
+    let mut millis = u64::try_from(duration.as_millis()).unwrap_or(u64::MAX);
+    if millis == 0 {
+        return "0s".to_owned();
+    }
+    let mut out = String::new();
+    for (name, scale) in DURATION_UNITS {
+        if millis >= scale {
+            out.push_str(&format!("{}{name}", millis / scale));
+            millis %= scale;
+        }
+    }
+    out
+}
+
+/// A size in bytes as written by an operator (see the module doc).
+pub(crate) fn parse_size(text: &str) -> Result<u64, String> {
+    let text = text.trim();
+    let invalid = || format!("{text:?} is not a size (bytes, or e.g. 512B, 64KiB, 1.5MiB, 1GiB)");
+    let number_len = text
+        .find(|c: char| !(c.is_ascii_digit() || c == '.'))
+        .unwrap_or(text.len());
+    let (number, unit) = text.split_at(number_len);
+    let scale = match unit.trim() {
+        "" => 1,
+        unit => SIZE_UNITS
+            .iter()
+            .find(|(name, _)| *name == unit)
+            .map(|(_, scale)| *scale)
+            .ok_or_else(invalid)?,
+    };
+    scaled(number, scale).ok_or_else(invalid)
+}
+
+/// A size the way [`parse_size`] reads it: the largest unit that divides it.
+pub(crate) fn format_size(bytes: u64) -> String {
+    SIZE_UNITS
+        .iter()
+        .find(|(_, scale)| bytes != 0 && bytes % scale == 0)
+        .map(|(name, scale)| format!("{}{name}", bytes / scale))
+        .unwrap_or_else(|| "0B".to_owned())
+}
+
+/// Serde for a `Duration` option: a string as [`parse_duration`] reads it, or a number
+/// of seconds (configuration files may write either); written as [`format_duration`].
+pub(crate) mod duration_option {
+    use std::time::Duration;
+
+    use serde::{Deserialize, Deserializer, Serializer, de::Error as _};
+    use serde_json::Value as JsonValue;
+
+    pub(crate) fn serialize<S: Serializer>(value: &Duration, s: S) -> Result<S::Ok, S::Error> {
+        s.serialize_str(&super::format_duration(*value))
+    }
+
+    pub(crate) fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<Duration, D::Error> {
+        match JsonValue::deserialize(d)? {
+            JsonValue::String(text) => super::parse_duration(&text).map_err(D::Error::custom),
+            JsonValue::Number(number) => {
+                super::parse_duration(&number.to_string()).map_err(D::Error::custom)
+            }
+            other => Err(D::Error::custom(format!("{other} is not a duration"))),
+        }
+    }
+}
+
+/// Serde for a size option: a string as [`parse_size`] reads it, or a number of bytes;
+/// written as [`format_size`].
+pub(crate) mod size_option {
+    use serde::{Deserialize, Deserializer, Serializer, de::Error as _};
+    use serde_json::Value as JsonValue;
+
+    pub(crate) fn serialize<S: Serializer>(value: &u64, s: S) -> Result<S::Ok, S::Error> {
+        s.serialize_str(&super::format_size(*value))
+    }
+
+    pub(crate) fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<u64, D::Error> {
+        match JsonValue::deserialize(d)? {
+            JsonValue::String(text) => super::parse_size(&text).map_err(D::Error::custom),
+            JsonValue::Number(number) => {
+                super::parse_size(&number.to_string()).map_err(D::Error::custom)
+            }
+            other => Err(D::Error::custom(format!("{other} is not a size"))),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn ms(text: &str) -> u64 {
+        parse_duration(text).unwrap().as_millis() as u64
+    }
+
+    #[test]
+    fn durations() {
+        assert_eq!(ms("30"), 30_000);
+        assert_eq!(ms("0.5"), 500);
+        assert_eq!(ms(".25"), 250);
+        assert_eq!(ms("500s"), 500_000);
+        assert_eq!(ms("500ms"), 500);
+        assert_eq!(ms("1h30m"), 5_400_000);
+        assert_eq!(ms("1d2h3s"), 93_603_000);
+        assert_eq!(ms("1d 2h 3s"), 93_603_000);
+        assert_eq!(ms("1.5h"), 5_400_000);
+        assert_eq!(ms("0.001"), 1);
+        assert_eq!(ms(" 7d "), 7 * 86_400_000);
+        assert_eq!(ms("0"), 0);
+    }
+
+    #[test]
+    fn bad_durations() {
+        for text in [
+            "", "s", "1x", "1h1d", "1s1s", "1ms1s", "-1", "1.0005", "0.0001s", "1..5s", "1h 2",
+            "1e3",
+        ] {
+            assert!(parse_duration(text).is_err(), "{text:?} parsed");
+        }
+    }
+
+    #[test]
+    fn durations_format_as_they_parse() {
+        for millis in [0, 1, 500, 1_000, 60_000, 5_400_000, 93_603_000, 86_400_001] {
+            let duration = Duration::from_millis(millis);
+            assert_eq!(
+                parse_duration(&format_duration(duration)).unwrap(),
+                duration
+            );
+        }
+        assert_eq!(format_duration(Duration::from_millis(93_603_000)), "1d2h3s");
+    }
+
+    #[test]
+    fn sizes() {
+        assert_eq!(parse_size("4096"), Ok(4096));
+        assert_eq!(parse_size("512B"), Ok(512));
+        assert_eq!(parse_size("64KiB"), Ok(65_536));
+        assert_eq!(parse_size("64 KiB"), Ok(65_536));
+        assert_eq!(parse_size("1.5MiB"), Ok(1_572_864));
+        assert_eq!(parse_size("1GiB"), Ok(1 << 30));
+        for text in ["", "KiB", "1.5", "1kib", "1KB", "0.3KiB", "-1"] {
+            assert!(parse_size(text).is_err(), "{text:?} parsed");
+        }
+        for bytes in [1, 1024, 1536, 65_536, 1 << 30] {
+            assert_eq!(parse_size(&format_size(bytes)), Ok(bytes));
+        }
+        assert_eq!(format_size(65_536), "64KiB");
+    }
+}