config_types/
memory.rs

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