Skip to main content

acorn/schema/hardware/
memory.rs

1//! Data model for harware memory resources, including memory units and (de)serialization logic
2use crate::error::{AcornError, AcornResult};
3use crate::prelude::*;
4use alloc::borrow::Cow;
5use core::{convert::TryInto, str::FromStr};
6use derive_more::Display;
7use schemars::JsonSchema;
8use serde::{Deserialize, Serialize};
9
10/// Memory unit for hardware resources
11#[derive(Clone, Copy, Debug, Display, PartialEq, Serialize)]
12pub enum MemoryUnit {
13    /// Gigabytes
14    GB,
15    /// Kilobytes
16    KB,
17    /// Megabytes
18    MB,
19    /// Terabytes
20    TB,
21}
22/// Memory amount with explicit unit
23///
24/// Supports string formats like `"80GB"`, `"64KB"`, `"512MB"`, `"1TB"` and plain number (defaults to GB).
25#[derive(Clone, Debug, PartialEq)]
26pub struct Memory {
27    /// Numeric amount of memory
28    pub amount: f64,
29    /// Unit of measurement
30    pub unit: MemoryUnit,
31}
32impl<'de> Deserialize<'de> for Memory {
33    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
34    where
35        D: serde::de::Deserializer<'de>,
36    {
37        struct MemoryVisitor;
38
39        impl<'de> serde::de::Visitor<'de> for MemoryVisitor {
40            type Value = Memory;
41
42            fn expecting(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
43                f.write_str(r#"a memory string (e.g. "80GB", "2.5GB", "512MB") or a number (treated as GB)"#)
44            }
45            fn visit_u64<E: serde::de::Error>(self, value: u64) -> Result<Memory, E> {
46                memory_from_number(value as f64).map_err(E::custom)
47            }
48            fn visit_i64<E: serde::de::Error>(self, value: i64) -> Result<Memory, E> {
49                memory_from_number(value as f64).map_err(E::custom)
50            }
51            fn visit_f64<E: serde::de::Error>(self, value: f64) -> Result<Memory, E> {
52                memory_from_number(value).map_err(E::custom)
53            }
54            fn visit_str<E: serde::de::Error>(self, value: &str) -> Result<Memory, E> {
55                value.parse().map_err(E::custom)
56            }
57        }
58        deserializer.deserialize_any(MemoryVisitor)
59    }
60}
61impl JsonSchema for Memory {
62    fn schema_name() -> Cow<'static, str> {
63        "Memory".into()
64    }
65    fn json_schema(_gen: &mut schemars::generate::SchemaGenerator) -> schemars::Schema {
66        #[allow(clippy::unwrap_used)]
67        serde_json::json!({"type": "string", "pattern": "^\\d+(\\.\\d+)?\\s*(GB|KB|MB|TB)$"})
68            .try_into()
69            .unwrap()
70    }
71    fn inline_schema() -> bool {
72        true
73    }
74}
75impl Serialize for Memory {
76    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
77        let s = format!("{}{}", self.amount, self.unit);
78        serializer.serialize_str(&s)
79    }
80}
81impl Memory {
82    /// Convert the memory amount to bytes using binary unit multipliers
83    ///
84    /// Returns `None` when the amount is not finite, is negative, or exceeds `u64`
85    pub fn checked_bytes(&self) -> Option<u64> {
86        let multiplier = match self.unit {
87            | MemoryUnit::KB => 1024_f64,
88            | MemoryUnit::MB => 1024_f64.powi(2),
89            | MemoryUnit::GB => 1024_f64.powi(3),
90            | MemoryUnit::TB => 1024_f64.powi(4),
91        };
92        let bytes = self.amount * multiplier;
93        (bytes.is_finite() && bytes >= 0.0 && bytes < u64::MAX as f64)
94            .then(|| format!("{bytes:.0}").parse::<u64>().ok())
95            .flatten()
96    }
97    /// Determine whether a byte count fits in this memory amount
98    pub fn can_contain(&self, bytes: u64) -> Option<bool> {
99        self.checked_bytes().map(|available| bytes <= available)
100    }
101    /// Create a Memory value in GB
102    pub fn gb(amount: impl Into<f64>) -> Self {
103        Memory {
104            amount: amount.into(),
105            unit: MemoryUnit::GB,
106        }
107    }
108    /// Create a Memory value in KB
109    pub fn kb(amount: impl Into<f64>) -> Self {
110        Memory {
111            amount: amount.into(),
112            unit: MemoryUnit::KB,
113        }
114    }
115    /// Create a Memory value in MB
116    pub fn mb(amount: impl Into<f64>) -> Self {
117        Memory {
118            amount: amount.into(),
119            unit: MemoryUnit::MB,
120        }
121    }
122    /// Create a Memory value in TB
123    pub fn tb(amount: impl Into<f64>) -> Self {
124        Memory {
125            amount: amount.into(),
126            unit: MemoryUnit::TB,
127        }
128    }
129}
130impl FromStr for Memory {
131    type Err = AcornError;
132    fn from_str(value: &str) -> AcornResult<Self> {
133        parse_memory_string(value)
134    }
135}
136impl<'de> Deserialize<'de> for MemoryUnit {
137    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
138    where
139        D: serde::de::Deserializer<'de>,
140    {
141        struct MemoryUnitVisitor;
142        impl<'de> serde::de::Visitor<'de> for MemoryUnitVisitor {
143            type Value = MemoryUnit;
144
145            fn expecting(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
146                f.write_str("a memory unit string (e.g. 'GB', 'KB', 'MB', 'TB')")
147            }
148            fn visit_str<E: serde::de::Error>(self, value: &str) -> Result<MemoryUnit, E> {
149                match value.trim().to_uppercase().as_str() {
150                    | "GB" | "G" | "GIB" => Ok(MemoryUnit::GB),
151                    | "KB" | "K" | "KIB" => Ok(MemoryUnit::KB),
152                    | "MB" | "M" | "MIB" => Ok(MemoryUnit::MB),
153                    | "TB" | "T" | "TIB" => Ok(MemoryUnit::TB),
154                    | other => Err(serde::de::Error::custom(format!("Invalid memory unit: '{other}'"))),
155                }
156            }
157        }
158        deserializer.deserialize_str(MemoryUnitVisitor)
159    }
160}
161impl JsonSchema for MemoryUnit {
162    fn schema_name() -> alloc::borrow::Cow<'static, str> {
163        "MemoryUnit".into()
164    }
165    fn json_schema(_gen: &mut schemars::generate::SchemaGenerator) -> schemars::Schema {
166        #[allow(clippy::unwrap_used)]
167        serde_json::json!({"type": "string", "enum": ["GB", "KB", "MB", "TB"]})
168            .try_into()
169            .unwrap()
170    }
171    fn inline_schema() -> bool {
172        true
173    }
174}
175fn memory_from_number(amount: f64) -> AcornResult<Memory> {
176    if !amount.is_finite() {
177        Err(AcornError::new("Memory amount must be finite"))
178    } else if amount < 0.0 {
179        Err(AcornError::new("Memory amount cannot be negative"))
180    } else {
181        Ok(Memory {
182            amount,
183            unit: MemoryUnit::GB,
184        })
185    }
186}
187fn parse_memory_string(value: &str) -> AcornResult<Memory> {
188    let s = value.trim();
189    match s.find(|c: char| !c.is_ascii_digit() && c != '.') {
190        | Some(split) => match (s.get(..split), s.get(split..)) {
191            | (Some(value), Some(unit)) => match value.trim().parse::<f64>() {
192                | Ok(amount) => memory_from_number(amount).and_then(|_| {
193                    MemoryUnit::deserialize(serde::de::value::StrDeserializer::<serde::de::value::Error>::new(unit.trim()))
194                        .map(|unit| Memory { amount, unit })
195                        .map_err(|why| AcornError::new(why.to_string()))
196                }),
197                | Err(_) => Err(AcornError::new(format!("Invalid memory amount — '{value}'"))),
198            },
199            | _ => Err(AcornError::new(format!("Invalid memory value — '{s}'"))),
200        },
201        | None => Err(AcornError::new(format!("Missing unit in memory value — '{s}'"))),
202    }
203}
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208
209    #[test]
210    fn test_memory_from_str_and_serde_share_parsing() {
211        let parsed = "1.5GB".parse::<Memory>().unwrap();
212        let deserialized = serde_json::from_str::<Memory>(r#""1.5GiB""#).unwrap();
213        assert_eq!(parsed, deserialized);
214        assert_eq!(parsed.checked_bytes(), Some(1_610_612_736));
215    }
216    #[test]
217    fn test_memory_binary_aliases_are_equivalent() {
218        let gb = "24GB".parse::<Memory>().unwrap();
219        let gib = "24GiB".parse::<Memory>().unwrap();
220        assert_eq!(gb.checked_bytes(), gib.checked_bytes());
221        assert_eq!(gb.can_contain(24 * 1024 * 1024 * 1024), Some(true));
222    }
223    #[test]
224    fn test_memory_rejects_invalid_values_and_checked_overflow() {
225        assert!("24XB".parse::<Memory>().is_err());
226        assert!("-1GB".parse::<Memory>().is_err());
227        assert_eq!(Memory::tb(f64::MAX).checked_bytes(), None);
228    }
229}