use std::mem::MaybeUninit;
pub struct MemBufStr<'a> {
len: &'a mut usize,
buffer: &'a mut [MaybeUninit<u8>],
}
impl<'a> MemBufStr<'a> {
pub unsafe fn wrap_uninit(
len: &'a mut usize,
buffer: &'a mut [MaybeUninit<u8>],
) -> MemBufStr<'a> {
MemBufStr { buffer, len }
}
pub unsafe fn wrap(len: &'a mut usize, buffer: &'a mut [u8]) -> MemBufStr<'a> {
MemBufStr {
#[allow(clippy::missing_transmute_annotations)]
buffer: std::mem::transmute(buffer),
len,
}
}
#[allow(clippy::missing_transmute_annotations)]
pub fn str(&self) -> &str {
unsafe { std::str::from_utf8_unchecked(std::mem::transmute(&self.buffer[..*self.len])) }
}
#[allow(clippy::missing_transmute_annotations)]
pub unsafe fn write(&mut self, buf: &[u8]) -> usize {
let len = utf8_max(buf, self.buffer.len() - *self.len);
unsafe {
std::ptr::copy_nonoverlapping(
buf.as_ptr(),
std::mem::transmute(self.buffer.as_mut_ptr().add(*self.len)),
len,
);
}
*self.len += len;
len
}
}
impl std::fmt::Write for MemBufStr<'_> {
fn write_str(&mut self, value: &str) -> std::fmt::Result {
unsafe { self.write(value.as_bytes()) };
Ok(())
}
}
#[derive(Clone, Debug)]
pub struct FixedBufStr<const N: usize> {
len: usize,
buffer: [MaybeUninit<u8>; N],
}
impl<const N: usize> Default for FixedBufStr<N> {
fn default() -> Self {
Self::new()
}
}
#[inline(always)]
fn utf8_max(buf: &[u8], max: usize) -> usize {
let buf_len = buf.len();
if buf_len <= max {
buf_len
} else if max == 0 {
0
} else if unsafe { buf.get_unchecked(max.unchecked_sub(1)) } & 0x80 == 0x00 {
max
} else {
let start = unsafe { max.unchecked_sub(1) };
let mut i = start;
unsafe {
while buf.get_unchecked(i) & 0xC0 == 0x80 {
i = i.unchecked_sub(1);
}
let n = start.unchecked_sub(i);
if (buf.get_unchecked(i) & 0xF0 == 0xF0 && n == 4)
|| (buf.get_unchecked(i) & 0xE0 == 0xE0 && n == 3)
|| (buf.get_unchecked(i) & 0xC0 == 0xC0 && n == 2)
{
max
} else {
i
}
}
}
}
impl<const N: usize> FixedBufStr<N> {
pub fn new() -> FixedBufStr<N> {
FixedBufStr {
buffer: unsafe { MaybeUninit::uninit().assume_init() },
len: 0,
}
}
#[allow(clippy::missing_transmute_annotations)]
pub fn str(&self) -> &str {
unsafe { std::str::from_utf8_unchecked(std::mem::transmute(&self.buffer[..self.len as _])) }
}
#[allow(clippy::missing_transmute_annotations)]
#[allow(clippy::should_implement_trait)]
pub fn from_str(value: &str) -> Self {
let mut buffer = FixedBufStr::new();
let len = utf8_max(value.as_bytes(), N);
unsafe {
std::ptr::copy_nonoverlapping(
value.as_ptr(),
std::mem::transmute(buffer.buffer.as_mut_ptr()),
len,
);
}
buffer.len = len as _;
buffer
}
#[allow(clippy::missing_transmute_annotations)]
pub unsafe fn write(&mut self, buf: &[u8]) -> usize {
let len = utf8_max(buf, N - self.len);
unsafe {
std::ptr::copy_nonoverlapping(
buf.as_ptr(),
std::mem::transmute(self.buffer.as_mut_ptr().add(self.len)),
len,
);
}
self.len += len;
len
}
}
impl<const N: usize> std::fmt::Write for FixedBufStr<N> {
fn write_str(&mut self, value: &str) -> std::fmt::Result {
unsafe { self.write(value.as_bytes()) };
Ok(())
}
}
pub struct IoToFmt<W: std::fmt::Write>(W);
impl<W: std::fmt::Write> IoToFmt<W> {
pub fn new(w: W) -> Self {
Self(w)
}
pub fn into_inner(self) -> W {
self.0
}
}
impl<W: std::fmt::Write> std::io::Write for IoToFmt<W> {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
let str = std::str::from_utf8(buf)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
self.0
.write_str(str)
.map(|_| str.len())
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
#[cfg(test)]
mod tests {
use crate::format::{FixedBufStr, MemBufStr};
use std::fmt::Write;
use std::mem::MaybeUninit;
#[test]
fn basic() {
let mut msg: FixedBufStr<64> = FixedBufStr::new();
let _ = write!(msg, "this");
let _ = write!(msg, " is");
let _ = write!(msg, " a");
let _ = write!(msg, " test");
assert_eq!(msg.str(), "this is a test");
}
#[test]
fn basic_mem() {
let mut buf: [MaybeUninit<u8>; 64] = unsafe { MaybeUninit::uninit().assume_init() };
let mut len = 0;
let mut msg = unsafe { MemBufStr::wrap_uninit(&mut len, &mut buf) };
let _ = write!(msg, "this");
let _ = write!(msg, " is");
let _ = write!(msg, " a");
let _ = write!(msg, " test");
assert_eq!(msg.str(), "this is a test");
}
#[test]
fn truncate_ascii() {
let mut msg: FixedBufStr<4> = FixedBufStr::new();
let _ = write!(msg, "this");
let _ = write!(msg, " is");
let _ = write!(msg, " a");
let _ = write!(msg, " test");
assert_eq!(msg.str().len(), 4);
assert_eq!(msg.str(), "this");
}
#[test]
fn truncate_ascii_mem() {
let mut buf = [0; 4];
let mut len = 0;
let mut msg = unsafe { MemBufStr::wrap(&mut len, &mut buf) };
let _ = write!(msg, "this");
let _ = write!(msg, " is");
let _ = write!(msg, " a");
let _ = write!(msg, " test");
assert_eq!(msg.str().len(), 4);
assert_eq!(msg.str(), "this");
}
#[test]
fn truncate_utf8_exact() {
let mut msg: FixedBufStr<3> = FixedBufStr::new();
let _ = write!(msg, "我");
assert_eq!(msg.str().len(), 3);
assert_eq!(msg.str(), "我");
}
#[test]
fn truncate_utf8_exact_mem() {
let mut buf = [0; 3];
let mut len = 0;
let mut msg = unsafe { MemBufStr::wrap(&mut len, &mut buf) };
let _ = write!(msg, "我");
assert_eq!(msg.str().len(), 3);
assert_eq!(msg.str(), "我");
}
#[test]
fn truncate_utf8_exact2() {
let mut msg: FixedBufStr<6> = FixedBufStr::new();
let _ = write!(msg, "我是");
assert_eq!(msg.str().len(), 6);
assert_eq!(msg.str(), "我是");
}
#[test]
fn truncate_utf8_exact2_mem() {
let mut buf = [0; 6];
let mut len = 0;
let mut msg = unsafe { MemBufStr::wrap(&mut len, &mut buf) };
let _ = write!(msg, "我是");
assert_eq!(msg.str().len(), 6);
assert_eq!(msg.str(), "我是");
}
#[test]
fn truncate_utf8_exact3() {
let mut msg: FixedBufStr<6> = FixedBufStr::new();
let _ = write!(msg, "我abcd");
assert_eq!(msg.str().len(), 6);
assert_eq!(msg.str(), "我abc");
}
#[test]
fn truncate_utf8_exact3_mem() {
let mut buf = [0; 6];
let mut len = 0;
let mut msg = unsafe { MemBufStr::wrap(&mut len, &mut buf) };
let _ = write!(msg, "我abcd");
assert_eq!(msg.str().len(), 6);
assert_eq!(msg.str(), "我abc");
}
#[test]
fn truncate_utf8() {
let mut msg: FixedBufStr<4> = FixedBufStr::new();
let _ = write!(msg, "我是");
assert_eq!(msg.str().len(), 3);
assert_eq!(msg.str(), "我");
}
#[test]
fn truncate_utf8_mem() {
let mut buf = [0; 4];
let mut len = 0;
let mut msg = unsafe { MemBufStr::wrap(&mut len, &mut buf) };
let _ = write!(msg, "我是");
assert_eq!(msg.str().len(), 3);
assert_eq!(msg.str(), "我");
}
}