path: root/src/units.rs

//! 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");
    }
}