human-bytesize-procmacro 0.0.0

//! A procedural macro for parsing and computing human-readable byte sizes with support for various units like KB, KiB, MB, MiB, and more.
//!
//! # Setup
//!
//! To use this crate, add the following entry to your `Cargo.toml` file in the `dependencies` section:
//!
//! ```toml
//! [dependencies]
//! human-bytesize-procmacro = "0.0.0"
//! ```
//!
//! Alternatively, you can use the [`cargo add`](https://doc.rust-lang.org/cargo/commands/cargo-add.html) subcommand:
//!
//! ```shell
//! cargo add human-bytesize-procmacro
//! ```
//!
//! # Usage
//!
//! Use the [`human_bytesize!`] macro to convert human-readable byte sizes into their byte equivalents by passing a size with its unit:
//!
//! ```rust
//! use human_bytesize_procmacro::human_bytesize;
//!
//! assert_eq!(human_bytesize!(10KB), 10 * 1000);
//! assert_eq!(human_bytesize!(16 KiB), 16 * 1024);
//! let variable = 160;
//! assert_eq!(human_bytesize!({ variable } MB), variable * 1000 * 1000);
//! ```
//!
//! # Supported units
//!
//! ## Decimal Units (Base 10)
//!
//! - B (Byte)
//! - K, KB (Kilobyte)
//! - M, MB (Megabyte)
//! - G, GB (Gigabyte)
//! - T, TB (Terabyte)
//! - P, PB (Petabyte)
//! - E, EB (Exabyte)
//! - Z, ZB (Zettabyte)
//! - Y, YB (Yottabyte)
//!
//! ## Binary Units (Base 2)
//!
//! - B (Byte)
//! - Ki, KiB (Kibibyte)
//! - Mi, MiB (Mebibyte)
//! - Gi, GiB (Gibibyte)
//! - Ti, TiB (Tebibyte)
//! - Pi, PiB (Pebibyte)
//! - Ei, EiB (Exbibyte)
//! - Zi, ZiB (Zebibyte)
//! - Yi, YiB (Yobibyte)
//!
//! # License
//!
//! This crate is licensed under the MIT License.

#![forbid(unsafe_code)]

use std::ops::Mul;

use nom::branch::alt;
use nom::bytes::complete::{is_not, tag};
use nom::character::complete::{alpha0, i128, space0};
use nom::combinator::{all_consuming, map, map_parser, map_res, opt, value};
use nom::sequence::{delimited, tuple};
use nom::{Finish, IResult};
use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_str, Error as SynError, Expr};

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Unit {
    Byte,
    Kilobyte,
    Kibibyte,
    Megabyte,
    Mebibyte,
    Gigabyte,
    Gibibyte,
    Terabyte,
    Tebibyte,
    Petabyte,
    Pebibyte,
    Exabyte,
    Exbibyte,
    Zettabyte,
    Zebibyte,
    Yottabyte,
    Yobibyte,
}

impl Unit {
    const fn multiplier(&self) -> i128 {
        match self {
            Self::Byte => 1,
            Self::Kilobyte => 1000,
            Self::Kibibyte => 1024,
            Self::Megabyte => 1000 * 1000,
            Self::Mebibyte => 1024 * 1024,
            Self::Gigabyte => 1000 * 1000 * 1000,
            Self::Gibibyte => 1024 * 1024 * 1024,
            Self::Terabyte => 1000 * 1000 * 1000 * 1000,
            Self::Tebibyte => 1024 * 1024 * 1024 * 1024,
            Self::Petabyte => 1000 * 1000 * 1000 * 1000 * 1000,
            Self::Pebibyte => 1024 * 1024 * 1024 * 1024 * 1024,
            Self::Exabyte => 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
            Self::Exbibyte => 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
            Self::Zettabyte => 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
            Self::Zebibyte => 1024 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
            Self::Yottabyte => 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
            Self::Yobibyte => 1024 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
        }
    }

