Start to implement Borsh Write by reusing the write function, wrap std

Author: af-afk

src/arrayvec.rs

@@ -1340,17 +1340,26 @@ where
     }
 }
 
+impl<const CAP: usize> ArrayVec<u8, CAP> {
+    /// `Write` appends written data to the end of the vector.
+    pub fn write(&mut self, data: &[u8]) -> usize {
+        let len = cmp::min(self.remaining_capacity(), data.len());
+        let _result = self.try_extend_from_slice(&data[..len]);
+        debug_assert!(_result.is_ok());
+        len
+    }
+}
+
 #[cfg(feature = "std")]
 /// `Write` appends written data to the end of the vector.
+/// Wraps the existing write function.
 ///
 /// Requires `features="std"`.
 impl<const CAP: usize> io::Write for ArrayVec<u8, CAP> {
     fn write(&mut self, data: &[u8]) -> io::Result<usize> {
-        let len = cmp::min(self.remaining_capacity(), data.len());
-        let _result = self.try_extend_from_slice(&data[..len]);
-        debug_assert!(_result.is_ok());
-        Ok(len)
+        Ok(ArrayVec::write(self, data))
     }
+
     fn flush(&mut self) -> io::Result<()> {
         Ok(())
     }
@@ -1419,6 +1428,33 @@ where
     }
 }
 
+#[cfg(all(feature = "borsh", not(feature = "std")))]
+/// Requires crate feature `"borsh"`
+impl<const CAP: usize> borsh::io::Write for ArrayVec<u8, CAP> {
+    fn write(&mut self, data: &[u8]) -> borsh::io::Result<usize> {
+        Ok(ArrayVec::write(self, data))
+    }
+
+    fn write_all(&mut self, mut buf: &[u8]) -> borsh::io::Result<()> {
+        while !buf.is_empty() {
+            match self.write(buf) {
+                0 => {
+                    return Err(borsh::io::Error::new(
+                        borsh::io::ErrorKind::WriteZero,
+                        "failed to write whole buffer",
+                    ));
+                }
+                n => buf = &buf[n..],
+            }
+        }
+        Ok(())
+    }
+
+    fn flush(&mut self) -> borsh::io::Result<()> {
+        Ok(())
+    }
+}
+
 #[cfg(feature = "borsh")]
 /// Requires crate feature `"borsh"`
 impl<T, const CAP: usize> borsh::BorshDeserialize for ArrayVec<T, CAP>

tests/tests.rs

@@ -490,7 +490,7 @@ fn test_write() {
     let mut v = ArrayVec::<_, 8>::new();
     write!(&mut v, "\x01\x02\x03").unwrap();
     assert_eq!(&v[..], &[1, 2, 3]);
-    let r = v.write(&[9; 16]).unwrap();
+    let r = v.write(&[9; 16]);
     assert_eq!(r, 5);
     assert_eq!(&v[..], &[1, 2, 3, 9, 9, 9, 9, 9]);
 }