bytelike_derive/
lib.rs

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    // Define units with their multipliers and descriptions
29    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    // Generate methods
44    let methods = units.iter().map(|(fn_name, multiplier, description)| {
45        // Create an identifier for the method name
46        let method_name = syn::Ident::new(fn_name, Span::call_site());
47
48        // Parse the multiplier into an expression
49        let multiplier_expr: syn::Expr = syn::parse_str(multiplier).unwrap();
50
51        // Generate the documentation comment
52        let doc_comment = format!("Construct `{}` given an amount of {}.", name, description);
53
54        // Generate the method using quote!
55        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            /// Provides `ByteLikeRange` with explicit lower and upper bounds.
247            pub fn range<I: Into<Self>>(start: I, stop: I) -> ::bytelike::ByteLikeRange<Self> {
248                ::bytelike::ByteLikeRange::new(Some(start), Some(stop))
249            }
250
251            /// Provides `ByteLikeRange` with explicit lower bound. Upper bound is set to `u64::MAX`.
252            pub fn range_start<I: Into<Self>>(start: I) -> ::bytelike::ByteLikeRange<Self> {
253                ::bytelike::ByteLikeRange::new(Some(start), None)
254            }
255
256            /// Provides `ByteLikeRange` with explicit lower bound. Upper bound is set to `u64::MAX`.
257            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            /// Returns the size as a string with an optional SI unit.
335            #[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            /// Returns the inner u64 value.
341            #[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}