| 1 | use std::{fmt, str::FromStr, time::Duration}; |
| 2 | |
| 3 | use serde::{Deserialize, Deserializer, Serialize, Serializer, de}; |
| 4 | |
| 5 | const KIB: u64 = 1024; |
| 6 | |
| 7 | #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] |
| 8 | pub struct Bytes(pub u64); |
| 9 | |
| 10 | impl Bytes { |
| 11 | pub const ZERO: Self = Self(0); |
| 12 | |
| 13 | pub const fn gib(n: u64) -> Self { |
| 14 | Self(n * KIB * KIB * KIB) |
| 15 | } |
| 16 | |
| 17 | pub const fn mib(n: u64) -> Self { |
| 18 | Self(n * KIB * KIB) |
| 19 | } |
| 20 | |
| 21 | pub const fn saturating_sub(self, other: Self) -> Self { |
| 22 | Self(self.0.saturating_sub(other.0)) |
| 23 | } |
| 24 | } |
| 25 | |
| 26 | impl std::ops::Add for Bytes { |
| 27 | type Output = Self; |
| 28 | |
| 29 | fn add(self, other: Self) -> Self { |
| 30 | Self(self.0.saturating_add(other.0)) |
| 31 | } |
| 32 | } |
| 33 | |
| 34 | impl std::ops::AddAssign for Bytes { |
| 35 | fn add_assign(&mut self, other: Self) { |
| 36 | *self = *self + other; |
| 37 | } |
| 38 | } |
| 39 | |
| 40 | impl std::iter::Sum for Bytes { |
| 41 | fn sum<I: Iterator<Item = Self>>(iter: I) -> Self { |
| 42 | iter.fold(Self::ZERO, |a, b| a + b) |
| 43 | } |
| 44 | } |
| 45 | |
| 46 | fn split_number(s: &str) -> Result<(f64, &str), &'static str> { |
| 47 | let s = s.trim(); |
| 48 | let end = s |
| 49 | .find(|c: char| !(c.is_ascii_digit() || c == '.')) |
| 50 | .unwrap_or(s.len()); |
| 51 | let (num, suffix) = s.split_at(end); |
| 52 | let num = num.parse::<f64>().map_err(|_| "expected a number")?; |
| 53 | if !num.is_finite() || num < 0.0 { |
| 54 | return Err("expected a non-negative number"); |
| 55 | } |
| 56 | Ok((num, suffix.trim())) |
| 57 | } |
| 58 | |
| 59 | impl FromStr for Bytes { |
| 60 | type Err = &'static str; |
| 61 | |
| 62 | fn from_str(s: &str) -> Result<Self, Self::Err> { |
| 63 | let (num, suffix) = split_number(s)?; |
| 64 | let scale = match suffix.to_ascii_lowercase().as_str() { |
| 65 | "" | "b" => 1, |
| 66 | "k" | "kb" | "kib" => KIB, |
| 67 | "m" | "mb" | "mib" => KIB.pow(2), |
| 68 | "g" | "gb" | "gib" => KIB.pow(3), |
| 69 | "t" | "tb" | "tib" => KIB.pow(4), |
| 70 | _ => return Err("expected a size like 512M or 4GiB"), |
| 71 | }; |
| 72 | Ok(Self((num * scale as f64) as u64)) |
| 73 | } |
| 74 | } |
| 75 | |
| 76 | impl fmt::Display for Bytes { |
| 77 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
| 78 | const UNITS: [&str; 5] = ["B", "KiB", "MiB", "GiB", "TiB"]; |
| 79 | let mut value = self.0 as f64; |
| 80 | let mut unit = 0; |
| 81 | while value >= 1024.0 && unit < UNITS.len() - 1 { |
| 82 | value /= 1024.0; |
| 83 | unit += 1; |
| 84 | } |
| 85 | if unit == 0 { |
| 86 | write!(f, "{}B", self.0) |
| 87 | } else { |
| 88 | write!(f, "{value:.1}{}", UNITS[unit]) |
| 89 | } |
| 90 | } |
| 91 | } |
| 92 | |
| 93 | impl pound::FromArg for Bytes { |
| 94 | fn from_arg(s: &str) -> Result<Self, pound::ValueError> { |
| 95 | s.parse().map_err(|e| pound::ValueError::new(s, e)) |
| 96 | } |
| 97 | } |
| 98 | |
| 99 | impl Serialize for Bytes { |
| 100 | fn serialize<S: Serializer>(&self, s: S) -> Result<S::Ok, S::Error> { |
| 101 | s.serialize_u64(self.0) |
| 102 | } |
| 103 | } |
| 104 | |
| 105 | impl<'de> Deserialize<'de> for Bytes { |
| 106 | fn deserialize<D: Deserializer<'de>>(d: D) -> Result<Self, D::Error> { |
| 107 | struct V; |
| 108 | impl de::Visitor<'_> for V { |
| 109 | type Value = Bytes; |
| 110 | |
| 111 | fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
| 112 | f.write_str("a byte count or a size string like \"4GiB\"") |
| 113 | } |
| 114 | |
| 115 | fn visit_u64<E: de::Error>(self, v: u64) -> Result<Bytes, E> { |
| 116 | Ok(Bytes(v)) |
| 117 | } |
| 118 | |
| 119 | fn visit_i64<E: de::Error>(self, v: i64) -> Result<Bytes, E> { |
| 120 | u64::try_from(v) |
| 121 | .map(Bytes) |
| 122 | .map_err(|_| E::custom("negative size")) |
| 123 | } |
| 124 | |
| 125 | fn visit_str<E: de::Error>(self, v: &str) -> Result<Bytes, E> { |
| 126 | v.parse().map_err(E::custom) |
| 127 | } |
| 128 | } |
| 129 | d.deserialize_any(V) |
| 130 | } |
| 131 | } |
| 132 | |
| 133 | #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord)] |
| 134 | pub struct Dur(pub Duration); |
| 135 | |
| 136 | impl Dur { |
| 137 | pub const fn secs(n: u64) -> Self { |
| 138 | Self(Duration::from_secs(n)) |
| 139 | } |
| 140 | } |
| 141 | |
| 142 | impl FromStr for Dur { |
| 143 | type Err = &'static str; |
| 144 | |
| 145 | fn from_str(s: &str) -> Result<Self, Self::Err> { |
| 146 | let mut rest = s.trim(); |
| 147 | if rest.is_empty() { |
| 148 | return Err("expected a duration like 90s or 5m"); |
| 149 | } |
| 150 | let mut total = 0.0; |
| 151 | while !rest.is_empty() { |
| 152 | let (num, tail) = split_number(rest)?; |
| 153 | let unit_len = tail |
| 154 | .find(|c: char| c.is_ascii_digit() || c == '.') |
| 155 | .unwrap_or(tail.len()); |
| 156 | let (unit, next) = tail.split_at(unit_len); |
| 157 | let scale = match unit.trim() { |
| 158 | "" | "s" | "sec" | "secs" => 1.0, |
| 159 | "ms" => 0.001, |
| 160 | "m" | "min" | "mins" => 60.0, |
| 161 | "h" | "hr" | "hrs" => 3600.0, |
| 162 | "d" | "day" | "days" => 86400.0, |
| 163 | _ => return Err("expected a duration like 90s or 5m"), |
| 164 | }; |
| 165 | total = num.mul_add(scale, total); |
| 166 | rest = next.trim(); |
| 167 | } |
| 168 | Ok(Self(Duration::from_secs_f64(total))) |
| 169 | } |
| 170 | } |
| 171 | |
| 172 | impl fmt::Display for Dur { |
| 173 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
| 174 | let s = self.0.as_secs(); |
| 175 | match s { |
| 176 | 0..60 => write!(f, "{s}s"), |
| 177 | 60..3600 if s.is_multiple_of(60) => write!(f, "{}m", s / 60), |
| 178 | 60..3600 => write!(f, "{}m{}s", s / 60, s % 60), |
| 179 | _ if s.is_multiple_of(3600) => write!(f, "{}h", s / 3600), |
| 180 | _ => write!(f, "{}h{}m", s / 3600, (s % 3600) / 60), |
| 181 | } |
| 182 | } |
| 183 | } |
| 184 | |
| 185 | impl pound::FromArg for Dur { |
| 186 | fn from_arg(s: &str) -> Result<Self, pound::ValueError> { |
| 187 | s.parse().map_err(|e| pound::ValueError::new(s, e)) |
| 188 | } |
| 189 | } |
| 190 | |
| 191 | impl Serialize for Dur { |
| 192 | fn serialize<S: Serializer>(&self, s: S) -> Result<S::Ok, S::Error> { |
| 193 | s.serialize_u64(self.0.as_secs()) |
| 194 | } |
| 195 | } |
| 196 | |
| 197 | impl<'de> Deserialize<'de> for Dur { |
| 198 | fn deserialize<D: Deserializer<'de>>(d: D) -> Result<Self, D::Error> { |
| 199 | struct V; |
| 200 | impl de::Visitor<'_> for V { |
| 201 | type Value = Dur; |
| 202 | |
| 203 | fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
| 204 | f.write_str("seconds or a duration string like \"5m\"") |
| 205 | } |
| 206 | |
| 207 | fn visit_u64<E: de::Error>(self, v: u64) -> Result<Dur, E> { |
| 208 | Ok(Dur::secs(v)) |
| 209 | } |
| 210 | |
| 211 | fn visit_i64<E: de::Error>(self, v: i64) -> Result<Dur, E> { |
| 212 | u64::try_from(v) |
| 213 | .map(Dur::secs) |
| 214 | .map_err(|_| E::custom("negative duration")) |
| 215 | } |
| 216 | |
| 217 | fn visit_str<E: de::Error>(self, v: &str) -> Result<Dur, E> { |
| 218 | v.parse().map_err(E::custom) |
| 219 | } |
| 220 | } |
| 221 | d.deserialize_any(V) |
| 222 | } |
| 223 | } |
| 224 | |
| 225 | #[derive(Clone, Copy, Debug, PartialEq)] |
| 226 | pub enum Amount<T> { |
| 227 | Abs(T), |
| 228 | Percent(f64), |
| 229 | } |
| 230 | |
| 231 | impl Amount<Bytes> { |
| 232 | pub fn resolve(self, total: Bytes) -> Bytes { |
| 233 | match self { |
| 234 | Self::Abs(b) => b, |
| 235 | Self::Percent(p) => Bytes((total.0 as f64 * p / 100.0) as u64), |
| 236 | } |
| 237 | } |
| 238 | } |
| 239 | |
| 240 | impl Amount<u32> { |
| 241 | pub fn resolve(self, total: u32) -> u32 { |
| 242 | match self { |
| 243 | Self::Abs(n) => n, |
| 244 | Self::Percent(p) => ((f64::from(total) * p / 100.0).round() as u32).max(1), |
| 245 | } |
| 246 | } |
| 247 | } |
| 248 | |
| 249 | impl<T: FromStr> FromStr for Amount<T> { |
| 250 | type Err = &'static str; |
| 251 | |
| 252 | fn from_str(s: &str) -> Result<Self, Self::Err> { |
| 253 | if let Some(p) = s.trim().strip_suffix('%') { |
| 254 | let p = p |
| 255 | .trim() |
| 256 | .parse::<f64>() |
| 257 | .map_err(|_| "expected a percentage")?; |
| 258 | if !(0.0..=100.0).contains(&p) { |
| 259 | return Err("percentage must be between 0 and 100"); |
| 260 | } |
| 261 | Ok(Self::Percent(p)) |
| 262 | } else { |
| 263 | s.trim() |
| 264 | .parse() |
| 265 | .map(Self::Abs) |
| 266 | .map_err(|_| "expected an amount or a percentage") |
| 267 | } |
| 268 | } |
| 269 | } |
| 270 | |
| 271 | impl<T: FromStr> pound::FromArg for Amount<T> { |
| 272 | fn from_arg(s: &str) -> Result<Self, pound::ValueError> { |
| 273 | s.parse().map_err(|e| pound::ValueError::new(s, e)) |
| 274 | } |
| 275 | } |
| 276 | |
| 277 | impl<'de, T: FromStr> Deserialize<'de> for Amount<T> { |
| 278 | fn deserialize<D: Deserializer<'de>>(d: D) -> Result<Self, D::Error> { |
| 279 | struct V<T>(std::marker::PhantomData<T>); |
| 280 | impl<T: FromStr> de::Visitor<'_> for V<T> { |
| 281 | type Value = Amount<T>; |
| 282 | |
| 283 | fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
| 284 | f.write_str("an amount or a percentage string like \"75%\"") |
| 285 | } |
| 286 | |
| 287 | fn visit_u64<E: de::Error>(self, v: u64) -> Result<Amount<T>, E> { |
| 288 | self.visit_str(&v.to_string()) |
| 289 | } |
| 290 | |
| 291 | fn visit_i64<E: de::Error>(self, v: i64) -> Result<Amount<T>, E> { |
| 292 | self.visit_str(&v.to_string()) |
| 293 | } |
| 294 | |
| 295 | fn visit_str<E: de::Error>(self, v: &str) -> Result<Amount<T>, E> { |
| 296 | v.parse().map_err(E::custom) |
| 297 | } |
| 298 | } |
| 299 | d.deserialize_any(V(std::marker::PhantomData)) |
| 300 | } |
| 301 | } |