From 71657b4f182c5b6851a0591fe054b7498f88675e Mon Sep 17 00:00:00 2001
From: Kris Nuttycombe <kris@nutty.land>
Date: Thu, 10 Mar 2022 11:22:24 -0700
Subject: [PATCH] Generalize vector and array writes & reads.

In a number of places, we transform other kinds of collections with
known length information into vectors simply to be able to use them with
`Vector::write` or `Vector::read`. We can avoid these extra allocations
by writing from iterators directly, and similarly by reading directly
into our desired collection types.
---
 components/zcash_encoding/src/lib.rs          | 70 ++++++++++++++++---
 .../src/merkle_tree/incremental.rs            |  2 -
 2 files changed, 59 insertions(+), 13 deletions(-)

diff --git a/components/zcash_encoding/src/lib.rs b/components/zcash_encoding/src/lib.rs
index 58c3e5389..79049b48a 100644
--- a/components/zcash_encoding/src/lib.rs
+++ b/components/zcash_encoding/src/lib.rs
@@ -12,6 +12,7 @@ use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 use nonempty::NonEmpty;
 use std::convert::TryFrom;
 use std::io::{self, Read, Write};
+use std::iter::FromIterator;
 
 /// The maximum allowed value representable as a `[CompactSize]`
 pub const MAX_COMPACT_SIZE: u32 = 0x02000000;
@@ -104,23 +105,34 @@ pub struct Vector;
 impl Vector {
     /// Reads a vector, assuming the encoding written by [`Vector::write`], using the provided
     /// function to decode each element of the vector.
-    pub fn read<R: Read, E, F>(mut reader: R, func: F) -> io::Result<Vec<E>>
+    pub fn read<R: Read, E, F>(reader: R, func: F) -> io::Result<Vec<E>>
+    where
+        F: Fn(&mut R) -> io::Result<E>,
+    {
+        Self::read_collected(reader, func)
+    }
+
+    /// Reads a CompactSize-prefixed series of elements into a collection, assuming the encoding
+    /// written by [`Vector::write`], using the provided function to decode each element.
+    pub fn read_collected<R: Read, E, F, O: FromIterator<E>>(
+        mut reader: R,
+        func: F,
+    ) -> io::Result<O>
     where
         F: Fn(&mut R) -> io::Result<E>,
     {
         let count: usize = CompactSize::read_t(&mut reader)?;
-        Array::read(reader, count, func)
+        Array::read_collected(reader, count, func)
     }
 
-    /// Writes a slice of values by writing [`CompactSize`]-encoded integer specifying the length of
-    /// the slice to the stream, followed by the encoding of each element of the slice as performed
-    /// by the provided function.
-    pub fn write<W: Write, E, F>(mut writer: W, vec: &[E], func: F) -> io::Result<()>
+    /// Writes a slice of values by writing [`CompactSize`]-encoded integer specifying the length
+    /// of the slice to the stream, followed by the encoding of each element of the slice as
+    /// performed by the provided function.
+    pub fn write<W: Write, E, F>(writer: W, vec: &[E], func: F) -> io::Result<()>
     where
         F: Fn(&mut W, &E) -> io::Result<()>,
     {
-        CompactSize::write(&mut writer, vec.len())?;
-        vec.iter().try_for_each(|e| func(&mut writer, e))
+        Self::write_sized(writer, vec.iter(), func)
     }
 
     /// Writes a NonEmpty container of values to the stream using the same encoding as
@@ -136,6 +148,29 @@ impl Vector {
         CompactSize::write(&mut writer, vec.len())?;
         vec.iter().try_for_each(|e| func(&mut writer, e))
     }
+
+    /// Writes at most `items_to_write` values from the provided iterator to the stream
+    /// in [`CompactSize`]-prefixed format.
+    ///
+    /// If not enough items are available, this will return an error; all available
+    /// elements will already have been written to the stream but the [`CompactSize`]
+    /// prefix that was written will be incorrect, and so the data written will not be
+    /// able to be correctly decoded as a vector.
+    pub fn write_sized<'a, W: Write, E: 'a, F, I: Iterator<Item = E> + ExactSizeIterator>(
+        mut writer: W,
+        items: I,
+        func: F,
+    ) -> io::Result<()>
+    where
+        F: Fn(&mut W, E) -> io::Result<()>,
+    {
+        CompactSize::write(&mut writer, items.len())?;
+        for item in items {
+            func(&mut writer, item)?;
+        }
+
+        Ok(())
+    }
 }
 
 /// Namespace for functions that perform encoding of array contents.
@@ -146,9 +181,22 @@ impl Vector {
 pub struct Array;
 
 impl Array {
-    /// Reads a vector, assuming the encoding written by [`Array::write`], using the provided
-    /// function to decode each element of the vector.
-    pub fn read<R: Read, E, F>(mut reader: R, count: usize, func: F) -> io::Result<Vec<E>>
+    /// Reads `count` elements from a stream into a vector, assuming the encoding written by
+    /// [`Array::write`], using the provided function to decode each element.
+    pub fn read<R: Read, E, F>(reader: R, count: usize, func: F) -> io::Result<Vec<E>>
+    where
+        F: Fn(&mut R) -> io::Result<E>,
+    {
+        Self::read_collected(reader, count, func)
+    }
+
+    /// Reads `count` elements into a collection, assuming the encoding written by
+    /// [`Array::write`], using the provided function to decode each element.
+    pub fn read_collected<R: Read, E, F, O: FromIterator<E>>(
+        mut reader: R,
+        count: usize,
+        func: F,
+    ) -> io::Result<O>
     where
         F: Fn(&mut R) -> io::Result<E>,
     {
diff --git a/zcash_primitives/src/merkle_tree/incremental.rs b/zcash_primitives/src/merkle_tree/incremental.rs
index 57946b035..c6cc09d58 100644
--- a/zcash_primitives/src/merkle_tree/incremental.rs
+++ b/zcash_primitives/src/merkle_tree/incremental.rs
@@ -59,8 +59,6 @@ pub fn read_leu64_usize<R: Read>(mut reader: R) -> io::Result<usize> {
 }
 
 pub fn write_position<W: Write>(mut writer: W, position: Position) -> io::Result<()> {
-    // The following `unwrap` is safe because we should never encounter
-    // a position that can't fit into a u64.
     write_usize_leu64(&mut writer, position.into())
 }