From e4782aec28be41761f0239a7c876e71c3ed04790 Mon Sep 17 00:00:00 2001 From: John Millikin Date: Thu, 16 Nov 2023 08:49:26 +0900 Subject: [PATCH] Add `core::fmt::WriteCursor` for formatting into a borrowed buffer. --- library/alloc/src/fmt.rs | 2 + library/core/src/fmt/mod.rs | 128 ++++++++++++++++++++++++++++++++++ library/core/tests/fmt/mod.rs | 34 +++++++++ library/core/tests/lib.rs | 1 + 4 files changed, 165 insertions(+) diff --git a/library/alloc/src/fmt.rs b/library/alloc/src/fmt.rs index 5b50ef7bf6c21..6e93111830743 100644 --- a/library/alloc/src/fmt.rs +++ b/library/alloc/src/fmt.rs @@ -557,6 +557,8 @@ pub use core::fmt::Alignment; pub use core::fmt::Error; #[unstable(feature = "debug_closure_helpers", issue = "117729")] pub use core::fmt::FormatterFn; +#[unstable(feature = "fmt_write_cursor", issue = "none")] +pub use core::fmt::WriteCursor; #[stable(feature = "rust1", since = "1.0.0")] pub use core::fmt::{write, Arguments}; #[stable(feature = "rust1", since = "1.0.0")] diff --git a/library/core/src/fmt/mod.rs b/library/core/src/fmt/mod.rs index e1b7b46a1ed2f..71a005d49da4e 100644 --- a/library/core/src/fmt/mod.rs +++ b/library/core/src/fmt/mod.rs @@ -231,6 +231,134 @@ impl Write for &mut W { } } +/// A `WriteCursor` implements [`fmt::Write`] by writing to an in-memory buffer. +/// +/// [`fmt::Write`]: self::Write +/// +/// # Examples +/// +/// ``` +/// #![feature(fmt_write_cursor)] +/// use std::fmt::WriteCursor; +/// +/// # fn test_write_cursor() -> std::fmt::Result { +/// let mut buf = [0u8; 5]; +/// let mut cursor = WriteCursor::new(&mut buf); +/// write!(cursor, "{}", 12345)?; +/// assert_eq!(cursor.as_str(), "12345"); +/// assert_eq!(&buf, b"12345"); +/// # Ok(()) +/// # } +/// ``` +#[doc(alias = "sprintf")] +#[doc(alias = "snprintf")] +#[unstable(feature = "fmt_write_cursor", issue = "none")] +pub struct WriteCursor<'a> { + buf: &'a mut [mem::MaybeUninit], + utf8_len: usize, +} + +#[unstable(feature = "fmt_write_cursor", issue = "none")] +impl Debug for WriteCursor<'_> { + fn fmt(&self, f: &mut Formatter<'_>) -> Result { + f.debug_struct("WriteCursor") + .field("capacity", &self.capacity()) + .field("written", &self.written()) + .finish() + } +} + +impl<'a> WriteCursor<'a> { + /// Creates a new cursor wrapping the provided buffer. + #[unstable(feature = "fmt_write_cursor", issue = "none")] + pub fn new(buf: &'a mut [u8]) -> WriteCursor<'a> { + // SAFETY: [T] and [MaybeUninit] have the same layout, and this + // code never writes uninitialized data to the buffer. + Self::new_uninit(unsafe { mem::transmute(buf) }) + } + + /// Creates a new cursor wrapping the provided uninitialized buffer. + #[unstable(feature = "fmt_write_cursor", issue = "none")] + pub fn new_uninit(buf: &'a mut [mem::MaybeUninit]) -> WriteCursor<'a> { + WriteCursor { buf, utf8_len: 0 } + } + + /// Returns the total capacity of the cursor's underlying buffer, in bytes. + #[unstable(feature = "fmt_write_cursor", issue = "none")] + pub fn capacity(&self) -> usize { + self.buf.len() + } + + /// Returns how many bytes have been written to the cursor's underlying buffer. + #[unstable(feature = "fmt_write_cursor", issue = "none")] + pub fn written(&self) -> usize { + self.utf8_len + } + + /// Returns the written portion of the cursor's underlying buffer as a + /// byte slice. + /// + /// The returned slice contains valid UTF-8. + #[unstable(feature = "fmt_write_cursor", issue = "none")] + pub fn as_bytes(&self) -> &[u8] { + let utf8 = &self.buf[..self.utf8_len]; + // SAFETY: The buffer is incrementally initialized by `write_str`, which + // updates `utf8_len`. + unsafe { mem::MaybeUninit::slice_assume_init_ref(utf8) } + } + + /// Consumes the cursor and returns the written portion of the cursor's + /// underlying buffer as a byte slice. + /// + /// The returned slice contains valid UTF-8. + #[unstable(feature = "fmt_write_cursor", issue = "none")] + pub fn into_bytes(self) -> &'a mut [u8] { + let utf8 = &mut self.buf[..self.utf8_len]; + // SAFETY: The buffer is incrementally initialized by `write_str`, which + // updates `utf8_len`. + unsafe { mem::MaybeUninit::slice_assume_init_mut(utf8) } + } + + /// Returns the written portion of the cursor's underlying buffer as a `&str`. + #[unstable(feature = "fmt_write_cursor", issue = "none")] + pub fn as_str(&self) -> &str { + // SAFETY: The buffer is only initialized by copying from `str` values. + unsafe { str::from_utf8_unchecked(self.as_bytes()) } + } + + /// Consumes the cursor and returns the written portion of the cursor's + /// underlying buffer as a `&str`. + #[unstable(feature = "fmt_write_cursor", issue = "none")] + pub fn into_str(self) -> &'a mut str { + // SAFETY: The buffer is only initialized by copying from `str` values. + unsafe { str::from_utf8_unchecked_mut(self.into_bytes()) } + } + + /// Allows the [`write!`] macro to write to a `WriteCursor`. + /// + /// This method should generally not be invoked manually. It exists so that + /// the `write!` macro can write to a `WriteCursor` without needing a + /// `use std::fmt::Write` declaration. + #[unstable(feature = "fmt_write_cursor", issue = "none")] + pub fn write_fmt(&mut self, args: Arguments<'_>) -> Result { + Write::write_fmt(self, args) + } +} + +#[unstable(feature = "fmt_write_cursor", issue = "none")] +impl Write for WriteCursor<'_> { + fn write_str(&mut self, s: &str) -> Result { + let b = s.as_bytes(); + let avail = &mut self.buf[self.utf8_len..]; + if b.len() > avail.len() { + return Err(Error); + } + mem::MaybeUninit::write_slice(&mut avail[..b.len()], b); + self.utf8_len += b.len(); + Ok(()) + } +} + /// Configuration for formatting. /// /// A `Formatter` represents various options related to formatting. Users do not diff --git a/library/core/tests/fmt/mod.rs b/library/core/tests/fmt/mod.rs index c1c80c46c78b7..a8e4e187155f6 100644 --- a/library/core/tests/fmt/mod.rs +++ b/library/core/tests/fmt/mod.rs @@ -43,3 +43,37 @@ fn pad_integral_resets() { assert_eq!(format!("{Bar:<03}"), "1 0051 "); } + +#[test] +fn write_cursor() { + use core::fmt::WriteCursor; + + let mut buf = [0u8; 5]; + + { + let mut c = WriteCursor::new(&mut buf); + assert!(write!(c, "{}", 12345).is_ok()); + assert_eq!(c.as_bytes(), b"12345"); + assert_eq!(c.as_str(), "12345"); + } + + // Writes either fully succeed or leave the cursor unchanged. A cursor that + // failed to write remains valid for smaller writes. + { + let mut c = WriteCursor::new(&mut buf); + assert!(write!(c, "{}", 123456).is_err()); + assert_eq!(c.written(), 0); + assert_eq!(c.as_str(), ""); + + assert!(write!(c, "{}", 1234).is_ok()); + assert_eq!(c.written(), 4); + assert_eq!(c.as_str(), "1234"); + + assert!(write!(c, "{}", 1).is_ok()); + assert_eq!(c.written(), 5); + assert_eq!(c.as_str(), "12341"); + + assert!(write!(c, "{}", 1).is_err()); + assert_eq!(c.written(), 5); + } +} diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs index c167240ba2439..6eff8fe116859 100644 --- a/library/core/tests/lib.rs +++ b/library/core/tests/lib.rs @@ -34,6 +34,7 @@ #![feature(extern_types)] #![feature(flt2dec)] #![feature(fmt_internals)] +#![feature(fmt_write_cursor)] #![feature(float_minimum_maximum)] #![feature(future_join)] #![feature(generic_assert_internals)]