luau_printf/

arg.rs

use super::printf_impl::Error;
use bstr::{BStr, BString, ByteSlice};
use std::result::Result;

/// Printf argument types.
/// Note no implementation of `ToArg` constructs the owned string variant;
/// callers can do so explicitly.
#[derive(Debug, PartialEq)]
pub enum Arg<'a> {
    Str(&'a BStr),
    String(BString),
    UInt(u64),
    SInt(i64),
    Float(f64),
    USizeRef(&'a mut usize), // for use with %n
}

impl<'a> Arg<'a> {
    pub fn string(bytes: &'a BStr) -> Self {
        Self::Str(bytes)
    }

    pub fn int(value: i32) -> Self {
        Self::SInt(i64::from(value))
    }

    pub fn sint(value: i64) -> Self {
        Self::SInt(value)
    }

    pub fn hex(value: u32) -> Self {
        Self::UInt(u64::from(value))
    }

    pub fn uint(value: u64) -> Self {
        Self::UInt(value)
    }

    pub fn float(value: f64) -> Self {
        Self::Float(value)
    }

    pub fn pointer(value: usize) -> Self {
        Self::UInt(value as u64)
    }

    pub fn set_count(&mut self, count: usize) -> Result<(), Error> {
        match self {
            Arg::USizeRef(p) => **p = count,
            _ => return Err(Error::BadArgType),
        }
        Ok(())
    }

    pub fn as_bstr(&self) -> Result<&BStr, Error> {
        match self {
            Arg::Str(s) => Ok(s),
            Arg::String(s) => Ok(s.as_bstr()),
            _ => Err(Error::BadArgType),
        }
    }

    // Return this value as an unsigned integer. Negative signed values will report overflow.
    pub fn as_uint(&self) -> Result<u64, Error> {
        match *self {
            Arg::UInt(u) => Ok(u),
            Arg::SInt(i) => i.try_into().map_err(|_| Error::Overflow),
            _ => Err(Error::BadArgType),
        }
    }

    // Return this value as a signed integer. Unsigned values > i64::MAX will report overflow.
    pub fn as_sint(&self) -> Result<i64, Error> {
        match *self {
            Arg::UInt(u) => u.try_into().map_err(|_| Error::Overflow),
            Arg::SInt(i) => Ok(i),
            _ => Err(Error::BadArgType),
        }
    }

    /// Unwraps [`Arg::UInt`] to [`u64`].
    /// Unwraps [`Arg::SInt`] and casts the [`i64`] to [`u64`].
    /// Calling this on other variants of `[Arg]` is an error.
    pub fn as_wrapping_sint(&self) -> Result<u64, Error> {
        match *self {
            Arg::UInt(u) => Ok(u),
            Arg::SInt(i) => Ok(i as u64),
            _ => Err(Error::BadArgType),
        }
    }

    // Note we allow passing ints as floats, even allowing precision loss.
    pub fn as_float(&self) -> Result<f64, Error> {
        #[allow(clippy::cast_precision_loss)]
        match *self {
            Arg::Float(f) => Ok(f),
            Arg::UInt(u) => Ok(u as f64),
            Arg::SInt(i) => Ok(i as f64),
            _ => Err(Error::BadArgType),
        }
    }

    pub fn as_uchar(&self) -> Result<u8, Error> {
        Ok(self.as_wrapping_sint()? as u8)
    }
}

/// Conversion from a raw value to a printf argument.
pub trait ToArg<'a> {
    fn to_arg(self) -> Arg<'a>;
}

impl<'a> ToArg<'a> for &'a str {
    fn to_arg(self) -> Arg<'a> {
        Arg::Str(self.as_bytes().as_bstr())
    }
}

impl<'a> ToArg<'a> for &'a String {
    fn to_arg(self) -> Arg<'a> {
        Arg::Str(self.as_bytes().as_bstr())
    }
}

impl<'a> ToArg<'a> for String {
    fn to_arg(self) -> Arg<'a> {
        Arg::String(BString::from(self))
    }
}

impl<'a> ToArg<'a> for &'a BStr {
    fn to_arg(self) -> Arg<'a> {
        Arg::Str(self)
    }
}

impl<'a> ToArg<'a> for &'a BString {
    fn to_arg(self) -> Arg<'a> {
        Arg::Str(self.as_bstr())
    }
}

impl<'a> ToArg<'a> for BString {
    fn to_arg(self) -> Arg<'a> {
        Arg::String(self)
    }
}

impl<'a> ToArg<'a> for &'a [u8] {
    fn to_arg(self) -> Arg<'a> {
        Arg::Str(self.as_bstr())
    }
}

impl<'a> ToArg<'a> for &'a Vec<u8> {
    fn to_arg(self) -> Arg<'a> {
        Arg::Str(self.as_slice().as_bstr())
    }
}

