luau-printf 0.732.0

/** Luau printf-compatible implementation, based on musl. */
use super::arg::Arg;
use super::fmt_fp::format_float;
use super::locale::Locale;
use bstr::{BStr, ByteSlice as _};
use std::io::{self, Write as IoWrite};
use std::mem;
use std::result::Result;

/// Possible errors from printf.
#[derive(Debug, PartialEq, Eq)]
pub enum Error {
    /// Invalid format string.
    BadFormatString,
    /// Too few arguments.
    MissingArg,
    /// Argument type doesn't match format specifier.
    BadArgType,
    /// Precision is too large to represent.
    Overflow,
    /// Error emitted by the output stream.
    Io(io::ErrorKind),
}

// Convenience conversion from io::Error.
impl From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::Io(err.kind())
    }
}

#[derive(Debug, Copy, Clone, Default)]
pub(super) struct ModifierFlags {
    pub alt_form: bool, // #
    pub zero_pad: bool, // 0
    pub left_adj: bool, // negative field width
    pub pad_pos: bool,  // space: blank before positive numbers
    pub mark_pos: bool, // +: sign before positive numbers
    pub grouped: bool,  // ': group indicator
}

impl ModifierFlags {
    // If c is a modifier byte, set the flag and return true.
    // Otherwise return false. Note we allow repeated modifier flags.
    fn try_set(&mut self, c: u8) -> bool {
        match c {
            b'#' => self.alt_form = true,
            b'0' => self.zero_pad = true,
            b'-' => self.left_adj = true,
            b' ' => self.pad_pos = true,
            b'+' => self.mark_pos = true,
            b'\'' => self.grouped = true,
            _ => return false,
        }
        true
    }
}

// The set of prefixes of conversion specifiers.
// Note that we mostly ignore prefixes - we take sizes of values from the arguments themselves.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
#[allow(non_camel_case_types)]
enum ConversionPrefix {
    Empty,
    hh,
    h,
    l,
    ll,
    j,
    t,
    z,
    L,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
#[allow(non_camel_case_types)]
#[rustfmt::skip]
pub(super) enum ConversionSpec {
    // Integers, with prefixes "hh", "h", "l", "ll", "j", "t", "z"
    // Note that we treat '%i' as '%d'.
    d, o, u, x, X,

    // USizeRef receiver, with same prefixes as ints
    n,

    // Float, with prefixes "l" and "L"
    a, A, e, E, f, F, g, G,

    // Pointer, no prefixes
    p,

    // Narrow byte or C string.
    c, s,
}

impl ConversionSpec {
    // Returns true if the given prefix is supported by this conversion specifier.
    fn supports_prefix(self, prefix: ConversionPrefix) -> bool {
        use ConversionPrefix::*;
        use ConversionSpec::*;
        if matches!(prefix, Empty) {
            // No prefix is always supported.
            return true;
        }
        match self {
            d | o | u | x | X | n => matches!(prefix, hh | h | l | ll | j | t | z),
            a | A | e | E | f | F | g | G => matches!(prefix, l | L),
            p => false,
            c | s => false,
        }
    }

    // Returns true if the conversion specifier is lowercase,
    // which affects certain rendering.
    #[inline]
    pub(super) fn is_lower(self) -> bool {
        use ConversionSpec::*;
        match self {
            d | o | u | x | n | a | e | f | g | p | c | s => true,
            X | A | E | F | G => false,
        }
    }

    // Returns a ConversionSpec from a byte, or None if none.
    fn from_byte(cc: u8) -> Option<Self> {
        use ConversionSpec::*;
        let res = match cc {
            b'd' | b'i' => d,
            b'o' => o,
            b'u' => u,
            b'x' => x,
            b'X' => X,
            b'n' => n,
            b'a' => a,
            b'A' => A,
            b'e' => e,
            b'E' => E,
            b'f' => f,
            b'F' => F,
            b'g' => g,
            b'G' => G,
            b'p' => p,
            b'c' => c,
            b's' => s,
            _ => return None,
        };
        Some(res)
    }
}

trait FormatString<'a> {
    fn is_empty(&self) -> bool;
    fn at(&self, index: usize) -> Option<u8>;
    fn advance_by(&mut self, n: usize);
    fn take_literal(&mut self) -> &'a BStr;
}

