tantivy_common/
byte_count.rs

1use std::iter::Sum;
2use std::ops::{Add, AddAssign};
3
4use serde::{Deserialize, Serialize};
5
6/// Indicates space usage in bytes
7#[derive(Copy, Clone, Default, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
8pub struct ByteCount(u64);
9
10impl std::fmt::Debug for ByteCount {
11    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
12        f.write_str(&self.human_readable())
13    }
14}
15
16impl std::fmt::Display for ByteCount {
17    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
18        f.write_str(&self.human_readable())
19    }
20}
21
22const SUFFIX_AND_THRESHOLD: [(&str, u64); 5] = [
23    ("KB", 1_000),
24    ("MB", 1_000_000),
25    ("GB", 1_000_000_000),
26    ("TB", 1_000_000_000_000),
27    ("PB", 1_000_000_000_000_000),
28];
29
30impl ByteCount {
31    #[inline]
32    pub fn get_bytes(&self) -> u64 {
33        self.0
34    }
35
36    pub fn human_readable(&self) -> String {
37        for (suffix, threshold) in SUFFIX_AND_THRESHOLD.iter().rev() {
38            if self.get_bytes() >= *threshold {
39                let unit_num = self.get_bytes() as f64 / *threshold as f64;
40                return format!("{unit_num:.2} {suffix}");
41            }
42        }
43        format!("{:.2} B", self.get_bytes())
44    }
45}
46
47impl From<u64> for ByteCount {
48    fn from(value: u64) -> Self {
49        ByteCount(value)
50    }
51}
52impl From<usize> for ByteCount {
53    fn from(value: usize) -> Self {
54        ByteCount(value as u64)
55    }
56}
57
58impl Sum for ByteCount {
59    #[inline]
60    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
61        iter.fold(ByteCount::default(), |acc, x| acc + x)
62    }
63}
64
65impl PartialEq<u64> for ByteCount {
66    #[inline]
67    fn eq(&self, other: &u64) -> bool {
68        self.get_bytes() == *other
69    }
70}
71
72impl PartialOrd<u64> for ByteCount {
73    #[inline]
74    fn partial_cmp(&self, other: &u64) -> Option<std::cmp::Ordering> {
75        self.get_bytes().partial_cmp(other)
76    }
77}
78
79impl Add for ByteCount {
80    type Output = Self;
81
82    #[inline]
83    fn add(self, other: Self) -> Self {
84        Self(self.get_bytes() + other.get_bytes())
85    }
86}
87
88impl AddAssign for ByteCount {
89    #[inline]
90    fn add_assign(&mut self, other: Self) {
91        *self = Self(self.get_bytes() + other.get_bytes());
92    }
93}
94
95#[cfg(test)]
96mod test {
97    use crate::ByteCount;
98
99    #[test]
100    fn test_bytes() {
101        assert_eq!(ByteCount::from(0u64).human_readable(), "0 B");
102        assert_eq!(ByteCount::from(300u64).human_readable(), "300 B");
103        assert_eq!(ByteCount::from(1_000_000u64).human_readable(), "1.00 MB");
104        assert_eq!(ByteCount::from(1_500_000u64).human_readable(), "1.50 MB");
105        assert_eq!(
106            ByteCount::from(1_500_000_000u64).human_readable(),
107            "1.50 GB"
108        );
109        assert_eq!(
110            ByteCount::from(3_213_000_000_000u64).human_readable(),
111            "3.21 TB"
112        );
113    }
114}