Skip to main content

disktest_lib/
util.rs

1// -*- coding: utf-8 -*-
2//
3// disktest - Storage tester
4//
5// Copyright 2020-2026 Michael Büsch <m@bues.ch>
6//
7// Licensed under the Apache License version 2.0
8// or the MIT license, at your option.
9// SPDX-License-Identifier: Apache-2.0 OR MIT
10//
11
12use anyhow as ah;
13use std::fmt::Write as _;
14use std::time::Duration;
15
16const EIB: u64 = 1024 * 1024 * 1024 * 1024 * 1024 * 1024;
17const PIB: u64 = 1024 * 1024 * 1024 * 1024 * 1024;
18const TIB: u64 = 1024 * 1024 * 1024 * 1024;
19const GIB: u64 = 1024 * 1024 * 1024;
20const MIB: u64 = 1024 * 1024;
21const KIB: u64 = 1024;
22
23const EB: u64 = 1000 * 1000 * 1000 * 1000 * 1000 * 1000;
24const PB: u64 = 1000 * 1000 * 1000 * 1000 * 1000;
25const TB: u64 = 1000 * 1000 * 1000 * 1000;
26const GB: u64 = 1000 * 1000 * 1000;
27const MB: u64 = 1000 * 1000;
28const KB: u64 = 1000;
29
30#[allow(clippy::cast_precision_loss)]
31pub fn prettybytes(count: u64, binary: bool, decimal: bool, bytes: bool) -> String {
32    let mut ret = String::new();
33
34    if !binary && !decimal {
35        return ret;
36    }
37
38    if count < KIB {
39        let _ = write!(ret, "{count} bytes");
40        return ret;
41    }
42
43    if binary {
44        let _ = match count {
45            EIB..=u64::MAX => write!(ret, "{:.4} EiB", ((count / TIB) as f64) / (MIB as f64)),
46            PIB..EIB => write!(ret, "{:.4} PiB", ((count / GIB) as f64) / (MIB as f64)),
47            TIB..PIB => write!(ret, "{:.4} TiB", ((count / MIB) as f64) / (MIB as f64)),
48            GIB..TIB => write!(ret, "{:.2} GiB", ((count / MIB) as f64) / (KIB as f64)),
49            MIB..GIB => write!(ret, "{:.1} MiB", (count as f64) / (MIB as f64)),
50            0..MIB => write!(ret, "{:.1} kiB", (count as f64) / (KIB as f64)),
51        };
52    }
53
54    let paren = if !ret.is_empty() && (decimal || bytes) {
55        ret.push_str(" (");
56        true
57    } else {
58        false
59    };
60
61    if decimal {
62        let _ = match count {
63            EB..=u64::MAX => write!(ret, "{:.4} EB", ((count / TB) as f64) / (MB as f64)),
64            PB..EB => write!(ret, "{:.4} PB", ((count / GB) as f64) / (MB as f64)),
65            TB..PB => write!(ret, "{:.4} TB", ((count / MB) as f64) / (MB as f64)),
66            GB..TB => write!(ret, "{:.2} GB", ((count / MB) as f64) / (KB as f64)),
67            MB..GB => write!(ret, "{:.1} MB", (count as f64) / (MB as f64)),
68            0..MB => write!(ret, "{:.1} kB", (count as f64) / (KB as f64)),
69        };
70    }
71
72    if bytes {
73        if decimal {
74            ret.push_str(", ");
75        }
76        let _ = write!(ret, "{count} bytes");
77    }
78
79    if paren {
80        ret.push(')');
81    }
82
83    ret
84}
85
86#[allow(clippy::cast_possible_truncation)]
87#[allow(clippy::cast_sign_loss)]
88fn try_one_parsebytes(s: &str, suffix: &str, factor: u64) -> ah::Result<u64> {
89    let Some(s) = s.strip_suffix(suffix) else {
90        return Err(ah::format_err!("Value suffix does not match."));
91    };
92    let s = s.trim();
93    if let Ok(value) = s.parse::<u64>() {
94        // Integer value.
95        let Some(prod) = value.checked_mul(factor) else {
96            return Err(ah::format_err!("Value integer overflow."));
97        };
98        Ok(prod)
99    } else if let Ok(value) = s.parse::<f64>() {
100        // Floating point value.
101        let factor = factor as f64;
102        if value.log2() + factor.log2() >= 61.0 {
103            return Err(ah::format_err!("Value float overflow."));
104        }
105        let value = (value * factor).round().max(0.0);
106        Ok(value as u64)
107    } else {
108        Err(ah::format_err!("Value is neither integer nor float."))
109    }
110}
111
112/// Parse a string describing a byte size into a number.
113pub fn parsebytes(s: &str) -> ah::Result<u64> {
114    let s = s.trim().to_lowercase();
115
116    if let Ok(v) = try_one_parsebytes(&s, "eib", EIB) {
117        Ok(v)
118    } else if let Ok(v) = try_one_parsebytes(&s, "pib", PIB) {
119        Ok(v)
120    } else if let Ok(v) = try_one_parsebytes(&s, "tib", TIB) {
121        Ok(v)
122    } else if let Ok(v) = try_one_parsebytes(&s, "gib", GIB) {
123        Ok(v)
124    } else if let Ok(v) = try_one_parsebytes(&s, "mib", MIB) {
125        Ok(v)
126    } else if let Ok(v) = try_one_parsebytes(&s, "kib", KIB) {
127        Ok(v)
128    } else if let Ok(v) = try_one_parsebytes(&s, "e", EIB) {
129        Ok(v)
130    } else if let Ok(v) = try_one_parsebytes(&s, "p", PIB) {
131        Ok(v)
132    } else if let Ok(v) = try_one_parsebytes(&s, "t", TIB) {
133        Ok(v)
134    } else if let Ok(v) = try_one_parsebytes(&s, "g", GIB) {
135        Ok(v)
136    } else if let Ok(v) = try_one_parsebytes(&s, "m", MIB) {
137        Ok(v)
138    } else if let Ok(v) = try_one_parsebytes(&s, "k", KIB) {
139        Ok(v)
140    } else if let Ok(v) = try_one_parsebytes(&s, "eb", EB) {
141        Ok(v)
142    } else if let Ok(v) = try_one_parsebytes(&s, "pb", PB) {
143        Ok(v)
144    } else if let Ok(v) = try_one_parsebytes(&s, "tb", TB) {
145        Ok(v)
146    } else if let Ok(v) = try_one_parsebytes(&s, "gb", GB) {
147        Ok(v)
148    } else if let Ok(v) = try_one_parsebytes(&s, "mb", MB) {
149        Ok(v)
150    } else if let Ok(v) = try_one_parsebytes(&s, "kb", KB) {
151        Ok(v)
152    } else if let Ok(v) = s.parse::<u64>() {
153        // byte count w/o suffix.
154        Ok(v)
155    } else {
156        Err(ah::format_err!("Cannot parse byte count: {s}"))
157    }
158}
159
160pub trait Hhmmss {
161    fn hhmmss(&self) -> String;
162}
163
164impl Hhmmss for Duration {
165    fn hhmmss(&self) -> String {
166        let secs = self.as_secs();
167        let secs_lim = (99 * 60 * 60) + (59 * 60) + 59;
168        let lim = if secs > secs_lim { ">" } else { "" };
169        let secs = secs.min(secs_lim);
170        let h = secs / (60 * 60);
171        let rem = secs % (60 * 60);
172        let m = rem / 60;
173        let s = rem % 60;
174        format!("{lim}{h:02}h:{m:02}m:{s:02}s")
175    }
176}
177
178/// Fold a byte vector into a smaller byte vector using XOR operation.
179/// If `output_size` is bigger than `input.len()`, the trailing bytes
180/// will be filled with zeros.
181pub fn fold(input: &[u8], output_size: usize) -> Vec<u8> {
182    let mut output = vec![0; output_size];
183
184    if output_size > 0 {
185        for (i, data) in input.iter().enumerate() {
186            output[i % output_size] ^= data;
187        }
188    }
189
190    output
191}
192
193#[cfg(test)]
194mod tests {
195    use super::*;
196
197    #[test]
198    fn test_prettybytes() {
199        assert_eq!(prettybytes(42, true, true, false), "42 bytes");
200        assert_eq!(
201            prettybytes(42 * 1024, true, true, false),
202            "42.0 kiB (43.0 kB)"
203        );
204        assert_eq!(
205            prettybytes(42 * 1024 * 1024, true, true, false),
206            "42.0 MiB (44.0 MB)"
207        );
208        assert_eq!(
209            prettybytes(42 * 1024 * 1024 * 1024, true, true, false),
210            "42.00 GiB (45.10 GB)"
211        );
212        assert_eq!(
213            prettybytes(42 * 1024 * 1024 * 1024 * 1024, true, true, false),
214            "42.0000 TiB (46.1795 TB)"
215        );
216        assert_eq!(
217            prettybytes(42 * 1024 * 1024 * 1024 * 1024 * 1024, true, true, false),
218            "42.0000 PiB (47.2878 PB)"
219        );
220        assert_eq!(
221            prettybytes(
222                2 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
223                true,
224                true,
225                false
226            ),
227            "2.0000 EiB (2.3058 EB)"
228        );
229
230        assert_eq!(prettybytes(42, true, false, false), "42 bytes");
231        assert_eq!(prettybytes(42, false, true, false), "42 bytes");
232        assert_eq!(prettybytes(42, false, false, false), "");
233
234        assert_eq!(prettybytes(42 * 1024, true, false, false), "42.0 kiB");
235        assert_eq!(prettybytes(42 * 1024, false, true, false), "43.0 kB");
236        assert_eq!(prettybytes(42 * 1024, false, false, false), "");
237
238        assert_eq!(
239            prettybytes(42 * 1024, true, true, true),
240            "42.0 kiB (43.0 kB, 43008 bytes)"
241        );
242        assert_eq!(
243            prettybytes(42 * 1024, true, false, true),
244            "42.0 kiB (43008 bytes)"
245        );
246    }
247
248    #[test]
249    fn test_parsebytes() {
250        // No suffix.
251        assert_eq!(parsebytes("42").unwrap(), 42);
252
253        // Binary suffix, integer.
254        assert_eq!(parsebytes("42kib").unwrap(), 42 * 1024);
255        assert_eq!(parsebytes("42 mib").unwrap(), 42 * 1024 * 1024);
256        assert_eq!(parsebytes(" 42 gib ").unwrap(), 42 * 1024 * 1024 * 1024);
257        assert_eq!(parsebytes("42Tib").unwrap(), 42 * 1024 * 1024 * 1024 * 1024);
258        assert_eq!(
259            parsebytes("42PiB").unwrap(),
260            42 * 1024 * 1024 * 1024 * 1024 * 1024
261        );
262        assert_eq!(
263            parsebytes("2 EIB ").unwrap(),
264            2 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024
265        );
266
267        // Binary suffix, fractional.
268        assert_eq!(parsebytes("42.5kib").unwrap(), 42 * 1024 + 1024 / 2);
269        assert_eq!(
270            parsebytes("42.5 mib").unwrap(),
271            42 * 1024 * 1024 + 1024 * 1024 / 2
272        );
273        assert_eq!(
274            parsebytes(" 42.5 gib ").unwrap(),
275            42 * 1024 * 1024 * 1024 + 1024 * 1024 * 1024 / 2
276        );
277        assert_eq!(
278            parsebytes("42.5Tib").unwrap(),
279            42 * 1024 * 1024 * 1024 * 1024 + 1024 * 1024 * 1024 * 1024 / 2
280        );
281        assert_eq!(
282            parsebytes("42.5PiB").unwrap(),
283            42 * 1024 * 1024 * 1024 * 1024 * 1024 + 1024 * 1024 * 1024 * 1024 * 1024 / 2
284        );
285        assert_eq!(
286            parsebytes("1.5 EIB ").unwrap(),
287            1024 * 1024 * 1024 * 1024 * 1024 * 1024 + 1024 * 1024 * 1024 * 1024 * 1024 * 1024 / 2
288        );
289
290        // Binary suffix, integer.
291        assert_eq!(parsebytes("42k").unwrap(), 42 * 1024);
292        assert_eq!(parsebytes("42 m").unwrap(), 42 * 1024 * 1024);
293        assert_eq!(parsebytes(" 42 g ").unwrap(), 42 * 1024 * 1024 * 1024);
294        assert_eq!(parsebytes("42T").unwrap(), 42 * 1024 * 1024 * 1024 * 1024);
295        assert_eq!(
296            parsebytes("42P").unwrap(),
297            42 * 1024 * 1024 * 1024 * 1024 * 1024
298        );
299        assert_eq!(
300            parsebytes("2 E ").unwrap(),
301            2 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024
302        );
303
304        // Binary suffix, fractional.
305        assert_eq!(parsebytes("42.5k").unwrap(), 42 * 1024 + 1024 / 2);
306        assert_eq!(
307            parsebytes("42.5 m").unwrap(),
308            42 * 1024 * 1024 + 1024 * 1024 / 2
309        );
310        assert_eq!(
311            parsebytes(" 42.5 g ").unwrap(),
312            42 * 1024 * 1024 * 1024 + 1024 * 1024 * 1024 / 2
313        );
314        assert_eq!(
315            parsebytes("42.5T").unwrap(),
316            42 * 1024 * 1024 * 1024 * 1024 + 1024 * 1024 * 1024 * 1024 / 2
317        );
318        assert_eq!(
319            parsebytes("42.5P").unwrap(),
320            42 * 1024 * 1024 * 1024 * 1024 * 1024 + 1024 * 1024 * 1024 * 1024 * 1024 / 2
321        );
322        assert_eq!(
323            parsebytes("1.5 E ").unwrap(),
324            1024 * 1024 * 1024 * 1024 * 1024 * 1024 + 1024 * 1024 * 1024 * 1024 * 1024 * 1024 / 2
325        );
326
327        // Decimal suffix, integer.
328        assert_eq!(parsebytes("42kb").unwrap(), 42 * 1000);
329        assert_eq!(parsebytes("42 mb").unwrap(), 42 * 1000 * 1000);
330        assert_eq!(parsebytes(" 42 gb ").unwrap(), 42 * 1000 * 1000 * 1000);
331        assert_eq!(parsebytes("42Tb").unwrap(), 42 * 1000 * 1000 * 1000 * 1000);
332        assert_eq!(
333            parsebytes("42PB").unwrap(),
334            42 * 1000 * 1000 * 1000 * 1000 * 1000
335        );
336        assert_eq!(
337            parsebytes("2 EB ").unwrap(),
338            2 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000
339        );
340
341        // Decimal suffix, fractional.
342        assert_eq!(parsebytes("42.5kb").unwrap(), 42 * 1000 + 1000 / 2);
343        assert_eq!(
344            parsebytes("42.5 mb").unwrap(),
345            42 * 1000 * 1000 + 1000 * 1000 / 2
346        );
347        assert_eq!(
348            parsebytes(" 42.5 gb ").unwrap(),
349            42 * 1000 * 1000 * 1000 + 1000 * 1000 * 1000 / 2
350        );
351        assert_eq!(
352            parsebytes("42.5Tb").unwrap(),
353            42 * 1000 * 1000 * 1000 * 1000 + 1000 * 1000 * 1000 * 1000 / 2
354        );
355        assert_eq!(
356            parsebytes("42.5PB").unwrap(),
357            42 * 1000 * 1000 * 1000 * 1000 * 1000 + 1000 * 1000 * 1000 * 1000 * 1000 / 2
358        );
359        assert_eq!(
360            parsebytes("1.5 EB ").unwrap(),
361            1000 * 1000 * 1000 * 1000 * 1000 * 1000 + 1000 * 1000 * 1000 * 1000 * 1000 * 1000 / 2
362        );
363    }
364
365    #[test]
366    fn test_hhmmss() {
367        assert_eq!(Duration::from_secs(0).hhmmss(), "00h:00m:00s");
368        assert_eq!(
369            Duration::from_secs((2 * 60 * 60) + (3 * 60) + 4).hhmmss(),
370            "02h:03m:04s"
371        );
372        assert_eq!(
373            Duration::from_secs((23 * 60 * 60) + (59 * 60) + 59).hhmmss(),
374            "23h:59m:59s"
375        );
376        assert_eq!(
377            Duration::from_secs((99 * 60 * 60) + (59 * 60) + 59).hhmmss(),
378            "99h:59m:59s"
379        );
380        assert_eq!(
381            Duration::from_secs((99 * 60 * 60) + (59 * 60) + 59 + 1).hhmmss(),
382            ">99h:59m:59s"
383        );
384    }
385
386    #[test]
387    fn test_fold() {
388        assert_eq!(fold(&[0x55, 0x55, 0xAA, 0xAA], 2), [0xFF, 0xFF]);
389        assert_eq!(fold(&[0x55, 0x55, 0x55, 0x55], 2), [0x00, 0x00]);
390        assert_eq!(fold(&[0x55, 0x55, 0xAA, 0x55], 2), [0xFF, 0x00]);
391        assert_eq!(fold(&[0x55, 0x55, 0x55, 0xAA], 2), [0x00, 0xFF]);
392        assert_eq!(
393            fold(&[0x98, 0xB1, 0x5B, 0x47, 0x8F, 0xF7, 0x9C, 0x6F], 3),
394            [0x43, 0x51, 0xAC]
395        );
396        assert_eq!(fold(&[0x12, 0x34, 0x56, 0x78], 4), [0x12, 0x34, 0x56, 0x78]);
397        assert_eq!(
398            fold(&[0x12, 0x34, 0x56, 0x78], 6),
399            [0x12, 0x34, 0x56, 0x78, 0x00, 0x00]
400        );
401        assert_eq!(fold(&[0x12, 0x34, 0x56, 0x78], 0), []);
402    }
403}
404
405// vim: ts=4 sw=4 expandtab