    fn parse(input: &str) -> IResult<&str, Self> {
        let mut parser = alt((
            Self::parse_byte,
            Self::parse_kilobyte,
            Self::parse_kibibyte,
            Self::parse_megabyte,
            Self::parse_mebibyte,
            Self::parse_gigabyte,
            Self::parse_gibibyte,
            Self::parse_terabyte,
            Self::parse_tebibyte,
            Self::parse_petabyte,
            Self::parse_pebibyte,
            Self::parse_exabyte,
            Self::parse_exbibyte,
            Self::parse_zettabyte,
            Self::parse_zebibyte,
            Self::parse_yottabyte,
            Self::parse_yobibyte,
        ));
        parser(input)
    }

    fn parse_byte(input: &str) -> IResult<&str, Self> {
        let parser = all_consuming(opt(tag("B")));
        let mut parser = value(Self::Byte, parser);
        parser(input)
    }

    fn parse_kilobyte(input: &str) -> IResult<&str, Self> {
        let parser = alt((all_consuming(tag("K")), all_consuming(tag("KB"))));
        let mut parser = value(Self::Kilobyte, parser);
        parser(input)
    }

    fn parse_kibibyte(input: &str) -> IResult<&str, Self> {
        let parser = alt((all_consuming(tag("Ki")), all_consuming(tag("KiB"))));
        let mut parser = value(Self::Kibibyte, parser);
        parser(input)
    }

    fn parse_megabyte(input: &str) -> IResult<&str, Self> {
        let parser = alt((all_consuming(tag("M")), all_consuming(tag("MB"))));
        let mut parser = value(Self::Megabyte, parser);
        parser(input)
    }

    fn parse_mebibyte(input: &str) -> IResult<&str, Self> {
        let parser = alt((all_consuming(tag("Mi")), all_consuming(tag("MiB"))));
        let mut parser = value(Self::Mebibyte, parser);
        parser(input)
    }

    fn parse_gigabyte(input: &str) -> IResult<&str, Self> {
        let parser = alt((all_consuming(tag("G")), all_consuming(tag("GB"))));
        let mut parser = value(Self::Gigabyte, parser);
        parser(input)
    }

    fn parse_gibibyte(input: &str) -> IResult<&str, Self> {
        let parser = alt((all_consuming(tag("Gi")), all_consuming(tag("GiB"))));
        let mut parser = value(Self::Gibibyte, parser);
        parser(input)
    }

    fn parse_terabyte(input: &str) -> IResult<&str, Self> {
        let parser = alt((all_consuming(tag("T")), all_consuming(tag("TB"))));
        let mut parser = value(Self::Terabyte, parser);
        parser(input)
    }

    fn parse_tebibyte(input: &str) -> IResult<&str, Self> {
        let parser = alt((all_consuming(tag("Ti")), all_consuming(tag("TiB"))));
        let mut parser = value(Self::Tebibyte, parser);
        parser(input)
    }

    fn parse_petabyte(input: &str) -> IResult<&str, Self> {
        let parser = alt((all_consuming(tag("P")), all_consuming(tag("PB"))));
        let mut parser = value(Self::Petabyte, parser);
        parser(input)
    }

    fn parse_pebibyte(input: &str) -> IResult<&str, Self> {
        let parser = alt((all_consuming(tag("Pi")), all_consuming(tag("PiB"))));
        let mut parser = value(Self::Pebibyte, parser);
        parser(input)
    }

    fn parse_exabyte(input: &str) -> IResult<&str, Self> {
        let parser = alt((all_consuming(tag("E")), all_consuming(tag("EB"))));
        let mut parser = value(Self::Exabyte, parser);
        parser(input)
    }

    fn parse_exbibyte(input: &str) -> IResult<&str, Self> {
        let parser = alt((all_consuming(tag("Ei")), all_consuming(tag("EiB"))));
        let mut parser = value(Self::Exbibyte, parser);
        parser(input)
    }

    fn parse_zettabyte(input: &str) -> IResult<&str, Self> {
        let parser = alt((all_consuming(tag("Z")), all_consuming(tag("ZB"))));
        let mut parser = value(Self::Zettabyte, parser);
        parser(input)
    }

    fn parse_zebibyte(input: &str) -> IResult<&str, Self> {
        let parser = alt((all_consuming(tag("Zi")), all_consuming(tag("ZiB"))));
        let mut parser = value(Self::Zebibyte, parser);
        parser(input)
    }

    fn parse_yottabyte(input: &str) -> IResult<&str, Self> {
        let parser = alt((all_consuming(tag("Y")), all_consuming(tag("YB"))));
        let mut parser = value(Self::Yottabyte, parser);
        parser(input)
    }

