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 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}