impl<'a> FormatString<'a> for &'a BStr {
    fn is_empty(&self) -> bool {
        self.len() == 0
    }

    fn at(&self, index: usize) -> Option<u8> {
        self.get(index).copied()
    }

    fn advance_by(&mut self, n: usize) {
        debug_assert!(
            n <= self.len(),
            "FormatString::advance_by(): index out of bounds"
        );
        *self = self[n..].as_bstr();
    }

    fn take_literal(&mut self) -> &'a BStr {
        let non_percents: usize = self.iter().take_while(|&&c| c != b'%').count();
        // Take only an even number of percents. Note we know these have byte length 1.
        let percent_pairs = self[non_percents..]
            .iter()
            .take_while(|&&c| c == b'%')
            .count()
            / 2;
        let (prefix, rest) = self.split_at(non_percents + percent_pairs * 2);
        *self = rest.as_bstr();
        // Trim half of the trailing percent characters from the prefix.
        prefix[..prefix.len() - percent_pairs].as_bstr()
    }
}

// Read an int from a format string, stopping at the first non-digit.
// Negative values are not supported.
// If there are no digits, return 0.
// Adjust the format string to point to the char after the int.
fn get_int<'a>(fmt: &mut impl FormatString<'a>) -> Result<usize, Error> {
    use Error::Overflow;
    let mut i: usize = 0;
    while let Some(digit) = fmt.at(0).and_then(|c| {
        if c.is_ascii_digit() {
            Some(c - b'0')
        } else {
            None
        }
    }) {
        i = i.checked_mul(10).ok_or(Overflow)?;
        i = i.checked_add(usize::from(digit)).ok_or(Overflow)?;
        fmt.advance_by(1);
    }
    Ok(i)
}

// Read a conversion prefix from a format string, advancing it.
fn get_prefix<'a>(fmt: &mut impl FormatString<'a>) -> ConversionPrefix {
    use ConversionPrefix as CP;
    let prefix = match fmt.at(0).unwrap_or(b'\0') {
        b'h' if fmt.at(1) == Some(b'h') => CP::hh,
        b'h' => CP::h,
        b'l' if fmt.at(1) == Some(b'l') => CP::ll,
        b'l' => CP::l,
        b'j' => CP::j,
        b't' => CP::t,
        b'z' => CP::z,
        b'L' => CP::L,
        _ => CP::Empty,
    };
    fmt.advance_by(match prefix {
        CP::Empty => 0,
        CP::hh | CP::ll => 2,
        _ => 1,
    });
    prefix
}

// Read an (optionally prefixed) format specifier, such as d, Lf, etc.
// Adjust the cursor to point to the char after the specifier.
fn get_specifier<'a>(fmt: &mut impl FormatString<'a>) -> Result<ConversionSpec, Error> {
    let prefix = get_prefix(fmt);
    let spec = fmt
        .at(0)
        .and_then(ConversionSpec::from_byte)
        .ok_or(Error::BadFormatString)?;
    if !spec.supports_prefix(prefix) {
        return Err(Error::BadFormatString);
    }
    fmt.advance_by(1);
    Ok(spec)
}

fn c_string_prefix(fmt: &BStr) -> &BStr {
    let len = fmt.iter().position(|&c| c == b'\0').unwrap_or(fmt.len());
    fmt[..len].as_bstr()
}

fn check_printf_count(count: usize) -> Result<usize, Error> {
    if count > i32::MAX as usize {
        return Err(Error::Overflow);
    }
    Ok(count)
}

fn add_printf_count(count: usize, add: usize) -> Result<usize, Error> {
    check_printf_count(count.checked_add(add).ok_or(Error::Overflow)?)
}

pub(crate) trait FormatSink {
    fn write_bytes(&mut self, bytes: &[u8]) -> Result<(), Error>;

