From dca393f1d7cf4987ba769d3217ceaf8a83f90607 Mon Sep 17 00:00:00 2001
From: Eric Ridge <eebbrr@gmail.com>
Date: Fri, 24 Sep 2021 11:58:46 -0600
Subject: [PATCH] Issue #151: impl std::io::Write for u8-based TinyVecs (#152)

* impl std::io::Write for u8-based TinyVec

* undo bad formatting

* fix formatting
---
 Cargo.toml       |  5 ++++-
 src/lib.rs       |  2 +-
 src/tinyvec.rs   | 14 ++++++++++++++
 tests/tinyvec.rs | 15 +++++++++++++++
 4 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 4ca9680..3863c3d 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -22,6 +22,9 @@ default = []
 # Provide things that utilize the `alloc` crate, namely `TinyVec`.
 alloc = ["tinyvec_macros"]
 
+# Provide things that require Rust's `std` module
+std = []
+
 # (not part of Vec!) Extra methods to let you grab the slice of memory after the
 # "active" portion of an `ArrayVec` or `SliceVec`.
 grab_spare_slice = []
@@ -65,7 +68,7 @@ serde_test = "1.0"
 
 [[test]]
 name = "tinyvec"
-required-features = ["alloc"]
+required-features = ["alloc", "std"]
 
 [[bench]]
 name = "macros"
diff --git a/src/lib.rs b/src/lib.rs
index 29271c7..68a7036 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,4 +1,4 @@
-#![no_std]
+#![cfg_attr(not(feature = "std"), no_std)]
 #![forbid(unsafe_code)]
 #![cfg_attr(
   feature = "nightly_slice_partition_dedup",
diff --git a/src/tinyvec.rs b/src/tinyvec.rs
index efee037..7623fa8 100644
--- a/src/tinyvec.rs
+++ b/src/tinyvec.rs
@@ -172,6 +172,20 @@ impl<A: Array, I: SliceIndex<[A::Item]>> IndexMut<I> for TinyVec<A> {
   }
 }
 
+#[cfg(feature = "std")]
+impl<A: Array<Item = u8>> std::io::Write for TinyVec<A> {
+  #[inline(always)]
+  fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+    self.extend_from_slice(buf);
+    Ok(buf.len())
+  }
+
+  #[inline(always)]
+  fn flush(&mut self) -> std::io::Result<()> {
+    Ok(())
+  }
+}
+
 #[cfg(feature = "serde")]
 #[cfg_attr(docs_rs, doc(cfg(feature = "serde")))]
 impl<A: Array> Serialize for TinyVec<A>
diff --git a/tests/tinyvec.rs b/tests/tinyvec.rs
index 1f26be4..ed1bcf9 100644
--- a/tests/tinyvec.rs
+++ b/tests/tinyvec.rs
@@ -394,3 +394,18 @@ fn TinyVec_pretty_debug() {
 
   assert_eq!(s, expected);
 }
+
+#[cfg(feature = "std")]
+#[test]
+fn TinyVec_std_io_write() {
+  use std::io::Write;
+  let mut tv: TinyVec<[u8; 3]> = TinyVec::new();
+
+  tv.write_all(b"foo").ok();
+  assert!(tv.is_inline());
+  assert_eq!(tv, tiny_vec![b'f', b'o', b'o']);
+
+  tv.write_all(b"bar").ok();
+  assert!(tv.is_heap());
+  assert_eq!(tv, tiny_vec![b'f', b'o', b'o', b'b', b'a', b'r']);
+}