impl<'a> ToArg<'a> for Vec<u8> {
    fn to_arg(self) -> Arg<'a> {
        Arg::String(BString::from(self))
    }
}

impl<'a, const N: usize> ToArg<'a> for &'a [u8; N] {
    fn to_arg(self) -> Arg<'a> {
        Arg::Str(self.as_bstr())
    }
}

impl<'a> ToArg<'a> for Arg<'a> {
    fn to_arg(self) -> Arg<'a> {
        self
    }
}

impl<'a> ToArg<'a> for &'a std::io::Error {
    fn to_arg(self) -> Arg<'a> {
        Arg::String(BString::from(self.to_string()))
    }
}

impl<'a> ToArg<'a> for f32 {
    fn to_arg(self) -> Arg<'a> {
        Arg::Float(self.into())
    }
}

impl<'a> ToArg<'a> for f64 {
    fn to_arg(self) -> Arg<'a> {
        Arg::Float(self)
    }
}

impl<'a> ToArg<'a> for &'a mut usize {
    fn to_arg(self) -> Arg<'a> {
        Arg::USizeRef(self)
    }
}

impl<'a, T> ToArg<'a> for &'a *const T {
    fn to_arg(self) -> Arg<'a> {
        Arg::UInt((*self) as usize as u64)
    }
}

/// All signed types.
macro_rules! impl_to_arg {
    ($($t:ty),*) => {
        $(
            impl<'a> ToArg<'a> for $t {
                fn to_arg(self) -> Arg<'a> {
                    Arg::SInt(self as i64)
                }
            }
        )*
    };
}
impl_to_arg!(i8, i16, i32, i64, isize);

/// All unsigned types.
macro_rules! impl_to_arg_u {
    ($($t:ty),*) => {
        $(
            impl<'a> ToArg<'a> for $t {
                fn to_arg(self) -> Arg<'a> {
                    Arg::UInt(self as u64)
                }
            }
        )*
    };
}
impl_to_arg_u!(u8, u16, u32, u64, usize);

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

    #[test]
    fn test_to_arg() {
        assert!(matches!("test".to_arg(), Arg::Str(_)));
        assert!(matches!((&String::from("test")).to_arg(), Arg::Str(_)));
        assert!(matches!(String::from("test").to_arg(), Arg::String(_)));
        assert!(matches!(b"test".to_arg(), Arg::Str(_)));
        assert!(matches!(b"test".as_slice().to_arg(), Arg::Str(_)));
        assert!(matches!((&BString::from("test")).to_arg(), Arg::Str(_)));
        assert!(matches!(BString::from("test").to_arg(), Arg::String(_)));
        assert!(matches!((&b"test".to_vec()).to_arg(), Arg::Str(_)));
        assert!(matches!(b"test".to_vec().to_arg(), Arg::String(_)));
        assert!(matches!(42f32.to_arg(), Arg::Float(_)));
        assert!(matches!(42f64.to_arg(), Arg::Float(_)));
        let mut usize_val: usize = 0;
        assert!(matches!((&mut usize_val).to_arg(), Arg::USizeRef(_)));
        assert!(matches!(42i8.to_arg(), Arg::SInt(42)));
        assert!(matches!(42i16.to_arg(), Arg::SInt(42)));
        assert!(matches!(42i32.to_arg(), Arg::SInt(42)));
        assert!(matches!(42i64.to_arg(), Arg::SInt(42)));
        assert!(matches!(42isize.to_arg(), Arg::SInt(42)));

        assert_eq!((-42i8).to_arg(), Arg::SInt(-42));
        assert_eq!((-42i16).to_arg(), Arg::SInt(-42));
        assert_eq!((-42i32).to_arg(), Arg::SInt(-42));
        assert_eq!((-42i64).to_arg(), Arg::SInt(-42));
        assert_eq!((-42isize).to_arg(), Arg::SInt(-42));

        assert!(matches!(42u8.to_arg(), Arg::UInt(42)));
        assert!(matches!(42u16.to_arg(), Arg::UInt(42)));
        assert!(matches!(42u32.to_arg(), Arg::UInt(42)));
        assert!(matches!(42u64.to_arg(), Arg::UInt(42)));
        assert!(matches!(42usize.to_arg(), Arg::UInt(42)));

        let ptr = std::ptr::from_ref(&42f32);
        assert!(matches!(ptr.to_arg(), Arg::UInt(_)));
    }

    #[test]
    fn test_negative_to_arg() {
        assert_eq!((-1_i8).to_arg().as_sint(), Ok(-1));
        assert_eq!((-1_i16).to_arg().as_sint(), Ok(-1));
        assert_eq!((-1_i32).to_arg().as_sint(), Ok(-1));
        assert_eq!((-1_i64).to_arg().as_sint(), Ok(-1));

        assert_eq!((u64::MAX).to_arg().as_sint(), Err(Error::Overflow));
    }
}