config_types/
memory.rs

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}