From ac678bb4356cb081e2fa042c69477a9b9d9a0e29 Mon Sep 17 00:00:00 2001
From: Matt Corallo <git@bluematt.me>
Date: Tue, 12 Sep 2023 17:47:37 +0000
Subject: [PATCH] [IO] Move to custom `Write` trait mirroring `std::io::Write`

In order to move towards our own I/O traits in the `rust-bitcoin`
ecosystem, we have to slowly replace our use of the `std` and
`core2` traits.

Here we take the first real step, replacing
`{std,core2}::io::Write` withour own `bitcoin_io::io::Write`. We
provide a blanket impl for our trait for all `std::io::Write`, if
the `std` feature is enabled, allowing users who use their own
streams or `std` streams to call `rust-bitcoin` methods directly.
---
 io/src/lib.rs | 57 +++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 55 insertions(+), 2 deletions(-)

diff --git a/io/src/lib.rs b/io/src/lib.rs
index e4a4a79e..85b2afe9 100644
--- a/io/src/lib.rs
+++ b/io/src/lib.rs
@@ -36,10 +36,63 @@ pub mod io {
     compile_error!("At least one of std or core2 must be enabled");
 
     #[cfg(feature = "std")]
-    pub use std::io::{Read, Write, sink, Cursor, Take, Error, ErrorKind, Result};
+    pub use std::io::{Read, sink, Cursor, Take, Error, ErrorKind, Result};
 
     #[cfg(not(feature = "std"))]
-    pub use core2::io::{Read, Write, Cursor, Take, Error, ErrorKind, Result};
+    pub use core2::io::{Read, Cursor, Take, Error, ErrorKind, Result};
 
+    /// A generic trait describing an output stream. See [`std::io::Write`] for more info.
+    pub trait Write {
+        fn write(&mut self, buf: &[u8]) -> Result<usize>;
+        fn flush(&mut self) -> Result<()>;
 
+        #[inline]
+        fn write_all(&mut self, mut buf: &[u8]) -> Result<()> {
+            while !buf.is_empty() {
+                match self.write(buf) {
+                    Ok(0) => return Err(Error::new(ErrorKind::UnexpectedEof, "")),
+                    Ok(len) => buf = &buf[len..],
+                    Err(e) if e.kind() == ErrorKind::Interrupted => {}
+                    Err(e) => return Err(e),
+                }
+            }
+            Ok(())
+        }
+    }
+
+    #[cfg(feature = "std")]
+    impl<W: std::io::Write> Write for W {
+        #[inline]
+        fn write(&mut self, buf: &[u8]) -> Result<usize> {
+            <W as std::io::Write>::write(self, buf)
+        }
+        #[inline]
+        fn flush(&mut self) -> Result<()> {
+            <W as std::io::Write>::flush(self)
+        }
+    }
+
+    #[cfg(all(feature = "alloc", not(feature = "std")))]
+    impl Write for alloc::vec::Vec<u8> {
+        #[inline]
+        fn write(&mut self, buf: &[u8]) -> Result<usize> {
+            self.extend_from_slice(buf);
+            Ok(buf.len())
+        }
+        #[inline]
+        fn flush(&mut self) -> Result<()> { Ok(()) }
+    }
+
+    #[cfg(not(feature = "std"))]
+    impl<'a> Write for &'a mut [u8] {
+        #[inline]
+        fn write(&mut self, buf: &[u8]) -> Result<usize> {
+            let cnt = core::cmp::min(self.len(), buf.len());
+            self[..cnt].copy_from_slice(&buf[..cnt]);
+            *self = &mut core::mem::take(self)[cnt..];
+            Ok(cnt)
+        }
+        #[inline]
+        fn flush(&mut self) -> Result<()> { Ok(()) }
+    }
 }