use core::fmt;
use crate::error::{Error, Result};
#[cfg(feature = "alloc")]
use alloc::string::String;
pub trait Write {
fn write_bytes(&mut self, bytes: &[u8]) -> Result<()>;
#[inline]
fn write_str(&mut self, s: &str) -> Result<()> {
self.write_bytes(s.as_bytes())
}
#[inline]
fn write_byte(&mut self, b: u8) -> Result<()> {
self.write_bytes(core::slice::from_ref(&b))
}
#[inline]
fn write_char(&mut self, c: char) -> Result<()> {
let mut buf = [0u8; 4];
self.write_bytes(c.encode_utf8(&mut buf).as_bytes())
}
}
impl Write for &mut dyn Write {
#[inline]
fn write_bytes(&mut self, bytes: &[u8]) -> Result<()> {
(**self).write_bytes(bytes)
}
}
#[derive(Debug)]
pub struct Buf<'a> {
data: &'a mut [u8],
len: usize,
}
impl<'a> Buf<'a> {
pub fn new(data: &'a mut [u8]) -> Buf<'a> {
Buf { data, len: 0 }
}
#[inline]
pub const fn len(&self) -> usize {
self.len
}
#[inline]
pub const fn is_empty(&self) -> bool {
self.len == 0
}
#[inline]
pub const fn capacity(&self) -> usize {
self.data.len()
}
#[inline]
pub fn as_bytes(&self) -> &[u8] {
&self.data[..self.len]
}
#[inline]
pub fn as_str(&self) -> &str {
core::str::from_utf8(self.as_bytes()).expect("Buf only ever holds valid UTF-8")
}
#[inline]
pub fn clear(&mut self) {
self.len = 0;
}
}
impl Write for Buf<'_> {
#[inline]
fn write_bytes(&mut self, bytes: &[u8]) -> Result<()> {
core::str::from_utf8(bytes).map_err(|_| Error::invalid("invalid utf-8"))?;
let end = self
.len
.checked_add(bytes.len())
.ok_or_else(Error::buffer_overflow)?;
if end > self.data.len() {
return Err(Error::buffer_overflow());
}
self.data[self.len..end].copy_from_slice(bytes);
self.len = end;
Ok(())
}
}
#[cfg(feature = "alloc")]
impl Write for alloc::string::String {
#[inline]
fn write_bytes(&mut self, bytes: &[u8]) -> Result<()> {
let s = core::str::from_utf8(bytes).map_err(|_| Error::invalid("invalid utf-8"))?;
self.push_str(s);
Ok(())
}
}
pub(crate) struct FmtSink<'a, 'b>(pub(crate) &'a mut fmt::Formatter<'b>);
impl Write for FmtSink<'_, '_> {
#[inline]
fn write_bytes(&mut self, bytes: &[u8]) -> Result<()> {
let s = core::str::from_utf8(bytes).map_err(|_| Error::invalid("invalid utf-8"))?;
self.0.write_str(s).map_err(|_| Error::invalid("fmt error"))
}
}
#[inline]
pub(crate) fn with_buf(
out: &mut [u8],
f: impl FnOnce(&mut Buf<'_>) -> Result<()>,
) -> Result<usize> {
let mut buf = Buf::new(out);
f(&mut buf)?;
Ok(buf.len())
}
#[cfg(feature = "alloc")]
pub(crate) fn alloc_string<F: FnOnce(&mut Buf<'_>) -> Result<()>>(f: F) -> Result<String> {
let mut storage = [0u8; 64];
let mut buf = Buf::new(&mut storage);
f(&mut buf)?;
let s = core::str::from_utf8(buf.as_bytes()).expect("formatter output is ASCII");
Ok(alloc::string::String::from(s))
}
#[inline]
pub(crate) fn write_padded(out: &mut dyn Write, value: i64, width: usize) -> Result<()> {
debug_assert!(value >= 0);
let mut digits = [0u8; 20];
let mut len = 0usize;
let mut n = value;
loop {
digits[len] = (n % 10) as u8;
n /= 10;
len += 1;
if n == 0 {
break;
}
}
for _ in len..width {
out.write_byte(b'0')?;
}
for i in (0..len).rev() {
out.write_byte(b'0' + digits[i])?;
}
Ok(())
}
#[inline]
pub(crate) fn write_u128(out: &mut dyn Write, mut value: u128) -> Result<()> {
let mut digits = [0u8; 39];
let mut len = 0usize;
loop {
digits[len] = (value % 10) as u8;
value /= 10;
len += 1;
if value == 0 {
break;
}
}
for i in (0..len).rev() {
out.write_byte(b'0' + digits[i])?;
}
Ok(())
}
#[inline]
pub(crate) fn write_signed_i128(out: &mut dyn Write, value: i128) -> Result<()> {
if value < 0 {
out.write_byte(b'-')?;
write_u128(out, value.unsigned_abs())
} else {
write_u128(out, value as u128)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn buf_writes_and_reports_capacity() {
let mut storage = [0u8; 8];
let mut buf = Buf::new(&mut storage);
assert!(buf.is_empty());
assert_eq!(buf.capacity(), 8);
buf.write_str("abc").unwrap();
buf.write_byte(b'd').unwrap();
buf.write_char('é').unwrap();
assert_eq!(buf.len(), 6);
assert_eq!(buf.as_str(), "abcdé");
assert!(buf.write_str("xyz").is_err());
assert_eq!(buf.as_str(), "abcdé");
assert!(buf.write_bytes(&[0xff]).is_err());
assert_eq!(buf.as_str(), "abcdé");
buf.clear();
assert!(buf.is_empty());
}
#[test]
fn padded_and_u128() {
let mut storage = [0u8; 64];
let mut buf = Buf::new(&mut storage);
write_padded(&mut buf, 42, 4).unwrap();
write_u128(&mut buf, 0).unwrap();
write_u128(&mut buf, u128::MAX).unwrap();
assert_eq!(buf.as_str(), "00420340282366920938463463374607431768211455");
}
}