1use proc_macro::TokenStream;
2use proc_macro2::Span;
3use quote::quote;
4use syn::{parse_macro_input, DeriveInput};
5
6#[proc_macro_derive(ByteLike)]
7pub fn bytelike(input: TokenStream) -> TokenStream {
8 let input_str = input.to_string();
9 let constructor = bytelike_constructor(input_str.parse().unwrap());
10 let display = bytelike_display(input_str.parse().unwrap());
11 let parse = bytelike_parse(input_str.parse().unwrap());
12 let ops = bytelike_ops(input_str.parse().unwrap());
13 let fromstr = bytelike_fromstr(input_str.parse().unwrap());
14
15 let mut combined = format!("{}{}{}{}{}", constructor, display, parse, ops, fromstr);
16 if cfg!(feature = "serde") {
17 let serde = bytelike_serde(input_str.parse().unwrap());
18 combined = format!("{}{}", combined, serde);
19 }
20 combined.parse().unwrap()
21}
22
23#[proc_macro_derive(ByteLikeConstructor)]
24pub fn bytelike_constructor(input: TokenStream) -> TokenStream {
25 let input = parse_macro_input!(input as DeriveInput);
26 let name = &input.ident;
27
28 let units = vec![
30 ("b", "1", "bytes"),
31 ("kb", "bytelike::KB", "kilobytes"),
32 ("kib", "bytelike::KIB", "kibibytes"),
33 ("mb", "bytelike::MB", "megabytes"),
34 ("mib", "bytelike::MIB", "mebibytes"),
35 ("gb", "bytelike::GB", "gigabytes"),
36 ("gib", "bytelike::GIB", "gibibytes"),
37 ("tb", "bytelike::TB", "terabytes"),
38 ("tib", "bytelike::TIB", "tebibytes"),
39 ("pb", "bytelike::PB", "petabytes"),
40 ("pib", "bytelike::PIB", "pebibytes"),
41 ];
42
43 let methods = units.iter().map(|(fn_name, multiplier, description)| {
45 let method_name = syn::Ident::new(fn_name, Span::call_site());
47
48 let multiplier_expr: syn::Expr = syn::parse_str(multiplier).unwrap();
50
51 let doc_comment = format!("Construct `{}` given an amount of {}.", name, description);
53
54 quote! {
56 #[doc = #doc_comment]
57 #[inline(always)]
58 pub const fn #method_name(size: u64) -> Self {
59 Self(size * #multiplier_expr)
60 }
61 }
62 });
63
64 let expanded = quote! {
65 impl #name {
66 #(#methods)*
67 }
68
69 impl From<u64> for #name {
70 fn from(size: u64) -> #name {
71 Self(size)
72 }
73 }
74 };
75
76 TokenStream::from(expanded)
77}
78
79#[proc_macro_derive(ByteLikeOps)]
80pub fn bytelike_ops(input: TokenStream) -> TokenStream {
81 let input = parse_macro_input!(input as DeriveInput);
82 let name = &input.ident;
83
84 let expanded = quote! {
85 impl core::ops::Add<#name> for #name {
86 type Output = #name;
87
88 #[inline(always)]
89 fn add(self, rhs: #name) -> #name {
90 #name(self.0 + rhs.0)
91 }
92 }
93
94 impl core::ops::AddAssign<#name> for #name {
95 #[inline(always)]
96 fn add_assign(&mut self, rhs: #name) {
97 self.0 += rhs.0
98 }
99 }
100
101 impl<T> core::ops::Add<T> for #name
102 where
103 T: Into<u64>,
104 {
105 type Output = #name;
106 #[inline(always)]
107 fn add(self, rhs: T) -> #name {
108 #name(self.0 + (rhs.into()))
109 }
110 }
111
112 impl<T> core::ops::AddAssign<T> for #name
113 where
114 T: Into<u64>,
115 {
116 #[inline(always)]
117 fn add_assign(&mut self, rhs: T) {
118 self.0 += rhs.into();
119 }
120 }
121
122 impl core::ops::Sub<#name> for #name {
123 type Output = #name;
124
125 #[inline(always)]
126 fn sub(self, rhs: #name) -> #name {
127 #name(self.0 - rhs.0)
128 }
129 }
130
131 impl core::ops::SubAssign<#name> for #name {
132 #[inline(always)]
133 fn sub_assign(&mut self, rhs: #name) {
134 self.0 -= rhs.0
135 }
136 }
137
138 impl<T> core::ops::Sub<T> for #name
139 where
140 T: Into<u64>,
141 {
142 type Output = #name;
143
144 #[inline(always)]
145 fn sub(self, rhs: T) -> #name {
146 #name(self.0 - (rhs.into()))
147 }
148 }
149
150 impl<T> core::ops::SubAssign<T> for #name
151 where
152 T: Into<u64>,
153 {
154 #[inline(always)]
155 fn sub_assign(&mut self, rhs: T) {
156 self.0 -= rhs.into();
157 }
158 }
159
160 impl<T> core::ops::Mul<T> for #name
161 where
162 T: Into<u64>,
163 {
164 type Output = #name;
165 #[inline(always)]
166 fn mul(self, rhs: T) -> #name {
167 #name(self.0 * rhs.into())
168 }
169 }
170
171 impl<T> core::ops::MulAssign<T> for #name
172 where
173 T: Into<u64>,
174 {
175 #[inline(always)]
176 fn mul_assign(&mut self, rhs: T) {
177 self.0 *= rhs.into();
178 }
179 }
180
181 impl core::ops::Add<#name> for u64 {
182 type Output = #name;
183 #[inline(always)]
184 fn add(self, rhs: #name) -> #name {
185 #name(rhs.0 + self)
186 }
187 }
188
189 impl core::ops::Add<#name> for u32 {
190 type Output = #name;
191 #[inline(always)]
192 fn add(self, rhs: #name) -> #name {
193 #name(rhs.0 + (self as u64))
194 }
195 }
196
197 impl core::ops::Add<#name> for u16 {
198 type Output = #name;
199 #[inline(always)]
200 fn add(self, rhs: #name) -> #name {
201 #name(rhs.0 + (self as u64))
202 }
203 }
204
205 impl core::ops::Add<#name> for u8 {
206 type Output = #name;
207 #[inline(always)]
208 fn add(self, rhs: #name) -> #name {
209 #name(rhs.0 + (self as u64))
210 }
211 }
212
213 impl core::ops::Mul<#name> for u64 {
214 type Output = #name;
215 #[inline(always)]
216 fn mul(self, rhs: #name) -> #name {
217 #name(rhs.0 * self)
218 }
219 }
220
221 impl core::ops::Mul<#name> for u32 {
222 type Output = #name;
223 #[inline(always)]
224 fn mul(self, rhs: #name) -> #name {
225 #name(rhs.0 * (self as u64))
226 }
227 }
228
229 impl core::ops::Mul<#name> for u16 {
230 type Output = #name;
231 #[inline(always)]
232 fn mul(self, rhs: #name) -> #name {
233 #name(rhs.0 * (self as u64))
234 }
235 }
236
237 impl core::ops::Mul<#name> for u8 {
238 type Output = #name;
239 #[inline(always)]
240 fn mul(self, rhs: #name) -> #name {
241 #name(rhs.0 * (self as u64))
242 }
243 }
244
245 impl #name {
246 pub fn range<I: Into<Self>>(start: I, stop: I) -> ::bytelike::ByteLikeRange<Self> {
248 ::bytelike::ByteLikeRange::new(Some(start), Some(stop))
249 }
250
251 pub fn range_start<I: Into<Self>>(start: I) -> ::bytelike::ByteLikeRange<Self> {
253 ::bytelike::ByteLikeRange::new(Some(start), None)
254 }
255
256 pub fn range_stop<I: Into<Self>>(stop: I) -> ::bytelike::ByteLikeRange<Self> {
258 ::bytelike::ByteLikeRange::new(None, Some(stop.into()))
259 }
260 }
261 };
262
263 TokenStream::from(expanded)
264}
265
266#[proc_macro_derive(ByteLikeDisplay)]
267pub fn bytelike_display(input: TokenStream) -> TokenStream {
268 let input = parse_macro_input!(input as DeriveInput);
269 let name = &input.ident;
270
271 let expanded = quote! {
272 impl core::fmt::Display for #name {
273 fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
274 f.pad(&bytelike::to_string(self.0, true))
275 }
276 }
277
278 impl core::fmt::Debug for #name {
279 fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
280 write!(f, "{}", self)
281 }
282 }
283 };
284
285 TokenStream::from(expanded)
286}
287
288#[proc_macro_derive(ByteLikeFromStr)]
289pub fn bytelike_fromstr(input: TokenStream) -> TokenStream {
290 let input = parse_macro_input!(input as DeriveInput);
291 let name = &input.ident;
292
293 let expanded = quote! {
294 impl core::str::FromStr for #name {
295 type Err = ::bytelike::String;
296
297 fn from_str(value: &str) -> core::result::Result<Self, Self::Err> {
298 if let Ok(v) = value.parse::<u64>() {
299 return Ok(Self(v));
300 }
301 let number = ::bytelike::take_while(value, |c| c.is_ascii_digit() || c == '.');
302 match number.parse::<f64>() {
303 Ok(v) => {
304 let suffix = ::bytelike::skip_while(value, |c| {
305 c.is_whitespace() || c.is_ascii_digit() || c == '.'
306 });
307 match suffix.parse::<::bytelike::Unit>() {
308 Ok(u) => Ok(Self((v * u64::from(u) as f64) as u64)),
309 Err(error) => Err(::bytelike::format!(
310 "couldn't parse {:?} into a known SI unit, {}",
311 suffix, error
312 )),
313 }
314 }
315 Err(error) => Err(::bytelike::format!(
316 "couldn't parse {:?} into a ByteSize, {}",
317 value, error
318 )),
319 }
320 }
321 }
322 };
323
324 TokenStream::from(expanded)
325}
326
327#[proc_macro_derive(ByteLikeParse)]
328pub fn bytelike_parse(input: TokenStream) -> TokenStream {
329 let input = parse_macro_input!(input as DeriveInput);
330 let name = &input.ident;
331
332 let expanded = quote! {
333 impl #name {
334 #[inline(always)]
336 pub fn to_string_as(&self, si_unit: bool) -> ::bytelike::String {
337 ::bytelike::to_string(self.0, si_unit)
338 }
339
340 #[inline(always)]
342 pub const fn as_u64(&self) -> u64 {
343 self.0
344 }
345 }
346 };
347
348 TokenStream::from(expanded)
349}
350
351#[proc_macro_derive(ByteLikeSerde)]
352pub fn bytelike_serde(input: TokenStream) -> TokenStream {
353 let input = parse_macro_input!(input as DeriveInput);
354 let name = &input.ident;
355
356 let expanded = quote! {
357 impl<'de> ::bytelike::serde::Deserialize<'de> for #name {
358 fn deserialize<D>(deserializer: D) -> core::result::Result<Self, D::Error>
359 where
360 D: ::bytelike::serde::Deserializer<'de>,
361 {
362 struct ByteSizeVistor;
363
364 impl<'de> ::bytelike::serde::de::Visitor<'de> for ByteSizeVistor {
365 type Value = #name;
366
367 fn expecting(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
368 formatter.write_str("an integer or string")
369 }
370
371 fn visit_i64<E: ::bytelike::serde::de::Error>(self, value: i64) -> core::result::Result<Self::Value, E> {
372 if let Ok(val) = u64::try_from(value) {
373 Ok(#name(val))
374 } else {
375 Err(E::invalid_value(
376 ::bytelike::serde::de::Unexpected::Signed(value),
377 &"integer overflow",
378 ))
379 }
380 }
381
382 fn visit_u64<E: ::bytelike::serde::de::Error>(self, value: u64) -> core::result::Result<Self::Value, E> {
383 Ok(#name(value))
384 }
385
386 fn visit_str<E: ::bytelike::serde::de::Error>(self, value: &str) -> core::result::Result<Self::Value, E> {
387 if let Ok(val) = value.parse() {
388 Ok(val)
389 } else {
390 Err(E::invalid_value(
391 ::bytelike::serde::de::Unexpected::Str(value),
392 &"parsable string",
393 ))
394 }
395 }
396 }
397
398 if deserializer.is_human_readable() {
399 deserializer.deserialize_any(ByteSizeVistor)
400 } else {
401 deserializer.deserialize_u64(ByteSizeVistor)
402 }
403 }
404 }
405 impl ::bytelike::serde::Serialize for #name {
406 fn serialize<S>(&self, serializer: S) -> core::result::Result<S::Ok, S::Error>
407 where
408 S: ::bytelike::serde::Serializer,
409 {
410 if serializer.is_human_readable() {
411 <str>::serialize(self.to_string().as_str(), serializer)
412 } else {
413 self.0.serialize(serializer)
414 }
415 }
416 }
417 };
418
419 TokenStream::from(expanded)
420}