1use std::time::Duration;
6
7const SIZE_UNITS: &[(&str, u64)] = &[
10 ("B", 1),
11 ("KiB", 1 << 10),
12 ("MiB", 1 << 20),
13 ("GiB", 1 << 30),
14 ("TiB", 1 << 40),
15 ("kB", 1_000),
16 ("KB", 1_000),
17 ("MB", 1_000_000),
18 ("GB", 1_000_000_000),
19 ("TB", 1_000_000_000_000),
20];
21
22const DURATION_UNITS: &[(&str, u64)] = &[("ms", 1), ("s", 1_000), ("m", 60_000), ("h", 3_600_000)];
23
24fn split(
26 text: &str,
27 units: &[(&str, u64)],
28 what: &str,
29 example: &str,
30) -> Result<(f64, u64), String> {
31 let text = text.trim();
32 let at = text
33 .find(|c: char| !(c.is_ascii_digit() || c == '.' || c == '_'))
34 .unwrap_or(text.len());
35 let (number, unit) = text.split_at(at);
36 let number = number.replace('_', "");
37 let value: f64 = number
38 .parse()
39 .map_err(|_| format!("\"{text}\" is not a {what}, such as {example}"))?;
40 let unit = unit.trim();
41 if unit.is_empty() {
42 if value == 0.0 {
43 return Ok((0.0, 1));
44 }
45 return Err(format!("\"{text}\" needs a unit, as in {example}"));
46 }
47 let factor = units
48 .iter()
49 .find(|(name, _)| *name == unit)
50 .map(|(_, f)| *f)
51 .ok_or_else(|| {
52 let names: Vec<&str> = units.iter().map(|(n, _)| *n).collect();
53 format!("\"{unit}\" is not a unit of {what}: {}", names.join(", "))
54 })?;
55 Ok((value, factor))
56}
57
58pub fn parse_size(text: &str) -> Result<u64, String> {
60 let (value, factor) = split(text, SIZE_UNITS, "size", "512MiB")?;
61 let bytes = value * factor as f64;
62 if !bytes.is_finite() || bytes > u64::MAX as f64 {
63 return Err(format!("\"{}\" is too large", text.trim()));
64 }
65 Ok(bytes.round() as u64)
66}
67
68pub fn parse_duration(text: &str) -> Result<Duration, String> {
70 let (value, factor) = split(text, DURATION_UNITS, "duration", "250ms")?;
71 let ms = value * factor as f64;
72 if !ms.is_finite() || ms > u64::MAX as f64 {
73 return Err(format!("\"{}\" is too long", text.trim()));
74 }
75 Ok(Duration::from_millis(ms.round() as u64))
76}
77
78pub fn format_size(bytes: u64) -> String {
80 if bytes == 0 {
81 return "0".to_string();
82 }
83 let (name, factor) = [
84 ("TiB", 1u64 << 40),
85 ("GiB", 1 << 30),
86 ("MiB", 1 << 20),
87 ("KiB", 1 << 10),
88 ]
89 .into_iter()
90 .find(|(_, f)| bytes.is_multiple_of(*f))
91 .unwrap_or(("B", 1));
92 format!("{}{name}", bytes / factor)
93}
94
95pub fn format_duration(duration: Duration) -> String {
97 let ms = u64::try_from(duration.as_millis()).unwrap_or(u64::MAX);
98 if ms == 0 {
99 return "0".to_string();
100 }
101 let (name, factor) = DURATION_UNITS
102 .iter()
103 .rev()
104 .find(|(_, f)| ms.is_multiple_of(*f))
105 .copied()
106 .unwrap_or(("ms", 1));
107 format!("{}{name}", ms / factor)
108}
109
110#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)]
112pub struct ByteSize(pub u64);
113
114impl ByteSize {
115 pub const fn mib(n: u64) -> Self {
116 Self(n << 20)
117 }
118 pub const fn kib(n: u64) -> Self {
119 Self(n << 10)
120 }
121 pub fn bytes(self) -> u64 {
122 self.0
123 }
124}
125
126impl std::fmt::Display for ByteSize {
127 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
128 f.write_str(&format_size(self.0))
129 }
130}
131
132impl std::str::FromStr for ByteSize {
133 type Err = String;
134 fn from_str(s: &str) -> Result<Self, String> {
135 parse_size(s).map(Self)
136 }
137}
138
139#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)]
141pub struct Interval(pub Duration);
142
143impl Interval {
144 pub const fn ms(n: u64) -> Self {
145 Self(Duration::from_millis(n))
146 }
147 pub fn duration(self) -> Duration {
148 self.0
149 }
150}
151
152impl std::fmt::Display for Interval {
153 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
154 f.write_str(&format_duration(self.0))
155 }
156}
157
158impl std::str::FromStr for Interval {
159 type Err = String;
160 fn from_str(s: &str) -> Result<Self, String> {
161 parse_duration(s).map(Self)
162 }
163}
164
165macro_rules! serde_as_text {
166 ($ty:ty) => {
167 impl serde::Serialize for $ty {
168 fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
169 s.serialize_str(&self.to_string())
170 }
171 }
172
173 impl<'de> serde::Deserialize<'de> for $ty {
174 fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
175 struct Visit;
176 impl serde::de::Visitor<'_> for Visit {
177 type Value = $ty;
178 fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
179 f.write_str("a value with its unit, in quotes")
180 }
181 fn visit_str<E: serde::de::Error>(self, v: &str) -> Result<$ty, E> {
182 v.parse().map_err(E::custom)
183 }
184 fn visit_i64<E: serde::de::Error>(self, v: i64) -> Result<$ty, E> {
186 self.visit_str(&v.to_string())
187 }
188 fn visit_u64<E: serde::de::Error>(self, v: u64) -> Result<$ty, E> {
189 self.visit_str(&v.to_string())
190 }
191 }
192 d.deserialize_any(Visit)
193 }
194 }
195 };
196}
197
198serde_as_text!(ByteSize);
199serde_as_text!(Interval);
200
201#[cfg(test)]
202mod tests {
203 use super::*;
204
205 #[test]
206 fn sizes_read_and_write_with_their_unit() {
207 assert_eq!(parse_size("512MiB"), Ok(512 << 20));
208 assert_eq!(parse_size(" 2GiB "), Ok(2 << 30));
209 assert_eq!(parse_size("1.5KiB"), Ok(1536));
210 assert_eq!(parse_size("100KB"), Ok(100_000));
211 assert_eq!(parse_size("0"), Ok(0));
212 assert!(parse_size("512").unwrap_err().contains("needs a unit"));
213 assert!(parse_size("5mb").unwrap_err().contains("not a unit"));
214 assert!(parse_size("MiB").is_err());
215 assert_eq!(format_size(512 << 20), "512MiB");
216 assert_eq!(format_size(2048 << 20), "2GiB");
217 assert_eq!(format_size(100 << 10), "100KiB");
218 assert_eq!(format_size(1000), "1000B");
219 assert_eq!(format_size(0), "0");
220 }
221
222 #[test]
223 fn durations_read_and_write_with_their_unit() {
224 assert_eq!(parse_duration("250ms"), Ok(Duration::from_millis(250)));
225 assert_eq!(parse_duration("1.5s"), Ok(Duration::from_millis(1500)));
226 assert_eq!(parse_duration("2m"), Ok(Duration::from_secs(120)));
227 assert_eq!(parse_duration("0"), Ok(Duration::ZERO));
228 assert!(parse_duration("250").unwrap_err().contains("needs a unit"));
229 assert_eq!(format_duration(Duration::from_millis(1500)), "1500ms");
230 assert_eq!(format_duration(Duration::from_secs(2)), "2s");
231 assert_eq!(format_duration(Duration::from_secs(60)), "1m");
232 }
233}