1use alloc::{borrow::ToOwned as _, format, string::String};
2use core::{fmt, str};
3
4use super::ByteSize;
5
6impl str::FromStr for ByteSize {
7 type Err = String;
8
9 fn from_str(value: &str) -> Result<Self, Self::Err> {
10 if let Ok(v) = value.parse::<u64>() {
11 return Ok(Self(v));
12 }
13 let number = take_while(value, |c| c.is_ascii_digit() || c == '.');
14 match number.parse::<f64>() {
15 Ok(v) => {
16 let suffix = skip_while(&value[number.len()..], char::is_whitespace);
17 match suffix.parse::<Unit>() {
18 Ok(u) if !number.contains('.') => match number.parse::<u64>() {
23 Ok(n) => Ok(Self(n.saturating_mul(u.factor()))),
24 Err(_) => Ok(Self((v * u) as u64)),
25 },
26 Ok(u) => Ok(Self((v * u) as u64)),
27 Err(error) => Err(format!(
28 "couldn't parse {suffix:?} into a known SI unit, {error}"
29 )),
30 }
31 }
32 Err(error) => Err(format!("couldn't parse {value:?} into a ByteSize, {error}")),
33 }
34 }
35}
36
37fn take_while<P>(s: &str, mut predicate: P) -> &str
38where
39 P: FnMut(char) -> bool,
40{
41 let offset = s
42 .chars()
43 .take_while(|ch| predicate(*ch))
44 .map(|ch| ch.len_utf8())
45 .sum();
46 &s[..offset]
47}
48
49fn skip_while<P>(s: &str, mut predicate: P) -> &str
50where
51 P: FnMut(char) -> bool,
52{
53 let offset: usize = s
54 .chars()
55 .skip_while(|ch| predicate(*ch))
56 .map(|ch| ch.len_utf8())
57 .sum();
58 &s[(s.len() - offset)..]
59}
60
61#[non_exhaustive]
74#[derive(Debug, Clone, PartialEq)]
75pub enum Unit {
76 Byte,
78
79 KiloByte,
82
83 MegaByte,
85
86 GigaByte,
88
89 TeraByte,
91
92 PetaByte,
94
95 ExaByte,
97
98 KibiByte,
101
102 MebiByte,
104
105 GibiByte,
107
108 TebiByte,
110
111 PebiByte,
113
114 ExbiByte,
116}
117
118impl Unit {
119 fn factor(&self) -> u64 {
120 match self {
121 Self::Byte => 1,
122 Self::KiloByte => crate::KB,
124 Self::MegaByte => crate::MB,
125 Self::GigaByte => crate::GB,
126 Self::TeraByte => crate::TB,
127 Self::PetaByte => crate::PB,
128 Self::ExaByte => crate::EB,
129 Self::KibiByte => crate::KIB,
131 Self::MebiByte => crate::MIB,
132 Self::GibiByte => crate::GIB,
133 Self::TebiByte => crate::TIB,
134 Self::PebiByte => crate::PIB,
135 Self::ExbiByte => crate::EIB,
136 }
137 }
138}
139
140mod impl_ops {
141 use super::Unit;
142 use core::ops;
143
144 impl ops::Add<u64> for Unit {
145 type Output = u64;
146
147 fn add(self, other: u64) -> Self::Output {
148 self.factor() + other
149 }
150 }
151
152 impl ops::Add<Unit> for u64 {
153 type Output = u64;
154
155 fn add(self, other: Unit) -> Self::Output {
156 self + other.factor()
157 }
158 }
159
160 impl ops::Mul<u64> for Unit {
161 type Output = u64;
162
163 fn mul(self, other: u64) -> Self::Output {
164 self.factor() * other
165 }
166 }
167
168 impl ops::Mul<Unit> for u64 {
169 type Output = u64;
170
171 fn mul(self, other: Unit) -> Self::Output {
172 self * other.factor()
173 }
174 }
175
176 impl ops::Add<f64> for Unit {
177 type Output = f64;
178
179 fn add(self, other: f64) -> Self::Output {
180 self.factor() as f64 + other
181 }
182 }
183
184 impl ops::Add<Unit> for f64 {
185 type Output = f64;
186
187 fn add(self, other: Unit) -> Self::Output {
188 other.factor() as f64 + self
189 }
190 }
191
192 impl ops::Mul<f64> for Unit {
193 type Output = f64;
194
195 fn mul(self, other: f64) -> Self::Output {
196 self.factor() as f64 * other
197 }
198 }
199
200 impl ops::Mul<Unit> for f64 {
201 type Output = f64;
202
203 fn mul(self, other: Unit) -> Self::Output {
204 other.factor() as f64 * self
205 }
206 }
207}
208
209impl str::FromStr for Unit {
210 type Err = UnitParseError;
211
212 fn from_str(unit: &str) -> Result<Self, Self::Err> {
213 match () {
214 _ if unit.eq_ignore_ascii_case("b") => Ok(Self::Byte),
215 _ if unit.eq_ignore_ascii_case("k") | unit.eq_ignore_ascii_case("kb") => {
216 Ok(Self::KiloByte)
217 }
218 _ if unit.eq_ignore_ascii_case("m") | unit.eq_ignore_ascii_case("mb") => {
219 Ok(Self::MegaByte)
220 }
221 _ if unit.eq_ignore_ascii_case("g") | unit.eq_ignore_ascii_case("gb") => {
222 Ok(Self::GigaByte)
223 }
224 _ if unit.eq_ignore_ascii_case("t") | unit.eq_ignore_ascii_case("tb") => {
225 Ok(Self::TeraByte)
226 }
227 _ if unit.eq_ignore_ascii_case("p") | unit.eq_ignore_ascii_case("pb") => {
228 Ok(Self::PetaByte)
229 }
230 _ if unit.eq_ignore_ascii_case("e") | unit.eq_ignore_ascii_case("eb") => {
231 Ok(Self::ExaByte)
232 }
233 _ if unit.eq_ignore_ascii_case("ki") | unit.eq_ignore_ascii_case("kib") => {
234 Ok(Self::KibiByte)
235 }
236 _ if unit.eq_ignore_ascii_case("mi") | unit.eq_ignore_ascii_case("mib") => {
237 Ok(Self::MebiByte)
238 }
239 _ if unit.eq_ignore_ascii_case("gi") | unit.eq_ignore_ascii_case("gib") => {
240 Ok(Self::GibiByte)
241 }
242 _ if unit.eq_ignore_ascii_case("ti") | unit.eq_ignore_ascii_case("tib") => {
243 Ok(Self::TebiByte)
244 }
245 _ if unit.eq_ignore_ascii_case("pi") | unit.eq_ignore_ascii_case("pib") => {
246 Ok(Self::PebiByte)
247 }
248 _ if unit.eq_ignore_ascii_case("ei") | unit.eq_ignore_ascii_case("eib") => {
249 Ok(Self::ExbiByte)
250 }
251 _ => Err(UnitParseError(to_string_truncate(unit))),
252 }
253 }
254}
255
256fn to_string_truncate(unit: &str) -> String {
258 const MAX_UNIT_LEN: usize = 3;
259
260 if unit.len() > MAX_UNIT_LEN {
261 if unit.is_char_boundary(3) {
264 format!("{}...", &unit[..3])
265 } else if unit.is_char_boundary(4) {
266 format!("{}...", &unit[..4])
267 } else if unit.is_char_boundary(5) {
268 format!("{}...", &unit[..5])
269 } else if unit.is_char_boundary(6) {
270 format!("{}...", &unit[..6])
271 } else {
272 unreachable!("char boundary will be within 4 bytes")
273 }
274 } else {
275 unit.to_owned()
276 }
277}
278
279#[derive(Debug)]
281pub struct UnitParseError(String);
282
283impl fmt::Display for UnitParseError {
284 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
285 write!(f, "Failed to parse unit \"{}\"", self.0)
286 }
287}
288
289#[cfg(feature = "std")]
290impl std::error::Error for UnitParseError {}
291
292#[cfg(test)]
293mod tests {
294 use alloc::string::ToString as _;
295
296 use super::*;
297
298 #[test]
299 fn truncating_error_strings() {
300 assert_eq!("", to_string_truncate(""));
301 assert_eq!("b", to_string_truncate("b"));
302 assert_eq!("ob", to_string_truncate("ob"));
303 assert_eq!("foo", to_string_truncate("foo"));
304
305 assert_eq!("foo...", to_string_truncate("foob"));
306 assert_eq!("foo...", to_string_truncate("foobar"));
307 }
308
309 #[test]
310 fn when_ok() {
311 fn parse(s: &str) -> u64 {
313 s.parse::<ByteSize>().unwrap().0
314 }
315
316 assert_eq!("0".parse::<ByteSize>().unwrap().0, 0);
317 assert_eq!(parse("0"), 0);
318 assert_eq!(parse("500"), 500);
319 assert_eq!(parse("1K"), Unit::KiloByte * 1);
320 assert_eq!(parse("1Ki"), Unit::KibiByte * 1);
321 assert_eq!(parse("1.5Ki"), (1.5 * Unit::KibiByte) as u64);
322 assert_eq!(parse("1KiB"), 1 * Unit::KibiByte);
323 assert_eq!(parse("1.5KiB"), (1.5 * Unit::KibiByte) as u64);
324 assert_eq!(parse("3 MB"), Unit::MegaByte * 3);
325 assert_eq!(parse("4 MiB"), Unit::MebiByte * 4);
326 assert_eq!(parse("6 GB"), 6 * Unit::GigaByte);
327 assert_eq!(parse("4 GiB"), 4 * Unit::GibiByte);
328 assert_eq!(parse("88TB"), 88 * Unit::TeraByte);
329 assert_eq!(parse("521TiB"), 521 * Unit::TebiByte);
330 assert_eq!(parse("8 PB"), 8 * Unit::PetaByte);
331 assert_eq!(parse("8P"), 8 * Unit::PetaByte);
332 assert_eq!(parse("12 PiB"), 12 * Unit::PebiByte);
333 }
334
335 #[test]
336 fn large_integer_byte_counts_are_exact() {
337 fn parse(s: &str) -> u64 {
339 s.parse::<ByteSize>().unwrap().0
340 }
341
342 assert_eq!(parse("9007199254740993B"), 9_007_199_254_740_993);
346 assert_eq!(parse("9007199254740993"), parse("9007199254740993B"));
347 }
348
349 #[test]
350 fn when_err() {
351 fn parse(s: &str) -> Result<ByteSize, String> {
353 s.parse::<ByteSize>()
354 }
355
356 assert!(parse("").is_err());
357 assert!(parse("a124GB").is_err());
358 assert!(parse("1.3 42.0 B").is_err());
359 assert!(parse("1.3 ... B").is_err());
360 assert!(parse("1 000 B").is_err());
363 }
364
365 #[test]
366 fn to_and_from_str() {
367 fn parse(s: &str) -> u64 {
369 s.parse::<ByteSize>().unwrap().0
370 }
371
372 assert_eq!(parse(&parse("128GB").to_string()), 128 * Unit::GigaByte);
373 assert_eq!(
374 parse(&ByteSize(parse("128.000 GiB")).to_string()),
375 128 * Unit::GibiByte,
376 );
377 }
378}