    fn parse_yobibyte(input: &str) -> IResult<&str, Self> {
        let parser = alt((all_consuming(tag("Yi")), all_consuming(tag("YiB"))));
        let mut parser = value(Self::Yobibyte, parser);
        parser(input)
    }
}

macro_rules! impl_mul {
    ($($t:ty),+ => $o:ty) => {
        $(
            impl Mul<Unit> for $t {
                type Output = $o;

                fn mul(self, rhs: Unit) -> Self::Output {
                    (self as Self::Output) * (rhs.multiplier() as Self::Output)
                }
            }
        )*
    };
}
impl_mul!(u8, u16, u32, u64, u128, usize => u128);
impl_mul!(i8, i16, i32, i64, i128, isize => i128);

enum Value {
    Number(i128),
    Expression(Expr),
}

impl Value {
    fn parse(input: &str) -> IResult<&str, Self> {
        let number_parser = map(i128, Self::Number);
        let identifier_parser = {
            let parser = delimited(tag("{"), is_not("}"), tag("}"));
            map_res(parser, |expr| -> Result<Self, SynError> {
                let expr = parse_str::<Expr>(expr)?;
                let expression = Self::Expression(expr);
                Ok(expression)
            })
        };
        let mut parser = alt((number_parser, identifier_parser));
        parser(input)
    }
}

fn parse_human_bytesize(input: &str) -> IResult<&str, (Value, Unit)> {
    let parser = tuple((space0, Value::parse, space0, map_parser(alpha0, Unit::parse), space0));
    let parser = all_consuming(parser);
    let mut parser = map(parser, |(_, value, _, unit, _)| (value, unit)); // .unwrap_or(Unit::Byte)
    parser(input)
}

