diff --git a/src/lib.rs b/src/lib.rs index dd8cc9661..173c13c89 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -115,3 +115,67 @@ fn abort() -> ! { panic!("abort"); } } + +/// extension traits +#[cfg(feature = "std")] +pub mod ext { + /// extension trait for BytesMut + pub trait BytesMutExt { + /// extension trait for BytesMut + fn write(&mut self, buf: &[u8]) -> std::io::Result; + /// extension trait for BytesMut + fn write_vectored(&mut self, bufs: &[std::io::IoSlice<'_>]) -> std::io::Result; + /// extension trait for BytesMut + fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()>; + /// extension trait for BytesMut + fn flush(&mut self) -> std::io::Result<()>; + } + impl BytesMutExt for super::BytesMut { + #[inline] + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.extend_from_slice(buf); + Ok(buf.len()) + } + + #[inline] + fn write_vectored(&mut self, bufs: &[std::io::IoSlice<'_>]) -> std::io::Result { + let len = bufs.iter().map(|b| b.len()).sum(); + self.reserve(len); + for buf in bufs { + self.extend_from_slice(buf); + } + Ok(len) + } + #[inline] + fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> { + self.extend_from_slice(buf); + Ok(()) + } + + #[inline] + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + impl std::io::Write for dyn BytesMutExt { + #[inline] + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.write(buf) + } + + #[inline] + fn write_vectored(&mut self, bufs: &[std::io::IoSlice<'_>]) -> std::io::Result { + self.write_vectored(bufs) + } + #[inline] + fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> { + self.write_all(buf) + } + + #[inline] + fn flush(&mut self) -> std::io::Result<()> { + self.flush() + } + } +} diff --git a/tests/test_buf.rs b/tests/test_buf.rs index fbad003a4..627484d1f 100644 --- a/tests/test_buf.rs +++ b/tests/test_buf.rs @@ -118,3 +118,37 @@ fn copy_to_bytes_overflow() { let _bytes = buf.copy_to_bytes(12); } + +#[cfg(feature = "std")] +#[test] +fn fmt_write_bytesmut_test() { + use bytes::BytesMut; + use std::fmt::Write; + let mut buf = BytesMut::with_capacity(64); + write!(&mut buf, "hello").ok(); + assert_eq!(&buf[..], b"hello"); +} + +#[cfg(feature = "std")] +#[test] +fn io_write_bytesmut_test() { + use bytes::ext::BytesMutExt; + use bytes::BytesMut; + use std::io::Write; + let mut buf = BytesMut::with_capacity(64); + write!(&mut buf as &mut dyn BytesMutExt, "hello").ok(); + assert_eq!(&buf[..], b"hello"); +} + +#[cfg(feature = "std")] +#[test] +fn fmt_write_bytesmut_with_both_write_in_scope_test() { + use bytes::ext::BytesMutExt; + use bytes::BytesMut; + #[allow(unused_imports)] + use std::fmt::Write as FmtWrite; + use std::io::Write as IoWrite; + let mut buf = BytesMut::with_capacity(64); + write!(&mut buf as &mut dyn BytesMutExt, "hello").ok(); + assert_eq!(&buf[..], b"hello"); +}