Made `bytes` an optional feature.

Fortune for collect's current commit: Future small blessing − 末小吉

Cargo.toml

@@ -10,7 +10,7 @@ default = ["jemalloc"]
 
 # TODO: mmap, memfd_create() ver
 
-# TODO: make `bytes` optional, and use Vec<u8> by default instead? idk...
+# bytes: use `bytes` crate for collecting instead of `std::vec`
 
 # Use jemalloc instead of system malloc.
 # Seems to reduce overall memory usage at the cost of a very small speed drop.
@@ -28,6 +28,6 @@ inherits="release"
 strip=false
 
 [dependencies]
-bytes = "1.1.0"
+bytes = {version = "1.1.0", optional = true }
 jemallocator = { version = "0.3.2", optional = true }
 libc = "0.2.122"

src/buffers.rs

@@ -0,0 +1,336 @@
+//! Buffers and helpers
+use super::*;
+use std::num::NonZeroUsize;
+
+#[cfg(feature="bytes")]
+/// Default mutable buffer
+#[allow(dead_code)]
+pub type DefaultMut = bytes::BytesMut;
+
+#[cfg(not(feature="bytes"))] 
+/// Default mutable buffer
+#[allow(dead_code)]
+pub type DefaultMut = Vec<u8>;
+
+/// Default immutable buffer
+#[allow(dead_code)]
+pub type Default = <DefaultMut as MutBuffer>::Frozen;
+
+
+/// Reader from a mutable reference of a `Buffer`.
+#[derive(Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
+pub struct BufferReader<'a, B: ?Sized>(&'a mut B, usize);
+
+
+/// Writer to a mutable reference of a `MutBuffer`.
+#[derive(Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
+pub struct BufferWriter<'a, B: ?Sized>(&'a mut B, usize);
+
+#[allow(dead_code)]
+const _: () = {
+    impl<'a, B: ?Sized + Buffer> BufferReader<'a, B>
+    {
+	#[inline(always)] 
+	pub fn get(&self) -> &B
+	{
+	    &self.0
+	}
+	#[inline(always)] 
+	pub fn get_mut(&mut self) -> &B
+	{
+	    &mut self.0
+	}
+	#[inline(always)] 
+	pub fn amount_read(&self) -> usize
+	{
+	    self.1
+	}
+    }
+    impl<'a, 'b: 'a, B: Buffer + 'b> BufferReader<'a, B>
+    {
+	#[inline] 
+	pub fn unsize(self) -> BufferReader<'a, (dyn Buffer + 'b)>
+	{
+	    BufferReader(self.0, self.1)
+	}
+    }
+
+    impl<'a, B: ?Sized + Buffer> BufferWriter<'a, B>
+    {
+	#[inline(always)] 
+	pub fn get(&self) -> &B
+	{
+	    &self.0
+	}
+	#[inline(always)] 
+	pub fn get_mut(&mut self) -> &B
+	{
+	    &mut self.0
+	}
+	#[inline(always)] 
+	pub fn amount_written(&self) -> usize
+	{
+	    self.1
+	}
+    }
+    impl<'a, 'b: 'a, B: Buffer + 'b> BufferWriter<'a, B>
+    {
+	#[inline] 
+	pub fn unsize(self) -> BufferWriter<'a, (dyn Buffer + 'b)>
+	{
+	    BufferWriter(self.0, self.1)
+	}
+    }
+};
+
+impl<'a, B: ?Sized + Buffer> io::Read for BufferReader<'a, B>
+{
+    #[inline] 
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+	let adv = self.0.copy_to_slice(self.1, buf);
+	self.1 += adv;
+	Ok(adv)
+    }
+}
+
+impl<'a, B: ?Sized + MutBuffer> io::Write for BufferWriter<'a, B>
+{
+    #[inline]
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+	let adv = self.0.copy_from_slice(self.1, buf);
+	self.1 += adv;
+	Ok(adv)
+    }
+    #[inline(always)] 
+    fn flush(&mut self) -> io::Result<()> {
+	Ok(())
+    }
+}
+
+/// An immutable contiguous buffer
+pub trait Buffer: AsRef<[u8]>
+{
+    #[inline] 
+    fn copy_to_slice(&self, st: usize, slice: &mut [u8]) -> usize
+    {
+	let by = self.as_ref();
+	if st >= by.len() {
+	    return 0;
+	}
+	
+	let by = &by[st..];
+	let len = std::cmp::min(by.len(), slice.len());
+	// SAFETY: We know `self`'s AsRef impl cannot overlap with `slice`, since `slice` is a mutable reference.
+	if len > 0 {
+	    unsafe {
+		std::ptr::copy_nonoverlapping(by.as_ptr(), slice.as_mut_ptr(), len)
+	    }
+	}
+	len
+    }
+
+}
+pub trait BufferExt: Buffer
+{
+    #[inline(always)] 
+    fn reader_from(&mut self, st: usize) -> BufferReader<'_, Self>
+    {
+	BufferReader(self, st)
+    }
+    #[inline]
+    fn reader(&mut self) -> BufferReader<'_, Self>
+    {
+	self.reader_from(0)
+    }
+}
+impl<B: Buffer> BufferExt for B{}
+
+impl<T: ?Sized> Buffer for T
+where T: AsRef<[u8]>
+{}
+
+/// A mutable contiguous buffer
+pub trait MutBuffer: AsMut<[u8]>
+{
+    type Frozen: Sized + Buffer;
+    
+    /// Make immutable
+    fn freeze(self) -> Self::Frozen;
+
+    #[inline] 
+    fn copy_from_slice(&mut self, st: usize, slice: &[u8]) -> usize
+    {
+	let by = self.as_mut();
+	if st >= by.len() {
+	    return 0;
+	}
+
+	let by = &mut by[st..];
+	let len = std::cmp::min(by.len(), slice.len());
+	
+	if len > 0 {
+	    // SAFETY: We know `self`'s AsRef impl cannot overlap with `slice`, since `slice` is a mutable reference.
+	    unsafe {
+		std::ptr::copy_nonoverlapping(slice.as_ptr(), by.as_mut_ptr(), len);
+	    }
+	}
+	len
+    }
+}
+
+pub trait MutBufferExt: MutBuffer
+{
+    #[inline(always)] 
+    fn writer_from(&mut self, st: usize) -> BufferWriter<'_, Self>
+    {
+	BufferWriter(self, st)
+    }
+    #[inline] 
+    fn writer(&mut self) -> BufferWriter<'_, Self>
+    {
+	self.writer_from(0)
+    }
+}
+impl<B: ?Sized + MutBuffer> MutBufferExt for B{}
+
+#[cfg(feature="bytes")]
+impl MutBuffer for bytes::BytesMut
+{
+    type Frozen = bytes::Bytes;
+    #[inline(always)] 
+    fn freeze(self) -> Self::Frozen {
+	bytes::BytesMut::freeze(self)
+    }
+}
+
+impl MutBuffer for Vec<u8>
+{
+    type Frozen = Box<[u8]>;
+    #[inline]
+    fn freeze(self) -> Self::Frozen {
+	self.into_boxed_slice()
+    }
+}
+
+/// A trait for buffers that can be allocated with a capacity
+pub trait WithCapacity: Sized
+{
+    fn wc_new() -> Self;
+    fn wc_with_capacity(_: usize) -> Self;
+}
+
+impl WithCapacity for Box<[u8]>
+{
+    #[inline(always)] 
+    fn wc_new() -> Self {
+	Vec::wc_new().into_boxed_slice()
+    }
+    #[inline(always)] 
+    fn wc_with_capacity(cap: usize) -> Self {
+	Vec::wc_with_capacity(cap).into_boxed_slice()
+    }
+}
+
+pub trait WithCapExt: WithCapacity
+{
+    fn maybe_with_capacity(maybe: Option<NonZeroUsize>) -> Self;
+    #[inline(always)]
+    fn try_with_capacity(cap: usize) -> Self
+    {
+	Self::maybe_with_capacity(NonZeroUsize::new(cap))
+    }
+}
+
+/// A type that can be used as a size for creating a `WithCapacity` buffer
+pub trait TryCreateBuffer
+{
+    fn create_buffer<T: WithCapacity>(&self) -> T;
+}
+
+impl TryCreateBuffer for Option<NonZeroUsize>
+{
+    #[inline(always)] 
+    fn create_buffer<T: WithCapacity>(&self) -> T {
+	T::maybe_with_capacity(*self)
+    }
+}
+
+impl TryCreateBuffer for usize
+{
+    #[inline(always)]
+    fn create_buffer<T: WithCapacity>(&self) -> T {
+	T::try_with_capacity(*self)
+    }
+}
+
+impl<T: WithCapacity> WithCapExt for T
+{
+    #[inline] 
+    fn maybe_with_capacity(maybe: Option<NonZeroUsize>) -> Self {
+	match maybe {
+	    Some(sz) => Self::wc_with_capacity(sz.into()),
+	    None => Self::wc_new()
+	}
+    }
+    
+}
+
+/// Implement `WithCapacity` for a type that supports it.
+macro_rules! cap_buffer  {
+    ($name:ty) => {
+	impl $crate::buffers::WithCapacity for $name
+	{
+	    #[inline(always)] 
+	    fn wc_new() -> Self
+	    {
+		Self::new()
+	    }
+	    #[inline(always)] 
+	    fn wc_with_capacity(cap: usize) -> Self
+	    {
+		Self::with_capacity(cap)
+	    }
+	}
+    };
+}
+
+pub mod prelude
+{
+    /// Export these items anonymously.
+    macro_rules! export_anon {
+	($($name:ident),+ $(,)?) => {
+	    $(
+		pub use super::$name as _;
+	    )*
+	};
+    }
+
+    // Causes conflicts for `.writer()`, so remove them from prelude.
+    #[cfg(feature="bytes")]
+    export_anon!{
+	WithCapExt,
+	//BufferExt,
+	//MutBufferExt,
+	WithCapExt,
+    }
+    
+    #[cfg(not(feature="bytes"))] 
+    export_anon!{
+	WithCapExt,
+	BufferExt,
+	MutBufferExt,
+	WithCapExt,
+    }
+    
+    pub use super::{
+	WithCapacity,
+	TryCreateBuffer,
+	MutBuffer,
+	Buffer,
+    };
+}
+
+pub(crate) use cap_buffer;
+
+#[cfg(feature="bytes")] buffers::cap_buffer!(bytes::BytesMut);
+cap_buffer!(Vec<u8>);

src/main.rs

@@ -14,8 +14,11 @@ use std::{
     num::NonZeroUsize,
 };
 
+mod buffers;
+use buffers::prelude::*;
+
+#[cfg(feature="bytes")]
 use bytes::{
-    BytesMut,
     Buf,
     BufMut,
 };
@@ -45,16 +48,14 @@ where R: AsRawFd
 fn main() -> io::Result<()> {
     let (bytes, read) = {
 	let stdin = io::stdin();
-	let mut bytes = match try_get_size(&stdin) {
+	let mut bytes: buffers::DefaultMut = try_get_size(&stdin).create_buffer();
-	    Some(sz) => BytesMut::with_capacity(sz.into()),
-	    None => BytesMut::new(),
-	};
 	
 	let read = io::copy(&mut stdin.lock(), &mut (&mut bytes).writer())?;
 	(bytes.freeze(), read as usize)
     };
 
-    let written = io::copy(&mut bytes.slice(..read).reader() , &mut io::stdout().lock())?;
+    let written = 
+	io::copy(&mut (&bytes[..read]).reader() , &mut io::stdout().lock())?;
 
     if read != written as usize {
 	return Err(io::Error::new(io::ErrorKind::BrokenPipe, format!("read {read} bytes, but only wrote {written}")));