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