1use core::fmt::Write;
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11#[repr(u8)]
12pub enum SizeUnit {
13 Auto = 0, Bytes = b'B',
15 KiB = b'K',
16 MiB = b'M',
17 GiB = b'G',
18 TiB = b'T',
19 PiB = b'P',
20 EiB = b'E',
21}
22
23impl TryFrom<u8> for SizeUnit {
24 type Error = ();
25 fn try_from(value: u8) -> Result<Self, Self::Error> {
26 match value.to_ascii_uppercase() {
27 b'B' => Ok(SizeUnit::Bytes),
28 b'K' => Ok(SizeUnit::KiB),
29 b'M' => Ok(SizeUnit::MiB),
30 b'G' => Ok(SizeUnit::GiB),
31 b'T' => Ok(SizeUnit::TiB),
32 b'P' => Ok(SizeUnit::PiB),
33 b'E' => Ok(SizeUnit::EiB),
34 0 => Ok(SizeUnit::Auto),
35 _ => Err(()),
36 }
37 }
38}
39
40impl TryFrom<char> for SizeUnit {
41 type Error = ();
42 fn try_from(value: char) -> Result<Self, Self::Error> {
43 if value.is_ascii() { Self::try_from(value as u8) } else { Err(()) }
44 }
45}
46
47impl SizeUnit {
48 pub fn to_str(self) -> &'static str {
49 match self {
50 SizeUnit::Auto => "",
51 SizeUnit::Bytes => "B",
52 SizeUnit::KiB => "K",
53 SizeUnit::MiB => "M",
54 SizeUnit::GiB => "G",
55 SizeUnit::TiB => "T",
56 SizeUnit::PiB => "P",
57 SizeUnit::EiB => "E",
58 }
59 }
60}
61
62pub const MAX_FORMAT_SIZE_LEN: usize = 22;
66
67struct SliceWriter<'a> {
68 slice: &'a mut [u8],
69 cursor: usize,
70}
71
72impl<'a> SliceWriter<'a> {
73 fn new(slice: &'a mut [u8]) -> Self {
74 Self { slice, cursor: 0 }
75 }
76
77 fn into_str(self) -> &'a str {
78 core::str::from_utf8(&self.slice[..self.cursor]).unwrap_or("")
79 }
80}
81
82impl<'a> core::fmt::Write for SliceWriter<'a> {
83 fn write_str(&mut self, s: &str) -> core::fmt::Result {
84 let bytes = s.as_bytes();
85 let keep = core::cmp::min(bytes.len(), self.slice.len() - self.cursor);
86 self.slice[self.cursor..self.cursor + keep].copy_from_slice(&bytes[..keep]);
87 self.cursor += keep;
88 Ok(())
89 }
90}
91
92fn rounding_divide(n: usize, d: usize) -> usize {
97 let round_up = (n % d) >= (d / 2);
100 n / d + if round_up { 1 } else { 0 }
101}
102
103pub fn format_size_fixed_rs(buf: &mut [u8], bytes: usize, mut unit: u8) -> &str {
118 if buf.is_empty() {
119 return "";
120 }
121
122 let mut writer = SliceWriter::new(buf);
123 let units = b"BKMGTPE";
124 let num_units = units.len();
125
126 let orig_bytes = bytes;
127 let mut prepended_question = false;
128
129 let mut bytes = bytes;
130
131 loop {
132 let mut ui = 0;
133 let mut divisor = 1;
134
135 while if unit != 0 {
145 ui < num_units && units[ui] != unit
146 } else {
147 bytes >= 10000 || (bytes != 0 && (bytes & 1023) == 0)
148 } {
149 ui += 1;
150 if ui >= num_units {
151 if !prepended_question {
154 let _ = writer.write_char('?');
155 prepended_question = true;
156 }
157 unit = 0;
158 bytes = orig_bytes;
159 break;
160 }
161 bytes /= 1024;
162 divisor *= 1024;
163 }
164
165 if ui < num_units {
166 if orig_bytes % divisor == 0 {
169 let _ = core::write!(&mut writer, "{}{}", bytes, units[ui] as char);
170 } else {
171 let mut int_part = orig_bytes / divisor;
180 let mut fractional_part = rounding_divide((orig_bytes % divisor) * 10, divisor);
181 if fractional_part >= 10 {
182 fractional_part = 0;
184 int_part += 1;
185 }
186 let _ = core::write!(
187 &mut writer,
188 "{}.{}{}",
189 int_part,
190 fractional_part,
191 units[ui] as char
192 );
193 }
194 break;
195 }
196 }
197
198 let s = writer.into_str();
199 assert!(s.len() <= MAX_FORMAT_SIZE_LEN);
200 s
201}
202
203pub fn format_size_rs(buf: &mut [u8], bytes: usize) -> &str {
205 format_size_fixed_rs(buf, bytes, 0)
206}
207
208#[derive(Debug, Clone, Copy, PartialEq, Eq)]
209struct EncodedSize<'a> {
210 integral: &'a str,
212
213 fractional: Option<&'a str>,
215
216 unit: SizeUnit,
217
218 scale: u64,
219}
220
221fn process_formatted_string(mut formatted: &str) -> Option<EncodedSize<'_>> {
222 if formatted.is_empty() {
223 return None;
224 }
225
226 let mut unit = SizeUnit::Bytes;
227 let mut scale = 1u64;
228
229 if let Some(last_char) = formatted.chars().next_back() {
230 if !last_char.is_ascii_digit() {
231 unit = SizeUnit::try_from(last_char).ok()?;
232 formatted = &formatted[..formatted.len() - last_char.len_utf8()];
233
234 match unit {
236 SizeUnit::Bytes => scale = 1,
237 SizeUnit::KiB => scale = 1 << 10,
238 SizeUnit::MiB => scale = 1 << 20,
239 SizeUnit::GiB => scale = 1 << 30,
240 SizeUnit::TiB => scale = 1 << 40,
241 SizeUnit::PiB => scale = 1 << 50,
242 SizeUnit::EiB => scale = 1 << 60,
243 _ => return None,
244 }
245 }
246 }
247
248 let split_at = formatted.find('.');
249 let (integral, fractional) = if let Some(split_at) = split_at {
250 let integral = &formatted[..split_at];
251 let fractional = &formatted[split_at + 1..];
252 if fractional.is_empty() {
254 return None;
255 }
256 (integral, Some(fractional))
257 } else {
258 (formatted, None)
259 };
260
261 if integral.is_empty() {
262 return None;
263 }
264
265 Some(EncodedSize { integral, fractional, unit, scale })
266}
267
268pub fn parse_size_bytes(formatted_bytes: &str) -> Option<u64> {
276 let encoded_size = process_formatted_string(formatted_bytes)?;
277
278 let mut integral: u64 = 0;
279 let mut base_10: u64 = 1;
280
281 for c in encoded_size.integral.chars().rev() {
282 if !c.is_ascii_digit() {
283 return None;
284 }
285 let val = c.to_digit(10)? as u64;
286 let scaled_val = val.checked_mul(base_10)?.checked_mul(encoded_size.scale)?;
287
288 integral = integral.checked_add(scaled_val)?;
289 base_10 = base_10.checked_mul(10)?;
290 }
291
292 let mut fractional: u64 = 0;
293 if let Some(frac_str) = encoded_size.fractional {
294 let mut frac_base_10: u64 = 1;
295 let mut carry: u64 = 0;
296
297 for c in frac_str.chars() {
309 if !c.is_ascii_digit() {
310 return None;
311 }
312 let val = c.to_digit(10)? as u64;
313 frac_base_10 = frac_base_10.checked_mul(10)?;
314 let scaled_value = val.checked_mul(encoded_size.scale)?;
315
316 let contrib = scaled_value / frac_base_10;
318 fractional = fractional.checked_add(contrib)?;
319
320 carry = carry.checked_mul(10)?.checked_add(scaled_value % frac_base_10)?;
322
323 let consumed_carry = carry / frac_base_10;
325 fractional = fractional.checked_add(consumed_carry)?;
326
327 carry %= frac_base_10;
329 }
330
331 if carry != 0 {
335 return None;
336 }
337 }
338
339 integral.checked_add(fractional)
340}
341
342#[cfg(test)]
343mod tests {
344 use super::*;
345
346 const KILO: usize = 1024;
347 const MEGA: usize = KILO * 1024;
348 const GIGA: usize = MEGA * 1024;
349 const TERA: usize = GIGA * 1024;
350 const PETA: usize = TERA * 1024;
351 const EXA: usize = PETA * 1024;
352
353 struct FormatSizeTestCase {
354 input: usize,
355 unit: u8,
356 expected_output: &'static str,
357 }
358
359 const FORMAT_SIZE_TEST_CASES: &[FormatSizeTestCase] = &[
360 FormatSizeTestCase { input: 0, unit: 0, expected_output: "0B" },
363 FormatSizeTestCase { input: 1, unit: 0, expected_output: "1B" },
364 FormatSizeTestCase { input: KILO - 1, unit: 0, expected_output: "1023B" },
369 FormatSizeTestCase { input: KILO, unit: 0, expected_output: "1K" },
370 FormatSizeTestCase { input: KILO + 1, unit: 0, expected_output: "1025B" },
371 FormatSizeTestCase { input: KILO * 9, unit: 0, expected_output: "9K" },
372 FormatSizeTestCase { input: KILO * 9 + 1, unit: 0, expected_output: "9217B" },
373 FormatSizeTestCase { input: KILO * 10, unit: 0, expected_output: "10K" },
374 FormatSizeTestCase { input: MEGA - KILO, unit: 0, expected_output: "1023K" },
376 FormatSizeTestCase { input: MEGA, unit: 0, expected_output: "1M" },
377 FormatSizeTestCase { input: MEGA + KILO, unit: 0, expected_output: "1025K" },
378 FormatSizeTestCase { input: MEGA * 9, unit: 0, expected_output: "9M" },
379 FormatSizeTestCase { input: MEGA * 9 + KILO, unit: 0, expected_output: "9217K" },
380 FormatSizeTestCase { input: MEGA * 10, unit: 0, expected_output: "10M" },
381 FormatSizeTestCase { input: MEGA, unit: 0, expected_output: "1M" },
383 FormatSizeTestCase { input: GIGA, unit: 0, expected_output: "1G" },
384 FormatSizeTestCase { input: TERA, unit: 0, expected_output: "1T" },
385 FormatSizeTestCase { input: PETA, unit: 0, expected_output: "1P" },
386 FormatSizeTestCase { input: EXA, unit: 0, expected_output: "1E" },
387 FormatSizeTestCase { input: MEGA - 1, unit: 0, expected_output: "1024.0K" },
390 FormatSizeTestCase { input: MEGA + MEGA / 3, unit: 0, expected_output: "1365.3K" },
391 FormatSizeTestCase { input: GIGA - 1, unit: 0, expected_output: "1024.0M" },
392 FormatSizeTestCase { input: TERA - 1, unit: 0, expected_output: "1024.0G" },
393 FormatSizeTestCase { input: PETA - 1, unit: 0, expected_output: "1024.0T" },
394 FormatSizeTestCase { input: EXA - 1, unit: 0, expected_output: "1024.0P" },
395 FormatSizeTestCase { input: usize::MAX, unit: 0, expected_output: "16.0E" },
396 FormatSizeTestCase { input: 9999, unit: 0, expected_output: "9999B" },
399 FormatSizeTestCase { input: 10000, unit: 0, expected_output: "9.8K" },
400 FormatSizeTestCase { input: KILO * 9999, unit: 0, expected_output: "9999K" },
401 FormatSizeTestCase { input: KILO * 9999 + 1, unit: 0, expected_output: "9999.0K" },
402 FormatSizeTestCase { input: KILO * 10000, unit: 0, expected_output: "9.8M" },
403 FormatSizeTestCase { input: 10700, unit: 0, expected_output: "10.4K" },
405 FormatSizeTestCase { input: 10701, unit: 0, expected_output: "10.5K" },
406 FormatSizeTestCase { input: 69887590, unit: 0, expected_output: "66.6M" },
407 FormatSizeTestCase { input: 69887591, unit: 0, expected_output: "66.7M" },
408 FormatSizeTestCase { input: 18389097998479209267, unit: 0, expected_output: "15.9E" },
409 FormatSizeTestCase { input: 18389097998479209268, unit: 0, expected_output: "16.0E" },
410 FormatSizeTestCase {
412 input: usize::MAX,
413 unit: b'B',
414 expected_output: "18446744073709551615B",
415 },
416 FormatSizeTestCase {
417 input: usize::MAX,
418 unit: b'K',
419 expected_output: "18014398509481984.0K",
420 },
421 FormatSizeTestCase { input: usize::MAX, unit: b'M', expected_output: "17592186044416.0M" },
422 FormatSizeTestCase { input: usize::MAX, unit: b'G', expected_output: "17179869184.0G" },
423 FormatSizeTestCase { input: usize::MAX, unit: b'T', expected_output: "16777216.0T" },
424 FormatSizeTestCase { input: usize::MAX, unit: b'P', expected_output: "16384.0P" },
425 FormatSizeTestCase { input: usize::MAX, unit: b'E', expected_output: "16.0E" },
426 FormatSizeTestCase { input: GIGA, unit: b'K', expected_output: "1048576K" },
428 FormatSizeTestCase { input: MEGA / 10, unit: b'M', expected_output: "0.1M" },
430 FormatSizeTestCase { input: GIGA, unit: b'q', expected_output: "?1G" },
432 FormatSizeTestCase { input: KILO, unit: b'q', expected_output: "?1K" },
433 FormatSizeTestCase { input: GIGA + 1, unit: b'#', expected_output: "?1.0G" },
434 FormatSizeTestCase { input: KILO + 1, unit: b'#', expected_output: "?1025B" },
435 ];
436
437 #[test]
438 fn test_format_size_fixed() {
439 let mut buf = [0u8; MAX_FORMAT_SIZE_LEN];
440 for (i, tc) in FORMAT_SIZE_TEST_CASES.iter().enumerate() {
441 buf.fill(0);
442 let res = format_size_fixed_rs(&mut buf, tc.input, tc.unit);
443 assert_eq!(
444 res, tc.expected_output,
445 "case {}, input={}, unit={}",
446 i, tc.input, tc.unit as char
447 );
448 }
449 }
450
451 #[test]
452 fn test_format_size_short_buf_truncates() {
453 let input = 1023 * KILO + 1;
454 let expected_output = "1023.0K";
455
456 let mut buf = [0x55u8; MAX_FORMAT_SIZE_LEN * 2];
457 for str_size in 0..=expected_output.len() {
458 buf.fill(0x55);
459 let res = format_size_rs(&mut buf[..str_size], input);
460 assert_eq!(res, &expected_output[..str_size]);
461 assert_eq!(buf[str_size], 0x55);
462 }
463 }
464
465 #[test]
466 fn test_format_size_bad_unit_short_buf_truncates() {
467 let mut buf = [0x55u8; MAX_FORMAT_SIZE_LEN];
468
469 buf.fill(0x55);
471 let res = format_size_fixed_rs(&mut buf[..0], GIGA, b'q');
472 assert_eq!(res, "");
473 assert_eq!(buf[0], 0x55);
474
475 buf.fill(0x55);
477 let res = format_size_fixed_rs(&mut buf[..1], GIGA, b'q');
478 assert_eq!(res, "?");
479 assert_eq!(buf[1], 0x55);
480
481 buf.fill(0x55);
483 let res = format_size_fixed_rs(&mut buf[..2], GIGA, b'q');
484 assert_eq!(res, "?1");
485 assert_eq!(buf[2], 0x55);
486
487 buf.fill(0x55);
489 let res = format_size_fixed_rs(&mut buf[..3], GIGA, b'q');
490 assert_eq!(res, "?1G");
491 assert_eq!(buf[3], 0x55);
492 }
493
494 #[test]
495 fn test_cpp_to_string() {
496 assert_eq!(SizeUnit::Auto.to_str(), "");
497 assert_eq!(SizeUnit::Bytes.to_str(), "B");
498 assert_eq!(SizeUnit::KiB.to_str(), "K");
499 assert_eq!(SizeUnit::MiB.to_str(), "M");
500 assert_eq!(SizeUnit::GiB.to_str(), "G");
501 assert_eq!(SizeUnit::TiB.to_str(), "T");
502 assert_eq!(SizeUnit::PiB.to_str(), "P");
503 assert_eq!(SizeUnit::EiB.to_str(), "E");
504 }
505
506 struct ParseTestCase {
507 expected_bytes: u64,
508 input: &'static str,
509 }
510
511 const PARSE_TEST_CASES: &[ParseTestCase] = &[
512 ParseTestCase { expected_bytes: 1234, input: "1234" },
514 ParseTestCase { expected_bytes: 1234, input: "1234b" },
515 ParseTestCase { expected_bytes: 1234, input: "1234B" },
516 ParseTestCase { expected_bytes: 1234 * 1024, input: "1234k" },
517 ParseTestCase { expected_bytes: 1234 * 1024, input: "1234K" },
518 ParseTestCase { expected_bytes: 1234 * 1024 * 1024, input: "1234m" },
519 ParseTestCase { expected_bytes: 1234 * 1024 * 1024, input: "1234M" },
520 ParseTestCase { expected_bytes: 1234 * 1024 * 1024 * 1024, input: "1234g" },
521 ParseTestCase { expected_bytes: 1234 * 1024 * 1024 * 1024, input: "1234G" },
522 ParseTestCase { expected_bytes: 1234 * 1024 * 1024 * 1024 * 1024, input: "1234t" },
523 ParseTestCase { expected_bytes: 1234 * 1024 * 1024 * 1024 * 1024, input: "1234T" },
524 ParseTestCase { expected_bytes: 5 * 1024 * 1024 * 1024 * 1024 * 1024, input: "5p" },
525 ParseTestCase { expected_bytes: 5 * 1024 * 1024 * 1024 * 1024 * 1024, input: "5P" },
526 ParseTestCase { expected_bytes: 2 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024, input: "2e" },
527 ParseTestCase { expected_bytes: 2 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024, input: "2E" },
528 ParseTestCase { expected_bytes: 10700, input: "10.4492187500k" },
530 ParseTestCase { expected_bytes: 10700, input: "10.4492187500K" },
531 ParseTestCase { expected_bytes: 10700 * 1024, input: "10.4492187500m" },
532 ParseTestCase { expected_bytes: 10700 * 1024, input: "10.4492187500M" },
533 ParseTestCase { expected_bytes: 10700 * 1024 * 1024, input: "10.4492187500g" },
534 ParseTestCase { expected_bytes: 10700 * 1024 * 1024, input: "10.4492187500G" },
535 ParseTestCase { expected_bytes: 10700 * 1024 * 1024 * 1024, input: "10.4492187500t" },
536 ParseTestCase { expected_bytes: 10700 * 1024 * 1024 * 1024, input: "10.4492187500T" },
537 ParseTestCase {
538 expected_bytes: 10700 * 1024 * 1024 * 1024 * 1024,
539 input: "10.4492187500p",
540 },
541 ParseTestCase {
542 expected_bytes: 10700 * 1024 * 1024 * 1024 * 1024,
543 input: "10.4492187500P",
544 },
545 ParseTestCase { expected_bytes: 1441151880758558720, input: "1.25e" },
546 ParseTestCase { expected_bytes: 1441151880758558720, input: "1.25E" },
547 ];
548
549 #[test]
550 fn test_parse_size_bytes() {
551 for tc in PARSE_TEST_CASES {
552 let res = parse_size_bytes(tc.input);
553 assert_eq!(res, Some(tc.expected_bytes), "input: {}", tc.input);
554 }
555 }
556
557 const INVALID_INPUTS: &[&str] = &["", "1..1", "1w", "b", "AM", "1.AM", "A.1M"];
558
559 #[test]
560 fn test_parse_size_bytes_invalid() {
561 for input in INVALID_INPUTS {
562 let res = parse_size_bytes(input);
563 assert_eq!(res, None, "input: {}", input);
564 }
565 }
566}