common_base/
readable_size.rs

1// Copyright (c) 2017-present, PingCAP, Inc. Licensed under Apache-2.0.
2
3// This file is copied from https://github.com/tikv/raft-engine/blob/0.3.0/src/util.rs
4
5use std::fmt::{self, Debug, Display, Write};
6use std::ops::{Div, Mul};
7use std::str::FromStr;
8
9use serde::de::{Unexpected, Visitor};
10use serde::{de, Deserialize, Deserializer, Serialize, Serializer};
11
12const UNIT: u64 = 1;
13
14const BINARY_DATA_MAGNITUDE: u64 = 1024;
15pub const B: u64 = UNIT;
16pub const KIB: u64 = B * BINARY_DATA_MAGNITUDE;
17pub const MIB: u64 = KIB * BINARY_DATA_MAGNITUDE;
18pub const GIB: u64 = MIB * BINARY_DATA_MAGNITUDE;
19pub const TIB: u64 = GIB * BINARY_DATA_MAGNITUDE;
20pub const PIB: u64 = TIB * BINARY_DATA_MAGNITUDE;
21
22#[derive(Clone, Copy, PartialEq, Eq, Ord, PartialOrd, Default)]
23pub struct ReadableSize(pub u64);
24
25impl ReadableSize {
26    pub const fn kb(count: u64) -> ReadableSize {
27        ReadableSize(count * KIB)
28    }
29
30    pub const fn mb(count: u64) -> ReadableSize {
31        ReadableSize(count * MIB)
32    }
33
34    pub const fn gb(count: u64) -> ReadableSize {
35        ReadableSize(count * GIB)
36    }
37
38    pub const fn as_mb(self) -> u64 {
39        self.0 / MIB
40    }
41
42    pub const fn as_bytes(self) -> u64 {
43        self.0
44    }
45}
46
47impl Div<u64> for ReadableSize {
48    type Output = ReadableSize;
49
50    fn div(self, rhs: u64) -> ReadableSize {
51        ReadableSize(self.0 / rhs)
52    }
53}
54
55impl Div<ReadableSize> for ReadableSize {
56    type Output = u64;
57
58    fn div(self, rhs: ReadableSize) -> u64 {
59        self.0 / rhs.0
60    }
61}
62
63impl Mul<u64> for ReadableSize {
64    type Output = ReadableSize;
65
66    fn mul(self, rhs: u64) -> ReadableSize {
67        ReadableSize(self.0 * rhs)
68    }
69}
70
71impl Serialize for ReadableSize {
72    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
73    where
74        S: Serializer,
75    {
76        let size = self.0;
77        let mut buffer = String::new();
78        if size == 0 {
79            write!(buffer, "{}KiB", size).unwrap();
80        } else if size % PIB == 0 {
81            write!(buffer, "{}PiB", size / PIB).unwrap();
82        } else if size % TIB == 0 {
83            write!(buffer, "{}TiB", size / TIB).unwrap();
84        } else if size % GIB as u64 == 0 {
85            write!(buffer, "{}GiB", size / GIB).unwrap();
86        } else if size % MIB as u64 == 0 {
87            write!(buffer, "{}MiB", size / MIB).unwrap();
88        } else if size % KIB as u64 == 0 {
89            write!(buffer, "{}KiB", size / KIB).unwrap();
90        } else {
91            return serializer.serialize_u64(size);
92        }
93        serializer.serialize_str(&buffer)
94    }
95}
96
97impl FromStr for ReadableSize {
98    type Err = String;
99
100    // This method parses value in binary unit.
101    fn from_str(s: &str) -> Result<ReadableSize, String> {
102        let size_str = s.trim();
103        if size_str.is_empty() {
104            return Err(format!("{:?} is not a valid size.", s));
105        }
106
107        if !size_str.is_ascii() {
108            return Err(format!("ASCII string is expected, but got {:?}", s));
109        }
110
111        // size: digits and '.' as decimal separator
112        let size_len = size_str
113            .to_string()
114            .chars()
115            .take_while(|c| char::is_ascii_digit(c) || ['.', 'e', 'E', '-', '+'].contains(c))
116            .count();
117
118        // unit: alphabetic characters
119        let (size, unit) = size_str.split_at(size_len);
120
121        let unit = match unit.trim() {
122            "K" | "KB" | "KiB" => KIB,
123            "M" | "MB" | "MiB" => MIB,
124            "G" | "GB" | "GiB" => GIB,
125            "T" | "TB" | "TiB" => TIB,
126            "P" | "PB" | "PiB" => PIB,
127            "B" | "" => B,
128            _ => {
129                return Err(format!(
130                    "only B, KB, KiB, MB, MiB, GB, GiB, TB, TiB, PB, and PiB are supported: {:?}",
131                    s
132                ));
133            }
134        };
135
136        match size.parse::<f64>() {
137            Ok(n) => Ok(ReadableSize((n * unit as f64) as u64)),
138            Err(_) => Err(format!("invalid size string: {:?}", s)),
139        }
140    }
141}
142
143impl Debug for ReadableSize {
144    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
145        write!(f, "{}", self)
146    }
147}
148
149impl Display for ReadableSize {
150    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
151        if self.0 >= PIB {
152            write!(f, "{:.1}PiB", self.0 as f64 / PIB as f64)
153        } else if self.0 >= TIB {
154            write!(f, "{:.1}TiB", self.0 as f64 / TIB as f64)
155        } else if self.0 >= GIB {
156            write!(f, "{:.1}GiB", self.0 as f64 / GIB as f64)
157        } else if self.0 >= MIB {
158            write!(f, "{:.1}MiB", self.0 as f64 / MIB as f64)
159        } else if self.0 >= KIB {
160            write!(f, "{:.1}KiB", self.0 as f64 / KIB as f64)
161        } else {
162            write!(f, "{}B", self.0)
163        }
164    }
165}
166
167impl<'de> Deserialize<'de> for ReadableSize {
168    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
169    where
170        D: Deserializer<'de>,
171    {
172        struct SizeVisitor;
173
174        impl<'de> Visitor<'de> for SizeVisitor {
175            type Value = ReadableSize;
176
177            fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
178                formatter.write_str("valid size")
179            }
180
181            fn visit_i64<E>(self, size: i64) -> Result<ReadableSize, E>
182            where
183                E: de::Error,
184            {
185                if size >= 0 {
186                    self.visit_u64(size as u64)
187                } else {
188                    Err(E::invalid_value(Unexpected::Signed(size), &self))
189                }
190            }
191
192            fn visit_u64<E>(self, size: u64) -> Result<ReadableSize, E>
193            where
194                E: de::Error,
195            {
196                Ok(ReadableSize(size))
197            }
198
199            fn visit_str<E>(self, size_str: &str) -> Result<ReadableSize, E>
200            where
201                E: de::Error,
202            {
203                size_str.parse().map_err(E::custom)
204            }
205        }
206
207        deserializer.deserialize_any(SizeVisitor)
208    }
209}
210
211#[cfg(test)]
212mod tests {
213    use super::*;
214
215    #[test]
216    fn test_readable_size() {
217        let s = ReadableSize::kb(2);
218        assert_eq!(s.0, 2048);
219        assert_eq!(s.as_mb(), 0);
220        let s = ReadableSize::mb(2);
221        assert_eq!(s.0, 2 * 1024 * 1024);
222        assert_eq!(s.as_mb(), 2);
223        let s = ReadableSize::gb(2);
224        assert_eq!(s.0, 2 * 1024 * 1024 * 1024);
225        assert_eq!(s.as_mb(), 2048);
226
227        assert_eq!((ReadableSize::mb(2) / 2).0, MIB);
228        assert_eq!((ReadableSize::mb(1) / 2).0, 512 * KIB);
229        assert_eq!(ReadableSize::mb(2) / ReadableSize::kb(1), 2048);
230    }
231
232    #[test]
233    fn test_parse_readable_size() {
234        #[derive(Serialize, Deserialize)]
235        struct SizeHolder {
236            s: ReadableSize,
237        }
238
239        let legal_cases = vec![
240            (0, "0KiB"),
241            (2 * KIB, "2KiB"),
242            (4 * MIB, "4MiB"),
243            (5 * GIB, "5GiB"),
244            (7 * TIB, "7TiB"),
245            (11 * PIB, "11PiB"),
246        ];
247        for (size, exp) in legal_cases {
248            let c = SizeHolder {
249                s: ReadableSize(size),
250            };
251            let res_str = toml::to_string(&c).unwrap();
252            let exp_str = format!("s = {:?}\n", exp);
253            assert_eq!(res_str, exp_str);
254            let res_size: SizeHolder = toml::from_str(&exp_str).unwrap();
255            assert_eq!(res_size.s.0, size);
256        }
257
258        let c = SizeHolder {
259            s: ReadableSize(512),
260        };
261        let res_str = toml::to_string(&c).unwrap();
262        assert_eq!(res_str, "s = 512\n");
263        let res_size: SizeHolder = toml::from_str(&res_str).unwrap();
264        assert_eq!(res_size.s.0, c.s.0);
265
266        let decode_cases = vec![
267            (" 0.5 PB", PIB / 2),
268            ("0.5 TB", TIB / 2),
269            ("0.5GB ", GIB / 2),
270            ("0.5MB", MIB / 2),
271            ("0.5KB", KIB / 2),
272            ("0.5P", PIB / 2),
273            ("0.5T", TIB / 2),
274            ("0.5G", GIB / 2),
275            ("0.5M", MIB / 2),
276            ("0.5K", KIB / 2),
277            ("23", 23),
278            ("1", 1),
279            ("1024B", KIB),
280            // units with binary prefixes
281            (" 0.5 PiB", PIB / 2),
282            ("1PiB", PIB),
283            ("0.5 TiB", TIB / 2),
284            ("2 TiB", TIB * 2),
285            ("0.5GiB ", GIB / 2),
286            ("787GiB ", GIB * 787),
287            ("0.5MiB", MIB / 2),
288            ("3MiB", MIB * 3),
289            ("0.5KiB", KIB / 2),
290            ("1 KiB", KIB),
291            // scientific notation
292            ("0.5e6 B", B * 500000),
293            ("0.5E6 B", B * 500000),
294            ("1e6B", B * 1000000),
295            ("8E6B", B * 8000000),
296            ("8e7", B * 80000000),
297            ("1e-1MB", MIB / 10),
298            ("1e+1MB", MIB * 10),
299            ("0e+10MB", 0),
300        ];
301        for (src, exp) in decode_cases {
302            let src = format!("s = {:?}", src);
303            let res: SizeHolder = toml::from_str(&src).unwrap();
304            assert_eq!(res.s.0, exp);
305        }
306
307        let illegal_cases = vec![
308            "0.5kb", "0.5kB", "0.5Kb", "0.5k", "0.5g", "b", "gb", "1b", "B", "1K24B", " 5_KB",
309            "4B7", "5M_",
310        ];
311        for src in illegal_cases {
312            let src_str = format!("s = {:?}", src);
313            assert!(toml::from_str::<SizeHolder>(&src_str).is_err(), "{}", src);
314        }
315    }
316}
common_base/readable_size.rs

common_base/
readable_size.rs