/// A procedural macro for converting human-readable byte size expressions into their corresponding byte values at compile time.
///
/// - **Macro evaluation**: Numeric literals are evaluated directly by the macro at compile time, producing a constant value.
/// - **Compiler evaluation**: Expressions are expanded into code, leaving the final computation to the Rust compiler at compile time.
///
/// Returns a value of type [`prim@i128`] or a multiplication expression that evaluates to [`prim@i128`], representing the computed size in bytes.
///
/// # Usage
///
/// - **Numeric literals**: Use numbers followed by a unit, e.g., `100KB` or `16 KiB`.
/// - **Expressions**: Enclose a valid Rust expression in curly braces, followed by the unit.
///
/// # Example
///
/// ```rust
/// use human_bytesize_procmacro::human_bytesize;
///
/// // Using numeric literals
/// assert_eq!(human_bytesize!(-4B), -4);
/// assert_eq!(human_bytesize!(10KB), 10 * 1000);
/// assert_eq!(human_bytesize!(16 KiB), 16 * 1024);
///
/// // Using expressions
/// let variable = 320;
/// assert_eq!(human_bytesize!({ variable } MiB), variable * 1024 * 1024);
///
/// const TOTAL_MEGABYTES: i128 = 16;
/// assert_eq!(human_bytesize!({ TOTAL_MEGABYTES } MB), TOTAL_MEGABYTES * 1000 * 1000);
/// ```
///
/// # Limitations
///
/// ## Exponent Notation Conflict
///
/// The `E` unit (e.g., `100EB` for exabytes) may conflict with Rust's scientific notation parser, which expects an exponent after `E`.
///
/// For example:
///
/// ```rust,compile_fail
/// # use human_bytesize_procmacro::human_bytesize;
/// let x = human_bytesize!(100EB); // This results in a parsing error.
/// ```
///
/// To work around this limitation, use a space between the number and the unit:
///
/// ```rust
/// # use human_bytesize_procmacro::human_bytesize;
/// let x = human_bytesize!(100 EB); // This compiles correctly.
/// ```
///
/// ## Error Handling
///
/// The macro will panic if the input format is invalid. Ensure you use one of the following formats:
///
/// - `<number><unit>`: Example: `100KB`, `16 KiB`
/// - `{<expression>}<unit>`: Example: `{variable}MB`, `{ CONST_VAL } G`
///
/// ## Common Pitfalls
///
/// - **Spacing**: While both `100KB` and `100 KB` are valid, ensure consistency in formatting.
/// - **Expression Wrapping**: When using a variable or an expression, wrap it in `{}` before specifying the unit.
#[proc_macro]
pub fn human_bytesize(input: TokenStream) -> TokenStream {
    let input = input.to_string();
    match parse_human_bytesize(&input).finish() {
        Ok((_, (Value::Number(value), unit))) => {
            let result = value * unit;
            let result = quote! {
                #result
            };
            TokenStream::from(result)
        },
        Ok((_, (Value::Expression(expr), unit))) => {
            let unit = unit.multiplier();
            let result = quote! {
                ((#expr) * #unit)
            };
            TokenStream::from(result)
        },
        Err(_) => panic!("Invalid format! Please use format like '100 KiB' or '100KiB'"),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_unit_parse() {
        // Valid use cases
        assert!(matches!(Unit::parse(""), Ok((_, Unit::Byte))));
        assert!(matches!(Unit::parse("B"), Ok((_, Unit::Byte))));
        assert!(matches!(Unit::parse("K"), Ok((_, Unit::Kilobyte))));
        assert!(matches!(Unit::parse("KB"), Ok((_, Unit::Kilobyte))));
        assert!(matches!(Unit::parse("Ki"), Ok((_, Unit::Kibibyte))));
        assert!(matches!(Unit::parse("KiB"), Ok((_, Unit::Kibibyte))));
        assert!(matches!(Unit::parse("M"), Ok((_, Unit::Megabyte))));
        assert!(matches!(Unit::parse("MB"), Ok((_, Unit::Megabyte))));
        assert!(matches!(Unit::parse("Mi"), Ok((_, Unit::Mebibyte))));
        assert!(matches!(Unit::parse("MiB"), Ok((_, Unit::Mebibyte))));
        assert!(matches!(Unit::parse("G"), Ok((_, Unit::Gigabyte))));
        assert!(matches!(Unit::parse("GB"), Ok((_, Unit::Gigabyte))));
        assert!(matches!(Unit::parse("Gi"), Ok((_, Unit::Gibibyte))));
        assert!(matches!(Unit::parse("GiB"), Ok((_, Unit::Gibibyte))));
        assert!(matches!(Unit::parse("T"), Ok((_, Unit::Terabyte))));
        assert!(matches!(Unit::parse("TB"), Ok((_, Unit::Terabyte))));
        assert!(matches!(Unit::parse("Ti"), Ok((_, Unit::Tebibyte))));
        assert!(matches!(Unit::parse("TiB"), Ok((_, Unit::Tebibyte))));
        assert!(matches!(Unit::parse("P"), Ok((_, Unit::Petabyte))));
        assert!(matches!(Unit::parse("PB"), Ok((_, Unit::Petabyte))));
        assert!(matches!(Unit::parse("Pi"), Ok((_, Unit::Pebibyte))));
        assert!(matches!(Unit::parse("PiB"), Ok((_, Unit::Pebibyte))));
        assert!(matches!(Unit::parse("E"), Ok((_, Unit::Exabyte))));
        assert!(matches!(Unit::parse("EB"), Ok((_, Unit::Exabyte))));
        assert!(matches!(Unit::parse("Ei"), Ok((_, Unit::Exbibyte))));
        assert!(matches!(Unit::parse("EiB"), Ok((_, Unit::Exbibyte))));
        assert!(matches!(Unit::parse("Z"), Ok((_, Unit::Zettabyte))));
        assert!(matches!(Unit::parse("ZB"), Ok((_, Unit::Zettabyte))));
        assert!(matches!(Unit::parse("Zi"), Ok((_, Unit::Zebibyte))));
        assert!(matches!(Unit::parse("ZiB"), Ok((_, Unit::Zebibyte))));
        assert!(matches!(Unit::parse("Y"), Ok((_, Unit::Yottabyte))));
        assert!(matches!(Unit::parse("YB"), Ok((_, Unit::Yottabyte))));
        assert!(matches!(Unit::parse("Yi"), Ok((_, Unit::Yobibyte))));
        assert!(matches!(Unit::parse("YiB"), Ok((_, Unit::Yobibyte))));

        // Invalid use cases
        assert!(matches!(Unit::parse(" "), Err(_)));
        assert!(matches!(Unit::parse(" B"), Err(_)));
        assert!(matches!(Unit::parse("   B"), Err(_)));
        assert!(matches!(Unit::parse("B "), Err(_)));
        assert!(matches!(Unit::parse("B       "), Err(_)));
        assert!(matches!(Unit::parse("XYZ"), Err(_)));
    }

    #[test]
    fn test_parse_human_bytesize() {
        // Valid use cases
        assert!(matches!(
            parse_human_bytesize("8"),
            Ok((_, (Value::Number(8), Unit::Byte)))
        ));
        assert!(matches!(
            parse_human_bytesize("34 B"),
            Ok((_, (Value::Number(34), Unit::Byte)))
        ));
        assert!(matches!(
            parse_human_bytesize("82KB"),
            Ok((_, (Value::Number(82), Unit::Kilobyte)))
        ));
        assert!(matches!(
            parse_human_bytesize("7 KiB"),
            Ok((_, (Value::Number(7), Unit::Kibibyte)))
        ));
        assert!(matches!(
            parse_human_bytesize("987 MB"),
            Ok((_, (Value::Number(987), Unit::Megabyte)))
        ));
        assert!(matches!(
            parse_human_bytesize("150MiB"),
            Ok((_, (Value::Number(150), Unit::Mebibyte)))
        ));
        assert!(matches!(
            parse_human_bytesize("99 GB"),
            Ok((_, (Value::Number(99), Unit::Gigabyte)))
        ));
        assert!(matches!(
            parse_human_bytesize("1 GiB"),
            Ok((_, (Value::Number(1), Unit::Gibibyte)))
        ));
        assert!(matches!(
            parse_human_bytesize("678TB"),
            Ok((_, (Value::Number(678), Unit::Terabyte)))
        ));
        assert!(matches!(
            parse_human_bytesize("123TiB"),
            Ok((_, (Value::Number(123), Unit::Tebibyte)))
        ));
        assert!(matches!(
            parse_human_bytesize("54 PB"),
            Ok((_, (Value::Number(54), Unit::Petabyte)))
        ));
        assert!(matches!(
            parse_human_bytesize("19 PiB"),
            Ok((_, (Value::Number(19), Unit::Pebibyte)))
        ));
        assert!(matches!(
            parse_human_bytesize("556 EB"),
            Ok((_, (Value::Number(556), Unit::Exabyte)))
        ));
        assert!(matches!(
            parse_human_bytesize("153EiB"),
            Ok((_, (Value::Number(153), Unit::Exbibyte)))
        ));
        assert!(matches!(
            parse_human_bytesize("4ZB"),
            Ok((_, (Value::Number(4), Unit::Zettabyte)))
        ));
        assert!(matches!(
            parse_human_bytesize("34ZiB"),
            Ok((_, (Value::Number(34), Unit::Zebibyte)))
        ));
        assert!(matches!(
            parse_human_bytesize("750YB"),
            Ok((_, (Value::Number(750), Unit::Yottabyte)))
        ));
        assert!(matches!(
            parse_human_bytesize("56YiB"),
            Ok((_, (Value::Number(56), Unit::Yobibyte)))
        ));
        assert!(matches!(
            parse_human_bytesize("{ var } EB"),
            Ok((_, (Value::Expression(_), Unit::Exabyte)))
        ));
        assert!(matches!(
            parse_human_bytesize("0  "),
            Ok((_, (Value::Number(0), Unit::Byte)))
        ));
        assert!(matches!(
            parse_human_bytesize("   13M"),
            Ok((_, (Value::Number(13), Unit::Megabyte)))
        ));
        assert!(matches!(
            parse_human_bytesize("230  G"),
            Ok((_, (Value::Number(230), Unit::Gigabyte)))
        ));
        assert!(matches!(
            parse_human_bytesize(" 2 PiB "),
            Ok((_, (Value::Number(2), Unit::Pebibyte)))
        ));

        // Invalid use cases
        assert!(matches!(parse_human_bytesize("18BB"), Err(_)));
        assert!(matches!(parse_human_bytesize("1 0TB"), Err(_)));
        assert!(matches!(parse_human_bytesize("10 XB"), Err(_)));
        assert!(matches!(parse_human_bytesize("ABC XYZ"), Err(_)));
    }
}