use core::fmt;
use core::str::FromStr;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
const KIB: u64 = 1024;
const MIB: u64 = 1024 * KIB;
const GIB: u64 = 1024 * MIB;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct MemSize(u64);
impl MemSize {
#[inline]
#[must_use]
pub const fn from_bytes(bytes: u64) -> Self {
Self(bytes)
}
#[inline]
#[must_use]
pub const fn bytes(self) -> u64 {
self.0
}
}
impl FromStr for MemSize {
type Err = ParseMemSizeError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
if s.is_empty() {
return Err(ParseMemSizeError::Empty);
}
let (digits, multiplier) = match s.as_bytes()[s.len() - 1] {
b'G' => (&s[..s.len() - 1], GIB),
b'M' => (&s[..s.len() - 1], MIB),
b'K' => (&s[..s.len() - 1], KIB),
_ => (s, 1),
};
if digits.is_empty() {
return Err(ParseMemSizeError::MissingDigits);
}
if !digits.bytes().all(|b| b.is_ascii_digit()) {
return Err(ParseMemSizeError::InvalidCharacter);
}
let value: u64 = digits.parse().map_err(|_| ParseMemSizeError::Overflow)?;
value
.checked_mul(multiplier)
.map(Self)
.ok_or(ParseMemSizeError::Overflow)
}
}
impl fmt::Display for MemSize {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self.0 {
0 => f.write_str("0"),
b if b % GIB == 0 => write!(f, "{}G", b / GIB),
b if b % MIB == 0 => write!(f, "{}M", b / MIB),
b if b % KIB == 0 => write!(f, "{}K", b / KIB),
b => write!(f, "{b}"),
}
}
}
impl Serialize for MemSize {
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
serializer.collect_str(self)
}
}
impl<'de> Deserialize<'de> for MemSize {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
let s = String::deserialize(deserializer)?;
s.parse().map_err(serde::de::Error::custom)
}
}
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ParseMemSizeError {
Empty,
MissingDigits,
InvalidCharacter,
Overflow,
}
impl fmt::Display for ParseMemSizeError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(match self {
Self::Empty => "memory size is empty",
Self::MissingDigits => "memory size has a unit suffix but no digits",
Self::InvalidCharacter => {
"memory size must be ASCII digits with an optional trailing G, M, or K"
}
Self::Overflow => "memory size in bytes overflows u64",
})
}
}
impl core::error::Error for ParseMemSizeError {}
#[cfg(feature = "schema")]
impl schemars::JsonSchema for MemSize {
fn schema_name() -> std::borrow::Cow<'static, str> {
"MemSize".into()
}
fn json_schema(_: &mut schemars::SchemaGenerator) -> schemars::Schema {
schemars::json_schema!({
"type": "string",
"pattern": r"^\d+(G|M|K)?$",
"description": "A byte quantity: digits, optionally suffixed G, M or K (binary units).",
})
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct UpDuration(core::time::Duration);
impl UpDuration {
#[inline]
#[must_use]
pub const fn from_millis(ms: u64) -> Self {
Self(core::time::Duration::from_millis(ms))
}
#[inline]
#[must_use]
pub const fn as_duration(self) -> core::time::Duration {
self.0
}
#[inline]
#[must_use]
pub const fn as_millis(self) -> u64 {
self.0.as_millis() as u64
}
}
impl FromStr for UpDuration {
type Err = ParseUpDurationError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
if s.is_empty() {
return Err(ParseUpDurationError::Empty);
}
let (digits, ms_per_unit) = match s.as_bytes()[s.len() - 1] {
b'h' => (&s[..s.len() - 1], 3_600_000),
b'm' => (&s[..s.len() - 1], 60_000),
b's' => (&s[..s.len() - 1], 1_000),
_ => (s, 1),
};
if digits.is_empty() {
return Err(ParseUpDurationError::MissingDigits);
}
if !digits.bytes().all(|b| b.is_ascii_digit()) {
return Err(ParseUpDurationError::InvalidCharacter);
}
let value: u64 = digits.parse().map_err(|_| ParseUpDurationError::Overflow)?;
value
.checked_mul(ms_per_unit)
.map(Self::from_millis)
.ok_or(ParseUpDurationError::Overflow)
}
}
impl fmt::Display for UpDuration {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let ms = self.as_millis();
match ms {
0 => f.write_str("0"),
v if v % 3_600_000 == 0 => write!(f, "{}h", v / 3_600_000),
v if v % 60_000 == 0 => write!(f, "{}m", v / 60_000),
v if v % 1_000 == 0 => write!(f, "{}s", v / 1_000),
v => write!(f, "{v}"),
}
}
}
impl Serialize for UpDuration {
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
serializer.collect_str(self)
}
}
impl<'de> Deserialize<'de> for UpDuration {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
let s = String::deserialize(deserializer)?;
s.parse().map_err(serde::de::Error::custom)
}
}
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ParseUpDurationError {
Empty,
MissingDigits,
InvalidCharacter,
Overflow,
}
impl fmt::Display for ParseUpDurationError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(match self {
Self::Empty => "duration is empty",
Self::MissingDigits => "duration has a unit suffix but no digits",
Self::InvalidCharacter => {
"duration must be ASCII digits with an optional trailing h, m, or s"
}
Self::Overflow => "duration in milliseconds overflows u64",
})
}
}
impl core::error::Error for ParseUpDurationError {}
#[cfg(feature = "schema")]
impl schemars::JsonSchema for UpDuration {
fn schema_name() -> std::borrow::Cow<'static, str> {
"UpDuration".into()
}
fn json_schema(_: &mut schemars::SchemaGenerator) -> schemars::Schema {
schemars::json_schema!({
"type": "string",
"pattern": r"^\d+(h|m|s)?$",
"description": "A duration: digits, optionally suffixed h, m or s. Plain digits are milliseconds.",
})
}
}
#[cfg(test)]
mod mem_size_tests {
use super::*;
#[test]
fn plain_digits_parse_as_bytes() {
assert_eq!("123".parse::<MemSize>().unwrap().bytes(), 123);
}
#[test]
fn units_are_binary() {
assert_eq!("7K".parse::<MemSize>().unwrap().bytes(), 7 * 1024);
assert_eq!("512M".parse::<MemSize>().unwrap().bytes(), 512 << 20);
assert_eq!("3G".parse::<MemSize>().unwrap().bytes(), 3 << 30);
}
#[test]
fn rejects_spec_violations() {
use ParseMemSizeError::*;
assert_eq!("".parse::<MemSize>(), Err(Empty));
assert_eq!("G".parse::<MemSize>(), Err(MissingDigits));
assert_eq!("512m".parse::<MemSize>(), Err(InvalidCharacter)); assert_eq!(" 512M".parse::<MemSize>(), Err(InvalidCharacter)); assert_eq!("1.5G".parse::<MemSize>(), Err(InvalidCharacter)); assert_eq!("512MB".parse::<MemSize>(), Err(InvalidCharacter)); assert_eq!("18446744073709551616".parse::<MemSize>(), Err(Overflow));
assert_eq!("17179869184G".parse::<MemSize>(), Err(Overflow));
}
#[test]
fn display_uses_largest_exact_unit_and_round_trips() {
for bytes in [
0u64,
1,
1023,
1024,
1536,
1 << 20,
(1 << 30) + 1024,
u64::MAX,
] {
let size = MemSize::from_bytes(bytes);
let reparsed: MemSize = size.to_string().parse().unwrap();
assert_eq!(reparsed, size, "display of {bytes} bytes must reparse");
}
assert_eq!(MemSize::from_bytes(3 << 30).to_string(), "3G");
assert_eq!(MemSize::from_bytes(1536).to_string(), "1536");
}
#[test]
fn serde_uses_string_form() {
let size: MemSize = serde_json::from_str("\"512M\"").unwrap();
assert_eq!(size.bytes(), 512 << 20);
assert_eq!(serde_json::to_string(&size).unwrap(), "\"512M\"");
assert!(serde_json::from_str::<MemSize>("\"512MB\"").is_err());
}
#[cfg(feature = "schema")]
#[test]
fn the_schema_pattern_agrees_with_from_str() {
let schema = serde_json::to_value(schemars::schema_for!(MemSize)).unwrap();
let pattern = schema["pattern"].as_str().unwrap();
let re = regex::Regex::new(pattern).unwrap();
for accepted in ["512M", "1G", "4096", "7K"] {
assert!(re.is_match(accepted), "pattern rejects {accepted}");
assert!(
accepted.parse::<MemSize>().is_ok(),
"FromStr rejects {accepted}"
);
}
for rejected in ["512MB", "512m", "1.5G", "", "M", "512T", "1P", "512g"] {
assert!(!re.is_match(rejected), "pattern accepts {rejected}");
assert!(
rejected.parse::<MemSize>().is_err(),
"FromStr accepts {rejected}"
);
}
}
}
#[cfg(test)]
mod up_duration_tests {
use super::*;
#[test]
fn plain_digits_are_milliseconds() {
assert_eq!("1600".parse::<UpDuration>().unwrap().as_millis(), 1600);
}
#[test]
fn units_seconds_minutes_hours() {
assert_eq!("30s".parse::<UpDuration>().unwrap().as_millis(), 30_000);
assert_eq!("5m".parse::<UpDuration>().unwrap().as_millis(), 300_000);
assert_eq!("2h".parse::<UpDuration>().unwrap().as_millis(), 7_200_000);
}
#[test]
fn rejects_spec_violations() {
use ParseUpDurationError::*;
assert_eq!("".parse::<UpDuration>(), Err(Empty));
assert_eq!("s".parse::<UpDuration>(), Err(MissingDigits));
assert_eq!("30S".parse::<UpDuration>(), Err(InvalidCharacter)); assert_eq!("1.5s".parse::<UpDuration>(), Err(InvalidCharacter));
assert_eq!("30 s".parse::<UpDuration>(), Err(InvalidCharacter));
assert_eq!("99999999999999999999h".parse::<UpDuration>(), Err(Overflow));
assert_eq!("9999999999999999h".parse::<UpDuration>(), Err(Overflow));
}
#[test]
fn display_round_trips() {
for ms in [
0u64, 1, 999, 1000, 1600, 30_000, 300_000, 7_200_000, 3_601_000,
] {
let d = UpDuration::from_millis(ms);
assert_eq!(d.to_string().parse::<UpDuration>().unwrap(), d, "{ms}ms");
}
assert_eq!(UpDuration::from_millis(30_000).to_string(), "30s");
assert_eq!(UpDuration::from_millis(1600).to_string(), "1600");
assert_eq!(UpDuration::from_millis(7_200_000).to_string(), "2h");
}
#[test]
fn serde_uses_string_form() {
let d: UpDuration = serde_json::from_str("\"30s\"").unwrap();
assert_eq!(d.as_millis(), 30_000);
assert_eq!(serde_json::to_string(&d).unwrap(), "\"30s\"");
}
#[cfg(feature = "schema")]
#[test]
fn the_schema_pattern_agrees_with_from_str() {
let schema = serde_json::to_value(schemars::schema_for!(UpDuration)).unwrap();
let pattern = schema["pattern"].as_str().unwrap();
let re = regex::Regex::new(pattern).unwrap();
for accepted in ["1600", "30s", "5m", "2h"] {
assert!(re.is_match(accepted), "pattern rejects {accepted}");
assert!(
accepted.parse::<UpDuration>().is_ok(),
"FromStr rejects {accepted}"
);
}
for rejected in ["30S", "1.5s", "30 s", "", "s", "30d", "30w"] {
assert!(!re.is_match(rejected), "pattern accepts {rejected}");
assert!(
rejected.parse::<UpDuration>().is_err(),
"FromStr accepts {rejected}"
);
}
}
}