    fn write_repeat(&mut self, byte: u8, count: usize) -> Result<(), Error> {
        assert!(matches!(byte, b'0' | b' '));
        const ZEROS: &[u8] = b"0000000000000000";
        const SPACES: &[u8] = b"                ";
        let bytes = if byte == b'0' { ZEROS } else { SPACES };
        let mut remaining = count;
        while remaining > 0 {
            let size = remaining.min(bytes.len());
            self.write_bytes(&bytes[..size])?;
            remaining -= size;
        }
        Ok(())
    }
}

struct IoSink<'a, W: IoWrite + ?Sized> {
    output: &'a mut W,
}

impl<W: IoWrite + ?Sized> FormatSink for IoSink<'_, W> {
    fn write_bytes(&mut self, bytes: &[u8]) -> Result<(), Error> {
        self.output.write_all(bytes)?;
        Ok(())
    }
}

struct SliceSink<'a> {
    buffer: &'a mut [u8],
    len: usize,
}

impl FormatSink for SliceSink<'_> {
    fn write_bytes(&mut self, bytes: &[u8]) -> Result<(), Error> {
        let remaining = self.buffer.len().saturating_sub(self.len);
        let stored = remaining.min(bytes.len());
        if stored != 0 {
            self.buffer[self.len..self.len + stored].copy_from_slice(&bytes[..stored]);
            self.len += stored;
        }
        Ok(())
    }

    fn write_repeat(&mut self, byte: u8, count: usize) -> Result<(), Error> {
        assert!(matches!(byte, b'0' | b' '));
        let remaining = self.buffer.len().saturating_sub(self.len);
        let stored = remaining.min(count);
        if stored != 0 {
            self.buffer[self.len..self.len + stored].fill(byte);
            self.len += stored;
        }
        Ok(())
    }
}

pub fn printf_locale_to_slice(
    buffer: &mut [u8],
    fmt: &BStr,
    locale: &Locale,
    args: &mut [Arg],
) -> Result<usize, Error> {
    let mut sink = SliceSink { buffer, len: 0 };
    format_locale(&mut sink, fmt, locale, args)
}

// Pad output by emitting `c` until `min_width` is reached.
pub(super) fn pad(
    f: &mut (impl FormatSink + ?Sized),
    c: u8,
    min_width: usize,
    current_width: usize,
) -> Result<(), Error> {
    assert!(matches!(c, b'0' | b' '));
    if current_width >= min_width {
        return Ok(());
    }
    f.write_repeat(c, min_width - current_width)
}

fn format_unsigned_digits(
    storage: &mut [u8; 64],
    mut value: u64,
    radix: u64,
    uppercase: bool,
) -> &[u8] {
    debug_assert!(matches!(radix, 8 | 10 | 16));
    debug_assert_ne!(value, 0);

    let digits = if uppercase {
        b"0123456789ABCDEF"
    } else {
        b"0123456789abcdef"
    };
    let mut index = storage.len();

    while value != 0 {
        index -= 1;
        storage[index] = digits[(value % radix) as usize];
        value /= radix;
    }

    &storage[index..]
}

/// Formats a byte string using the provided format specifiers, arguments, and locale.
///
/// # Parameters
/// - `f`: The receiver of formatted output.
/// - `fmt`: The format string being parsed.
/// - `locale`: The locale to use for number formatting.
/// - `args`: Iterator over the arguments to format.
///
/// # Returns
/// A `Result` which is `Ok` containing the number of bytes written on success, or an `Error`.
///
/// # Example
///
/// ```
/// use luau_printf::{locale, sprintf_locale, ToArg};
///
/// let mut output = Vec::new();
/// let fmt = luau_printf::BStr::new("%'0.2f");
/// let mut args = [1234567.89_f64.to_arg()];
///
/// let result = sprintf_locale(&mut output, fmt, &locale::EN_US_LOCALE, &mut args);
///
/// assert_eq!(result, Ok(12));
/// assert_eq!(output.as_slice(), b"1,234,567.89");
/// ```
pub fn sprintf_locale<W: IoWrite + ?Sized>(
    f: &mut W,
    fmt: &BStr,
    locale: &Locale,
    args: &mut [Arg],
) -> Result<usize, Error> {
    let mut sink = IoSink { output: f };
    format_locale(&mut sink, fmt, locale, args)
}

