graft_core/
byte_unit.rs

1use std::{
2    fmt::{self, Debug, Display},
3    ops::{Add, Div, Mul, Range, Rem, Shl, Shr, Sub},
4    str::FromStr,
5};
6
7use serde::{Deserialize, Deserializer, Serialize};
8use thiserror::Error;
9
10struct NamedByteUnit {
11    value: ByteUnit,
12    suffix: &'static str,
13}
14
15const KB: NamedByteUnit = NamedByteUnit { value: ByteUnit::KB, suffix: "KB" };
16const MB: NamedByteUnit = NamedByteUnit { value: ByteUnit::MB, suffix: "MB" };
17const GB: NamedByteUnit = NamedByteUnit { value: ByteUnit::GB, suffix: "GB" };
18const TB: NamedByteUnit = NamedByteUnit { value: ByteUnit::TB, suffix: "TB" };
19const PB: NamedByteUnit = NamedByteUnit { value: ByteUnit::PB, suffix: "PB" };
20const EB: NamedByteUnit = NamedByteUnit { value: ByteUnit::EB, suffix: "EB" };
21
22#[derive(Clone, Copy, Eq, Ord)]
23pub struct ByteUnit(u64);
24
25impl ByteUnit {
26    pub const ZERO: ByteUnit = ByteUnit(0);
27    pub const MAX: ByteUnit = ByteUnit(u64::MAX);
28
29    pub const KB: ByteUnit = ByteUnit(1 << 10);
30    pub const MB: ByteUnit = ByteUnit(1 << 20);
31    pub const GB: ByteUnit = ByteUnit(1 << 30);
32    pub const TB: ByteUnit = ByteUnit(1 << 40);
33    pub const PB: ByteUnit = ByteUnit(1 << 50);
34    pub const EB: ByteUnit = ByteUnit(1 << 60);
35
36    pub const fn new(bytes: u64) -> ByteUnit {
37        ByteUnit(bytes)
38    }
39
40    pub const fn size_of<T>() -> ByteUnit {
41        ByteUnit(std::mem::size_of::<T>() as u64)
42    }
43
44    pub const fn as_u64(&self) -> u64 {
45        self.0
46    }
47
48    pub const fn as_u32(&self) -> u32 {
49        self.0 as u32
50    }
51
52    pub const fn as_u16(&self) -> u16 {
53        self.0 as u16
54    }
55
56    pub const fn as_usize(&self) -> usize {
57        self.0 as usize
58    }
59
60    const fn as_f64(&self) -> f64 {
61        self.0 as f64
62    }
63
64    pub const fn is_power_of_two(&self) -> bool {
65        self.0.is_power_of_two()
66    }
67
68    pub const fn from_kb(kb: u64) -> ByteUnit {
69        ByteUnit(kb.saturating_mul(KB.value.0))
70    }
71
72    pub const fn from_mb(mb: u64) -> ByteUnit {
73        ByteUnit(mb.saturating_mul(MB.value.0))
74    }
75
76    pub const fn from_gb(gb: u64) -> ByteUnit {
77        ByteUnit(gb.saturating_mul(GB.value.0))
78    }
79
80    pub const fn from_tb(tb: u64) -> ByteUnit {
81        ByteUnit(tb.saturating_mul(TB.value.0))
82    }
83
84    pub const fn from_pb(pb: u64) -> ByteUnit {
85        ByteUnit(pb.saturating_mul(PB.value.0))
86    }
87
88    pub const fn from_eb(eb: u64) -> ByteUnit {
89        ByteUnit(eb.saturating_mul(EB.value.0))
90    }
91
92    /// Returns the absolute difference between two `ByteUnit` values.
93    pub const fn diff(self, other: ByteUnit) -> ByteUnit {
94        if self.0 > other.0 {
95            ByteUnit(self.0 - other.0)
96        } else {
97            ByteUnit(other.0 - self.0)
98        }
99    }
100
101    /// Returns a range representing the byte range from `self` to `end`.
102    pub const fn range(self, end: ByteUnit) -> Range<usize> {
103        (self.0 as usize)..(end.0 as usize)
104    }
105}
106
107impl Display for ByteUnit {
108    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
109        let value = *self;
110        for unit in &[EB, PB, TB, GB, MB, KB] {
111            if value >= unit.value {
112                let whole = value / unit.value;
113                let rem = value % unit.value;
114                let frac = rem.as_f64() / unit.value.as_f64();
115                if frac < 0.005 {
116                    write!(f, "{} {}", whole.0, unit.suffix)?;
117                } else if frac >= 0.95 {
118                    write!(f, "{} {}", whole.0 + 1, unit.suffix)?;
119                } else {
120                    write!(f, "{}.{:02.0} {}", whole.0, frac * 100.0, unit.suffix)?;
121                }
122                return Ok(());
123            }
124        }
125
126        // If we reach this point, the value is smaller than 1KB
127        write!(f, "{} B", value.0)
128    }
129}
130
131impl Debug for ByteUnit {
132    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
133        write!(f, "{}", self.0)
134    }
135}
136
137#[derive(Debug, Error)]
138pub enum ByteUnitParseError {
139    #[error("Invalid format: got {0}, expected <number> [<unit>]")]
140    InvalidFormat(String),
141
142    #[error("Invalid number: {0}")]
143    InvalidNumber(String),
144
145    #[error("Unknown unit: {0}")]
146    UnknownUnit(String),
147}
148
149impl FromStr for ByteUnit {
150    type Err = ByteUnitParseError;
151
152    fn from_str(s: &str) -> Result<Self, Self::Err> {
153        let s = s.trim();
154        let mut parts = s.split_ascii_whitespace();
155        let number = parts
156            .next()
157            .ok_or(ByteUnitParseError::InvalidFormat(s.to_string()))?;
158
159        let unit = parts.next().unwrap_or("B").to_uppercase();
160
161        let value = number
162            .parse::<u64>()
163            .map_err(|_| ByteUnitParseError::InvalidNumber(number.into()))?;
164
165        Ok(match unit.as_str() {
166            "B" => ByteUnit(value),
167            "KB" => ByteUnit::from_kb(value),
168            "MB" => ByteUnit::from_mb(value),
169            "GB" => ByteUnit::from_gb(value),
170            "TB" => ByteUnit::from_tb(value),
171            "PB" => ByteUnit::from_pb(value),
172            "EB" => ByteUnit::from_eb(value),
173            _ => return Err(ByteUnitParseError::UnknownUnit(unit)),
174        })
175    }
176}
177
178impl<'de> Deserialize<'de> for ByteUnit {
179    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
180    where
181        D: Deserializer<'de>,
182    {
183        if deserializer.is_human_readable() {
184            let s = String::deserialize(deserializer)?;
185            FromStr::from_str(&s).map_err(serde::de::Error::custom)
186        } else {
187            let bytes = u64::deserialize(deserializer)?;
188            Ok(ByteUnit(bytes))
189        }
190    }
191}
192
193impl Serialize for ByteUnit {
194    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
195    where
196        S: serde::Serializer,
197    {
198        if serializer.is_human_readable() {
199            self.to_string().serialize(serializer)
200        } else {
201            self.0.serialize(serializer)
202        }
203    }
204}
205
206impl<T: Into<ByteUnit> + Copy> PartialEq<T> for ByteUnit {
207    fn eq(&self, other: &T) -> bool {
208        self.0 == (*other).into().0
209    }
210}
211
212impl<T: Into<ByteUnit> + Copy> PartialOrd<T> for ByteUnit {
213    fn partial_cmp(&self, other: &T) -> Option<std::cmp::Ordering> {
214        self.0.partial_cmp(&(*other).into().0)
215    }
216}
217
218impl<T: Into<ByteUnit>> Mul<T> for ByteUnit {
219    type Output = Self;
220
221    #[inline(always)]
222    fn mul(self, rhs: T) -> Self::Output {
223        ByteUnit(self.0.saturating_mul(rhs.into().0))
224    }
225}
226
227impl<T: Into<ByteUnit>> Add<T> for ByteUnit {
228    type Output = Self;
229
230    #[inline(always)]
231    fn add(self, rhs: T) -> Self::Output {
232        ByteUnit(self.0.saturating_add(rhs.into().0))
233    }
234}
235impl<T: Into<ByteUnit>> Sub<T> for ByteUnit {
236    type Output = Self;
237
238    #[inline(always)]
239    fn sub(self, rhs: T) -> Self::Output {
240        ByteUnit(self.0.saturating_sub(rhs.into().0))
241    }
242}
243
244impl<T: Into<ByteUnit>> Div<T> for ByteUnit {
245    type Output = Self;
246
247    #[inline(always)]
248    fn div(self, rhs: T) -> Self::Output {
249        ByteUnit(self.0.saturating_div(rhs.into().0))
250    }
251}
252
253impl<T: Into<ByteUnit>> Rem<T> for ByteUnit {
254    type Output = Self;
255
256    #[inline(always)]
257    fn rem(self, rhs: T) -> Self::Output {
258        let value = rhs.into().0;
259        match value {
260            0 => ByteUnit(0),
261            _ => ByteUnit(self.0 % value),
262        }
263    }
264}
265
266impl<T: Into<ByteUnit>> Shl<T> for ByteUnit {
267    type Output = Self;
268
269    #[inline(always)]
270    fn shl(self, rhs: T) -> Self::Output {
271        ByteUnit(self.0 << rhs.into().0)
272    }
273}
274
275impl<T: Into<ByteUnit>> Shr<T> for ByteUnit {
276    type Output = ByteUnit;
277
278    #[inline(always)]
279    fn shr(self, rhs: T) -> Self::Output {
280        ByteUnit(self.0 >> rhs.into().0)
281    }
282}
283
284macro_rules! impl_arith_op {
285    ($T:ident, $Trait:ident, $func:ident, $op:tt) => (
286        impl $Trait<ByteUnit> for $T {
287            type Output = ByteUnit;
288
289            #[inline(always)]
290            fn $func(self, rhs: ByteUnit) -> Self::Output {
291                ByteUnit::from(self) $op rhs
292            }
293        }
294    )
295}
296
297macro_rules! impl_primitive {
298    ($T:ident) => {
299        impl From<$T> for ByteUnit {
300            fn from(bytes: $T) -> ByteUnit {
301                ByteUnit(bytes as u64)
302            }
303        }
304
305        impl PartialEq<ByteUnit> for $T {
306            fn eq(&self, other: &ByteUnit) -> bool {
307                *self as u64 == other.0
308            }
309        }
310
311        impl PartialOrd<ByteUnit> for $T {
312            fn partial_cmp(&self, other: &ByteUnit) -> Option<std::cmp::Ordering> {
313                (*self as u64).partial_cmp(&other.0)
314            }
315        }
316
317        impl_arith_op!($T, Mul, mul, *);
318        impl_arith_op!($T, Div, div, /);
319        impl_arith_op!($T, Rem, rem, %);
320        impl_arith_op!($T, Add, add, +);
321        impl_arith_op!($T, Sub, sub, -);
322        impl_arith_op!($T, Shl, shl, <<);
323        impl_arith_op!($T, Shr, shr, >>);
324    };
325}
326
327impl_primitive!(u8);
328impl_primitive!(u16);
329impl_primitive!(u32);
330impl_primitive!(u64);
331impl_primitive!(usize);
332
333impl_primitive!(i8);
334impl_primitive!(i16);
335impl_primitive!(i32);
336impl_primitive!(i64);
337impl_primitive!(isize);
338
339#[cfg(test)]
340mod tests {
341    use super::*;
342
343    #[graft_test::test]
344    fn test_sanity() {
345        assert_eq!(0, ByteUnit::ZERO);
346        assert_eq!(ByteUnit::ZERO, 0);
347        assert_eq!(ByteUnit::ZERO, ByteUnit::ZERO);
348        assert!(0 < ByteUnit::KB);
349        assert!(ByteUnit::KB > 0);
350    }
351
352    #[graft_test::test]
353    fn test_display() {
354        assert_eq!(format!("{}", ByteUnit::ZERO), "0 B");
355        assert_eq!(format!("{}", ByteUnit::MAX), "16 EB");
356
357        for unit in &[KB, MB, GB, TB, PB, EB] {
358            assert_eq!(format!("{}", unit.value), format!("1 {}", unit.suffix));
359            assert_eq!(format!("{}", 7 * unit.value), format!("7 {}", unit.suffix));
360        }
361
362        assert_eq!(
363            format!("{}", (7 * ByteUnit::MB) + (132 * ByteUnit::KB)),
364            "7.13 MB"
365        );
366        assert_eq!(
367            format!("{}", (7 * ByteUnit::MB) + (512 * ByteUnit::KB)),
368            "7.50 MB"
369        );
370        assert_eq!(
371            format!("{}", (7 * ByteUnit::MB) + ((1024 - 1) * ByteUnit::KB)),
372            "8 MB"
373        );
374    }
375
376    #[graft_test::test]
377    fn test_const() {
378        const X: ByteUnit = ByteUnit::from_kb(4);
379        let arr: [u8; X.as_usize()] = [0; X.as_usize()];
380        assert_eq!(arr.len(), 4 * 1024);
381    }
382
383    #[graft_test::test]
384    fn test_parse() {
385        let cases = [
386            ByteUnit::new(0),
387            ByteUnit::from_kb(10),
388            ByteUnit::from_mb(10),
389            ByteUnit::from_gb(10),
390            ByteUnit::from_tb(10),
391            ByteUnit::from_pb(10),
392            ByteUnit::from_eb(10),
393        ];
394
395        // for each case, check that it roundtrips through display then parse
396        for &unit in &cases {
397            let s = format!("{unit}");
398            let parsed = s.parse::<ByteUnit>().unwrap();
399            println!("parsed `{s}` into {parsed}");
400            assert_eq!(unit, parsed);
401        }
402
403        // some valid cases that exercise the parser's flexibility
404        let nonstandard_cases = [
405            ("0", "no unit", ByteUnit::new(0)),
406            ("    0", "no unit prefix whitespace", ByteUnit::new(0)),
407            ("0 ", "no unit trailing whitespace", ByteUnit::new(0)),
408            (" 0 ", "no unit both whitespace", ByteUnit::new(0)),
409            ("0  kb", "lowercase", ByteUnit::new(0)),
410            ("5  kb", "lowercase kb", ByteUnit::from_kb(5)),
411            (" 5 \t kb  ", "lots of whitespace", ByteUnit::from_kb(5)),
412        ];
413
414        // check that each case parses
415        for &(s, desc, expected) in &nonstandard_cases {
416            let parsed = s.parse::<ByteUnit>().unwrap();
417            println!("parsed `{s}` into {parsed}");
418            assert_eq!(parsed, expected, "{desc}");
419        }
420
421        let invalid_cases = [
422            ("", "empty string", "Invalid format"),
423            ("    ", "only whitespace", "Invalid format"),
424            ("12.2", "decimal number", "Invalid number"),
425            ("12.2 kb", "decimal number", "Invalid number"),
426            ("123 xb", "unknown unit", "Unknown unit"),
427        ];
428
429        // check that each case fails to parse, and the error contains the expected message
430        for &(s, desc, expected) in &invalid_cases {
431            let parsed = s.parse::<ByteUnit>();
432            println!("parsed `{s}` into {parsed:?}");
433            assert!(parsed.is_err(), "{}", desc);
434            assert!(
435                parsed.unwrap_err().to_string().contains(expected),
436                "{}",
437                desc
438            );
439        }
440    }
441}
graft_core/byte_unit.rs

graft_core/
byte_unit.rs
