1use regex::Regex;
2use serde::{Deserialize, Serialize};
3
4#[derive(Debug, Clone, Default)]
5pub struct ByteSizeConf(u64);
6
7impl ByteSizeConf {
8 pub fn as_bytes(&self) -> u64 {
9 self.0
10 }
11
12 pub fn as_kilobytes(&self) -> u64 {
13 self.0 / 1000
14 }
15
16 pub fn as_kibibytes(&self) -> u64 {
17 self.0 / 1024
18 }
19
20 pub fn as_megabytes(&self) -> u64 {
21 self.0 / 10u64.pow(6)
22 }
23
24 pub fn as_mebibytes(&self) -> u64 {
25 self.0 / 2u64.pow(20)
26 }
27
28 pub fn as_gigabytes(&self) -> u64 {
29 self.0 / 10u64.pow(9)
30 }
31
32 pub fn as_gibibytes(&self) -> u64 {
33 self.0 / 2u64.pow(30)
34 }
35
36 fn to_string(&self) -> String {
37 let size = self.0;
38 if size < 1024 {
39 format!("{} bytes", size)
40 } else if size < 2u64.pow(20) {
41 format!("{} KiB", size / 1024)
42 } else if size < 2u64.pow(30) {
43 format!("{} MiB", size / 2u64.pow(20))
44 } else {
45 format!("{} GiB", size / 2u64.pow(30))
46 }
47 }
48}
49
50impl<'de> Deserialize<'de> for ByteSizeConf {
51 fn deserialize<D>(deserializer: D) -> Result<ByteSizeConf, D::Error>
52 where
53 D: serde::Deserializer<'de>,
54 {
55 let s = String::deserialize(deserializer)?;
56 let size = parse_byte_size(&s).map_err(serde::de::Error::custom)?;
57 Ok(ByteSizeConf(size))
58 }
59}
60
61impl Serialize for ByteSizeConf {
62 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
63 where
64 S: serde::Serializer,
65 {
66 self.to_string().serialize(serializer)
67 }
68}
69
70impl std::fmt::Display for ByteSizeConf {
71 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
72 write!(f, "{}", self.to_string())
73 }
74}
75
76fn parse_byte_size(s: &str) -> Result<u64, String> {
77 let re = Regex::new(r"^(\d+)\s*(b|kb|ki|mb|mi|gb|gi|tb|ti)$").map_err(|_| "Invalid regex")?;
78 let s = s.to_lowercase();
79 let caps = re
80 .captures(s.as_str())
81 .ok_or_else(|| format!("Invalid byte size: {}", s))?;
82 let size = caps
83 .get(1)
84 .ok_or_else(|| "Invalid byte size".to_string())?
85 .as_str()
86 .parse()
87 .map_err(|_| "Invalid byte size".to_string())?;
88 let unit = caps
89 .get(2)
90 .ok_or_else(|| "Invalid byte size".to_string())?
91 .as_str()
92 .to_lowercase();
93 let unit = unit.as_str();
94
95 match unit {
96 "b" => Ok(size),
97 "kb" => Ok(size * 1000),
98 "ki" => Ok(size * 2u64.pow(10)),
99 "mb" => Ok(size * 10u64.pow(6)),
100 "mi" => Ok(size * 2u64.pow(20)),
101 "gb" => Ok(size * 10u64.pow(9)),
102 "gi" => Ok(size * 2u64.pow(30)),
103 "tb" => Ok(size * 10u64.pow(12)),
104 "ti" => Ok(size * 2u64.pow(40)),
105 _ => Err("Invalid byte size".to_string()),
106 }
107}
108
109#[cfg(test)]
110mod tests {
111 use super::*;
112
113 #[test]
114 fn test_parse_byte_size() {
115 assert_eq!(parse_byte_size("1b").unwrap(), 1);
116 assert_eq!(parse_byte_size("1kb").unwrap(), 1000);
117 assert_eq!(parse_byte_size("1ki").unwrap(), 1024);
118 assert_eq!(parse_byte_size("1mb").unwrap(), 1000 * 1000);
119 assert_eq!(parse_byte_size("1mi").unwrap(), 1024 * 1024);
120 assert_eq!(parse_byte_size("1gb").unwrap(), 1000 * 1000 * 1000);
121 assert_eq!(parse_byte_size("1gi").unwrap(), 1024 * 1024 * 1024);
122 assert_eq!(parse_byte_size("1tb").unwrap(), 1000 * 1000 * 1000 * 1000);
123 assert_eq!(parse_byte_size("1ti").unwrap(), 1024 * 1024 * 1024 * 1024);
124 }
125
126 #[test]
127 fn test_byte_size_display() {
128 assert_eq!(ByteSizeConf(1).to_string(), "1 bytes");
129 assert_eq!(ByteSizeConf(1024).to_string(), "1 KiB");
130 assert_eq!(ByteSizeConf(1024 * 1024).to_string(), "1 MiB");
131 assert_eq!(ByteSizeConf(1024 * 1024 * 1024).to_string(), "1 GiB");
132 }
133}