From cdaaae14448919d497a8172bdf9dba90e778eb7f Mon Sep 17 00:00:00 2001
From: Elie ROUDNINSKI <elie@cloudflare.com>
Date: Thu, 14 Sep 2023 12:26:02 +0200
Subject: [PATCH] Implement `std::io::Write` for `Vec<'bump, u8>`

---
 Cargo.toml             |  1 +
 README.md              |  9 ++++++++-
 src/collections/vec.rs | 22 ++++++++++++++++++++++
 src/lib.rs             |  2 +-
 tests/all/vec.rs       | 27 +++++++++++++++++++++++++++
 5 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 351e185..ea73ca2 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -46,6 +46,7 @@ default = []
 collections = []
 boxed = []
 allocator_api = []
+std = []
 
 # [profile.bench]
 # debug = true
diff --git a/README.md b/README.md
index 6a83b8c..85e5c46 100644
--- a/README.md
+++ b/README.md
@@ -155,7 +155,14 @@ in its space itself.
 
 ### `#![no_std]` Support
 
-Bumpalo is a `no_std` crate. It depends only on the `alloc` and `core` crates.
+Bumpalo is a `no_std` crate by default. It depends only on the `alloc` and `core` crates.
+
+### `std` Support
+
+You can optionally decide to enable the `std` feature in order to enable some
+std only trait implementations for some collections:
+
+* `std::io::Write` for `Vec<'bump, u8>`
 
 ### Thread support
 
diff --git a/src/collections/vec.rs b/src/collections/vec.rs
index 312aa05..66c821b 100644
--- a/src/collections/vec.rs
+++ b/src/collections/vec.rs
@@ -104,6 +104,8 @@ use core::ops::{Index, IndexMut, RangeBounds};
 use core::ptr;
 use core::ptr::NonNull;
 use core::slice;
+#[cfg(feature = "std")]
+use std::io;
 
 unsafe fn arith_offset<T>(p: *const T, offset: isize) -> *const T {
     p.offset(offset)
@@ -2612,3 +2614,23 @@ where
         }
     }
 }
+
+#[cfg(feature = "std")]
+impl<'bump> io::Write for Vec<'bump, u8> {
+    #[inline]
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        self.extend_from_slice(buf);
+        Ok(buf.len())
+    }
+
+    #[inline]
+    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
+        self.extend_from_slice(buf);
+        Ok(())
+    }
+
+    #[inline]
+    fn flush(&mut self) -> io::Result<()> {
+        Ok(())
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index 221204f..a173321 100755
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,7 +1,7 @@
 #![doc = include_str!("../README.md")]
 #![deny(missing_debug_implementations)]
 #![deny(missing_docs)]
-#![no_std]
+#![cfg_attr(not(feature = "std"), no_std)]
 #![cfg_attr(feature = "allocator_api", feature(allocator_api))]
 
 #[doc(hidden)]
diff --git a/tests/all/vec.rs b/tests/all/vec.rs
index c069f7e..2c46780 100644
--- a/tests/all/vec.rs
+++ b/tests/all/vec.rs
@@ -105,3 +105,30 @@ fn test_vec_items_get_dropped() {
     }
     assert_eq!("Dropped!Dropped!", buffer.borrow().deref());
 }
+
+#[cfg(feature = "std")]
+#[test]
+fn test_vec_write() {
+    use std::io::Write;
+
+    let b = Bump::new();
+    let mut v = bumpalo::vec![in &b];
+
+    assert_eq!(v.write(&[]).unwrap(), 0);
+
+    v.flush().unwrap();
+
+    assert_eq!(v.write(&[1]).unwrap(), 1);
+
+    v.flush().unwrap();
+
+    v.write_all(&[]).unwrap();
+
+    v.flush().unwrap();
+
+    v.write_all(&[2, 3]).unwrap();
+
+    v.flush().unwrap();
+
+    assert_eq!(v, &[1, 2, 3]);
+}