fn format_locale(
    f: &mut (impl FormatSink + ?Sized),
    fmt: &BStr,
    locale: &Locale,
    args: &mut [Arg],
) -> Result<usize, Error> {
    use ConversionSpec as CS;
    let mut s = c_string_prefix(fmt);
    let mut args = args.iter_mut();
    let mut out_len: usize = 0;
    let mut float_buf = None;
    'main: while !s.is_empty() {
        // Handle literal text and %% format specifiers.
        let lit = s.take_literal();
        if !lit.is_empty() {
            f.write_bytes(lit.as_ref())?;
            out_len = add_printf_count(out_len, lit.len())?;
            continue 'main;
        }

        // Consume the % at the start of the format specifier.
        debug_assert_eq!(s.at(0), Some(b'%'));
        s.advance_by(1);

        // Read modifier flags. '-' and '0' flags are mutually exclusive.
        let mut flags = ModifierFlags::default();
        while flags.try_set(s.at(0).unwrap_or(b'\0')) {
            s.advance_by(1);
        }
        if flags.left_adj {
            flags.zero_pad = false;
        }

        // Read field width. We do not support $.
        let desired_width = if s.at(0) == Some(b'*') {
            let arg_width = args.next().ok_or(Error::MissingArg)?.as_sint()?;
            s.advance_by(1);
            if arg_width < 0 {
                flags.left_adj = true;
            }
            arg_width
                .unsigned_abs()
                .try_into()
                .map_err(|_| Error::Overflow)?
        } else {
            get_int(&mut s)?
        };
        check_printf_count(desired_width)?;

        // Optionally read precision. We do not support $.
        let mut desired_precision: Option<usize> = if s.at(0) == Some(b'.') && s.at(1) == Some(b'*')
        {
            // "A negative precision is treated as though it were missing."
            // Here we assume the precision is always signed.
            s.advance_by(2);
            let p = args.next().ok_or(Error::MissingArg)?.as_sint()?;
            p.try_into().ok()
        } else if s.at(0) == Some(b'.') {
            s.advance_by(1);
            Some(get_int(&mut s)?)
        } else {
            None
        };
        if let Some(precision) = desired_precision {
            check_printf_count(precision)?;
        }

        // Read out the format specifier and arg.
        let conv_spec = get_specifier(&mut s)?;
        let arg = args.next().ok_or(Error::MissingArg)?;
        let mut prefix = b"".as_slice();

        // Thousands grouping only works for d,u,i,f,F.
        // 'i' is mapped to 'd'.
        if flags.grouped && !matches!(conv_spec, CS::d | CS::u | CS::f | CS::F) {
            return Err(Error::BadFormatString);
        }

        // Disable zero-pad if we have an explicit precision.
        // "If a precision is given with a numeric conversion (d, i, o, u, i, x, and X),
        // the 0 flag is ignored." p is included here.
        let spec_is_numeric = matches!(conv_spec, CS::d | CS::u | CS::o | CS::p | CS::x | CS::X);
        if spec_is_numeric && desired_precision.is_some() {
            flags.zero_pad = false;
        }

        // Apply the formatting. Some cases continue the main loop.
        // Note that numeric conversions must leave 'body' empty if the value is 0.
        let mut body_storage = [0u8; 64];
        let body = match conv_spec {
            CS::n => {
                arg.set_count(out_len)?;
                continue 'main;
            }
            CS::e | CS::f | CS::g | CS::a | CS::E | CS::F | CS::G | CS::A => {
                // Floating point types handle output on their own.
                let float = arg.as_float()?;
                let buf = float_buf.get_or_insert_with(|| Vec::with_capacity(64));
                buf.clear();
                let len = format_float(
                    f,
                    float,
                    desired_width,
                    desired_precision,
                    flags,
                    locale,
                    conv_spec,
                    buf,
                )?;
                out_len = add_printf_count(out_len, len)?;
                continue 'main;
            }
            CS::p => {
                const PTR_HEX_DIGITS: usize = 2 * mem::size_of::<*const u8>();
                desired_precision = desired_precision.map(|p| p.max(PTR_HEX_DIGITS));
                let uint = arg.as_uint()?;
                if uint == 0 {
                    &[][..]
                } else {
                    prefix = b"0x";
                    format_unsigned_digits(&mut body_storage, uint, 16, false)
                }
            }
            CS::x | CS::X => {
                // If someone passes us a negative value, format it with the width
                // we were given.
                let lower = conv_spec.is_lower();
                let uint = arg.as_wrapping_sint()?;
                if uint == 0 {
                    &[][..]
                } else {
                    if flags.alt_form {
                        prefix = if lower { b"0x" } else { b"0X" };
                    }
                    format_unsigned_digits(&mut body_storage, uint, 16, !lower)
                }
            }
            CS::o => {
                let uint = arg.as_uint()?;
                let body = if uint == 0 {
                    &[][..]
                } else {
                    format_unsigned_digits(&mut body_storage, uint, 8, false)
                };
                if flags.alt_form && desired_precision.unwrap_or(0) <= body.len() + 1 {
                    desired_precision = Some(body.len() + 1);
                }
                body
            }
            CS::u => {
                let uint = arg.as_uint()?;
                if uint == 0 {
                    &[][..]
                } else {
                    format_unsigned_digits(&mut body_storage, uint, 10, false)
                }
            }
            CS::d => {
                let arg_i = arg.as_sint()?;
                if arg_i < 0 {
                    prefix = b"-";
                } else if flags.mark_pos {
                    prefix = b"+";
                } else if flags.pad_pos {
                    prefix = b" ";
                }
                if arg_i == 0 {
                    &[][..]
                } else {
                    format_unsigned_digits(&mut body_storage, arg_i.unsigned_abs(), 10, false)
                }
            }
            CS::c => {
                flags.zero_pad = false;
                body_storage[0] = arg.as_uchar()?;
                &body_storage[..1]
            }
            CS::s => {
                let s = arg.as_bstr()?;
                flags.zero_pad = false;
                let scan_limit =
                    desired_precision.map_or(s.len(), |precision| precision.min(s.len()));
                let len = s[..scan_limit]
                    .iter()
                    .position(|&c| c == b'\0')
                    .unwrap_or(scan_limit);
                desired_precision = Some(len);
                &s[..len]
            }
        };
        // Numeric output should be empty iff the value is 0.
        if spec_is_numeric && body.is_empty() {
            debug_assert_eq!(arg.as_uint().unwrap(), 0);
        }

        // Decide if we want to apply thousands grouping to the body, and compute its size.
        // Note we have already errored out if grouped is set and this is non-numeric.
        let wants_grouping = flags.grouped && locale.thousands_sep.is_some();
        let body_width = match wants_grouping {
            // We assume that text representing numbers is ASCII, so len == width.
            true => body.len() + locale.separator_count(body.len()),
            false => body.len(),
        };

        // Resolve the precision.
        // In the case of a non-numeric conversion, update the precision to at least the
        // length of the string.
        let desired_precision = if !spec_is_numeric {
            desired_precision.unwrap_or(body_width)
        } else {
            desired_precision.unwrap_or(1).max(body_width)
        };

        let prefix_width = prefix.len();
        let unpadded_width = prefix_width
            .checked_add(desired_precision)
            .ok_or(Error::Overflow)?;
        let width = desired_width.max(unpadded_width);

        // Pad on the left with spaces to the desired width?
        if !flags.left_adj && !flags.zero_pad {
            pad(f, b' ', width, unpadded_width)?;
        }

        // Output any prefix.
        f.write_bytes(prefix)?;

        // Pad after the prefix with zeros to the desired width?
        if !flags.left_adj && flags.zero_pad {
            pad(f, b'0', width, unpadded_width)?;
        }

        // Pad on the left to the given precision?
        // TODO: why pad with 0 here?
        pad(f, b'0', desired_precision, body_width)?;

        // Output the actual value, perhaps with grouping.
        if wants_grouping {
            f.write_bytes(&locale.apply_grouping(body))?;
        } else {
            f.write_bytes(body)?;
        }

        // Pad on the right with spaces if we are left adjusted?
        if flags.left_adj {
            pad(f, b' ', width, unpadded_width)?;
        }

        out_len = add_printf_count(out_len, width)?;
    }

    Ok(out_len)
}