1use core::fmt::Write as _;
9
10use super::Rate;
11
12#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
14#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
15#[cfg_attr(feature = "serde", serde(rename_all = "lowercase"))]
16pub enum ByteUnits {
17 #[default]
19 Iec,
20 Si,
22}
23
24impl ByteUnits {
25 const fn divisor(self) -> u64 {
26 match self {
27 Self::Iec => 1024,
28 Self::Si => 1000,
29 }
30 }
31
32 const fn suffixes(self) -> &'static [&'static str] {
34 match self {
35 Self::Iec => &["B", "KiB", "MiB", "GiB", "TiB", "PiB", "EiB"],
36 Self::Si => &["B", "kB", "MB", "GB", "TB", "PB", "EB"],
37 }
38 }
39
40 const fn short_suffixes(self) -> &'static [&'static str] {
42 match self {
43 Self::Iec | Self::Si => &["B", "K", "M", "G", "T", "P", "E"],
47 }
48 }
49}
50
51fn scale(bytes: u64, units: ByteUnits, max_index: usize) -> (f64, usize) {
53 let divisor = units.divisor();
54 let mut value = bytes as f64;
55 let mut index = 0usize;
56 let divisor_f = divisor as f64;
57 while value >= divisor_f && index < max_index {
58 value /= divisor_f;
59 index += 1;
60 }
61 (value, index)
62}
63
64#[must_use]
69pub fn format_bytes(bytes: u64, units: ByteUnits) -> String {
70 let suffixes = units.suffixes();
71 let (value, index) = scale(bytes, units, suffixes.len().saturating_sub(1));
72 let suffix = suffixes.get(index).copied().unwrap_or("B");
73 let mut out = String::with_capacity(10);
74 if index == 0 {
75 let _ = write!(out, "{bytes} {suffix}");
77 } else if value < 10.0 {
78 let _ = write!(out, "{value:.1} {suffix}");
79 } else {
80 let _ = write!(out, "{value:.0} {suffix}");
81 }
82 out
83}
84
85#[must_use]
89pub fn format_bytes_compact(bytes: u64, units: ByteUnits) -> String {
90 let suffixes = units.short_suffixes();
91 let (value, index) = scale(bytes, units, suffixes.len().saturating_sub(1));
92 let suffix = suffixes.get(index).copied().unwrap_or("B");
93 let mut out = String::with_capacity(6);
94 if index == 0 {
95 let _ = write!(out, "{bytes}{suffix}");
96 } else if value < 10.0 {
97 let _ = write!(out, "{value:.1}{suffix}");
98 } else {
99 let _ = write!(out, "{value:.0}{suffix}");
100 }
101 out
102}
103
104#[must_use]
106pub fn format_byte_rate(rate: Rate, units: ByteUnits) -> String {
107 #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
110 let whole = rate.per_second().min(u64::MAX as f64) as u64;
111 let mut out = format_bytes_compact(whole, units);
112 out.push_str("/s");
113 out
114}
115
116pub const MAX_COMPACT_BYTES_WIDTH: u16 = 5;
121
122pub const MAX_BYTE_RATE_WIDTH: u16 = MAX_COMPACT_BYTES_WIDTH + 2;
124
125#[cfg(test)]
126mod tests {
127 use super::*;
128
129 #[test]
130 fn plain_bytes_have_no_decimal() {
131 assert_eq!(format_bytes(0, ByteUnits::Iec), "0 B");
132 assert_eq!(format_bytes(512, ByteUnits::Iec), "512 B");
133 }
134
135 #[test]
136 fn iec_and_si_boundaries_differ_as_documented() {
137 assert_eq!(format_bytes(1024, ByteUnits::Iec), "1.0 KiB");
138 assert_eq!(format_bytes(1024, ByteUnits::Si), "1.0 kB");
139 assert_eq!(format_bytes(1000, ByteUnits::Si), "1.0 kB");
140 assert_eq!(format_bytes(1000, ByteUnits::Iec), "1000 B");
141 }
142
143 #[test]
144 fn crossing_a_unit_boundary_does_not_widen_the_column() {
145 for bytes in [1023u64, 1024, 1025, 1_048_575, 1_048_576] {
147 let rendered = format_bytes_compact(bytes, ByteUnits::Iec);
148 assert!(
149 rendered.chars().count() <= MAX_COMPACT_BYTES_WIDTH as usize,
150 "{bytes} rendered as {rendered:?}, wider than the reserved width"
151 );
152 }
153 }
154
155 #[test]
156 fn the_largest_counter_still_fits_the_reserved_width() {
157 let rendered = format_bytes_compact(u64::MAX, ByteUnits::Iec);
158 assert!(
159 rendered.chars().count() <= MAX_COMPACT_BYTES_WIDTH as usize,
160 "u64::MAX rendered as {rendered:?}"
161 );
162 assert!(
163 rendered.ends_with('E'),
164 "expected exbibytes, got {rendered:?}"
165 );
166 }
167
168 #[test]
169 fn rates_carry_the_consistent_per_second_suffix() {
170 let rate = Rate::new(42.0 * 1024.0 * 1024.0).expect("valid");
171 assert_eq!(format_byte_rate(rate, ByteUnits::Iec), "42M/s");
172 assert_eq!(format_byte_rate(Rate::ZERO, ByteUnits::Iec), "0B/s");
173 }
174
175 #[test]
176 fn every_byte_rate_fits_the_reserved_width() {
177 for per_second in [0.0, 1.0, 999.0, 1024.0, 1.5e9, 9.9e18] {
178 let rate = Rate::new(per_second).expect("valid");
179 let rendered = format_byte_rate(rate, ByteUnits::Iec);
180 assert!(
181 rendered.chars().count() <= MAX_BYTE_RATE_WIDTH as usize,
182 "{per_second} rendered as {rendered:?}"
183 );
184 }
185 }
186}
187
188#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
190pub enum ByteSizeParseError {
191 #[error("expected a number followed by a unit (B, KiB, MiB, GiB, kB, MB, GB), got {input:?}")]
193 Empty {
194 input: String,
196 },
197 #[error("{value:?} is not a whole number in {input:?}")]
199 NotANumber {
200 input: String,
202 value: String,
204 },
205 #[error(
207 "unknown size unit {unit:?} in {input:?}; expected one of B, KiB, MiB, GiB, kB, MB, GB"
208 )]
209 UnknownUnit {
210 input: String,
212 unit: String,
214 },
215 #[error("size {input:?} is too large")]
217 Overflow {
218 input: String,
220 },
221}
222
223pub fn parse_bytes(input: &str) -> Result<u64, ByteSizeParseError> {
236 let trimmed = input.trim();
237 let split = trimmed
238 .char_indices()
239 .find(|(_, c)| !c.is_ascii_digit() && *c != '_' && *c != '.')
240 .map_or(trimmed.len(), |(index, _)| index);
241 let (number, unit) = trimmed.split_at(split);
242 let number: String = number.chars().filter(|c| *c != '_').collect();
243
244 let (whole_text, fraction_text) = match number.split_once('.') {
245 Some((whole, fraction)) => (whole, fraction),
246 None => (number.as_str(), ""),
247 };
248 if whole_text.is_empty() || fraction_text.contains('.') {
249 return Err(ByteSizeParseError::Empty {
250 input: trimmed.to_owned(),
251 });
252 }
253 let amount: u64 = whole_text
254 .parse()
255 .map_err(|_| ByteSizeParseError::NotANumber {
256 input: trimmed.to_owned(),
257 value: number.clone(),
258 })?;
259 let fraction: u64 = if fraction_text.is_empty() {
260 0
261 } else {
262 fraction_text
263 .parse()
264 .map_err(|_| ByteSizeParseError::NotANumber {
265 input: trimmed.to_owned(),
266 value: number.clone(),
267 })?
268 };
269 let fraction_scale = 10_u64
270 .checked_pow(u32::try_from(fraction_text.len()).unwrap_or(u32::MAX))
271 .ok_or(ByteSizeParseError::Overflow {
272 input: trimmed.to_owned(),
273 })?;
274
275 let unit = unit.trim();
276 let multiplier: u64 = match unit {
277 "" | "B" | "b" => 1,
278 "KiB" | "kib" | "KIB" | "K" | "k" => 1024,
279 "MiB" | "mib" | "MIB" | "M" | "m" => 1024 * 1024,
280 "GiB" | "gib" | "GIB" | "G" | "g" => 1024 * 1024 * 1024,
281 "TiB" | "tib" | "TIB" | "T" | "t" => 1024_u64.pow(4),
282 "kB" | "KB" | "kb" => 1_000,
283 "MB" | "mB" | "mb" => 1_000_000,
284 "GB" | "gB" | "gb" => 1_000_000_000,
285 "TB" | "tB" | "tb" => 1_000_000_000_000,
286 other => {
287 return Err(ByteSizeParseError::UnknownUnit {
288 input: trimmed.to_owned(),
289 unit: other.to_owned(),
290 });
291 }
292 };
293
294 let total = u128::from(amount) * u128::from(multiplier)
297 + u128::from(fraction) * u128::from(multiplier) / u128::from(fraction_scale);
298 u64::try_from(total).map_err(|_| ByteSizeParseError::Overflow {
299 input: trimmed.to_owned(),
300 })
301}
302
303#[cfg(test)]
304mod parse_tests {
305 use super::*;
306
307 #[test]
308 fn parses_both_unit_families() {
309 assert_eq!(parse_bytes("32MiB"), Ok(32 * 1024 * 1024));
310 assert_eq!(parse_bytes("32MB"), Ok(32_000_000));
311 assert_eq!(parse_bytes("1KiB"), Ok(1024));
312 assert_eq!(parse_bytes("1kB"), Ok(1_000));
313 assert_eq!(parse_bytes("1GiB"), Ok(1024 * 1024 * 1024));
314 }
315
316 #[test]
317 fn a_single_decimal_fraction_is_exact_integer_arithmetic() {
318 assert_eq!(parse_bytes("1.5GiB"), Ok(1_610_612_736));
319 assert_eq!(parse_bytes("1.0KiB"), Ok(1024));
320 assert_eq!(parse_bytes("0.5MiB"), Ok(512 * 1024));
321 assert_eq!(parse_bytes("2.25GiB"), Ok(2_415_919_104));
322 assert_eq!(parse_bytes("10.9"), Ok(10));
324 }
325
326 #[test]
327 fn a_malformed_fraction_is_rejected() {
328 assert!(parse_bytes("1.2.3MiB").is_err());
329 assert!(matches!(
330 parse_bytes(".5MiB"),
331 Err(ByteSizeParseError::Empty { .. })
332 ));
333 }
334
335 #[test]
336 fn a_bare_number_is_bytes() {
337 assert_eq!(parse_bytes("1024"), Ok(1024));
338 assert_eq!(parse_bytes("0"), Ok(0));
339 }
340
341 #[test]
342 fn whitespace_and_underscores_are_tolerated() {
343 assert_eq!(parse_bytes(" 32 MiB "), Ok(32 * 1024 * 1024));
344 assert_eq!(parse_bytes("1_048_576"), Ok(1_048_576));
345 }
346
347 #[test]
348 fn the_i_distinguishes_the_families_case_insensitively() {
349 assert_eq!(parse_bytes("2MiB"), Ok(2 * 1024 * 1024));
350 assert_eq!(parse_bytes("2mib"), Ok(2 * 1024 * 1024));
351 assert_eq!(parse_bytes("2MB"), Ok(2_000_000));
352 assert_eq!(parse_bytes("2mb"), Ok(2_000_000));
353 }
354
355 #[test]
356 fn errors_quote_the_offending_input() {
357 let error = parse_bytes("32 gigglebytes").expect_err("bad unit");
358 assert!(error.to_string().contains("gigglebytes"), "{error}");
359 assert!(matches!(
360 parse_bytes("MiB"),
361 Err(ByteSizeParseError::Empty { .. })
362 ));
363 assert!(matches!(
364 parse_bytes("99999999999999999999999"),
365 Err(ByteSizeParseError::NotANumber { .. })
366 ));
367 assert!(matches!(
368 parse_bytes("99999999999999999TiB"),
369 Err(ByteSizeParseError::Overflow { .. })
370 ));
371 }
372
373 #[test]
374 fn formatting_round_trips_through_parsing_for_iec_sizes() {
375 for bytes in [0u64, 512, 1024, 32 * 1024 * 1024, 4 * 1024 * 1024 * 1024] {
376 let rendered = format_bytes(bytes, ByteUnits::Iec).replace(' ', "");
377 let reparsed = parse_bytes(&rendered).expect("re-parse");
378 let drift = reparsed.abs_diff(bytes);
381 assert!(
382 drift <= bytes / 100 + 1,
383 "{bytes} rendered as {rendered} reparsed as {reparsed}"
384 );
385 }
386 }
387}