common_base/
readable_size.rs

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

common_base/
readable_size.rs
