Merge rust-bitcoin#5086: consensus_encoding: Add CompactSizeEncoder and refactor WitnessEncoder

3051b3e primitives: Refactor `WitnessEncoder` to use CompactSizeEncoder (jrakibi)
982a805 consensus_encoding: Export `CompactSizeEncoder` (jrakibi)
aa00732 consensus_encoding: Add `CompactSizeEncoder` implementation (jrakibi)

Pull request description:

  Add `CompactSizeEncoder` implementation and simplify `WitnessEncoder` by using the new `CompactSizeEncoder` + `Encoder2` composition.
  
  Discussed in: rust-bitcoin#4992 (comment)
  
  Closes: rust-bitcoin#5077


ACKs for top commit:
  nyonson:
    ACK 3051b3e
  apoelstra:
    ACK 3051b3e; successfully ran local tests
  tcharding:
    ACK 3051b3e


Tree-SHA512: f74cd5aec981e902b6ae62ffd43a70145ac6dcbc48182e7062abf40010f49af4ee6313ad8ea69447a9d4b2b43b07d9d896892a22ab521847328c16d94b5f588e

Author: apoelstra

consensus_encoding/src/encode/encoders.rs

@@ -12,7 +12,7 @@
 //!
 
 use internals::array_vec::ArrayVec;
-use internals::compact_size;
+use internals::{compact_size, ToU64};
 
 use super::{Encodable, Encoder};
 
@@ -250,6 +250,27 @@ impl<A: Encoder, B: Encoder, C: Encoder, D: Encoder, E: Encoder, F: Encoder> Enc
     fn advance(&mut self) -> bool { self.inner.advance() }
 }
 
+/// Encoder for a compact size encoded integer.
+pub struct CompactSizeEncoder {
+    buf: Option<ArrayVec<u8, SIZE>>,
+}
+
+impl CompactSizeEncoder {
+    /// Constructs a new `CompactSizeEncoder`.
+    pub fn new(value: impl ToU64) -> Self { Self { buf: Some(compact_size::encode(value)) } }
+}
+
+impl Encoder for CompactSizeEncoder {
+    #[inline]
+    fn current_chunk(&self) -> Option<&[u8]> { self.buf.as_ref().map(|b| &b[..]) }
+
+    #[inline]
+    fn advance(&mut self) -> bool {
+        self.buf = None;
+        false
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -651,4 +672,42 @@ mod tests {
         assert!(!encoder.advance());
         assert_eq!(encoder.current_chunk(), None);
     }
+    #[test]
+    fn encode_compact_size() {
+        // 1-byte
+        let mut e = CompactSizeEncoder::new(0x10u64);
+        assert_eq!(e.current_chunk(), Some(&[0x10][..]));
+        assert!(!e.advance());
+        assert_eq!(e.current_chunk(), None);
+
+        let mut e = CompactSizeEncoder::new(0xFCu64);
+        assert_eq!(e.current_chunk(), Some(&[0xFC][..]));
+        assert!(!e.advance());
+
+        // 0xFD + u16
+        let mut e = CompactSizeEncoder::new(0x00FDu64);
+        assert_eq!(e.current_chunk(), Some(&[0xFD, 0xFD, 0x00][..]));
+        assert!(!e.advance());
+
+        let mut e = CompactSizeEncoder::new(0x0FFFu64);
+        assert_eq!(e.current_chunk(), Some(&[0xFD, 0xFF, 0x0F][..]));
+        assert!(!e.advance());
+
+        // 0xFE + u32
+        let mut e = CompactSizeEncoder::new(0x0001_0000u64);
+        assert_eq!(e.current_chunk(), Some(&[0xFE, 0x00, 0x00, 0x01, 0x00][..]));
+        assert!(!e.advance());
+
+        let mut e = CompactSizeEncoder::new(0x0F0F_0F0Fu64);
+        assert_eq!(e.current_chunk(), Some(&[0xFE, 0x0F, 0x0F, 0x0F, 0x0F][..]));
+        assert!(!e.advance());
+
+        // 0xFF + u64
+        let mut e = CompactSizeEncoder::new(0x0000_F0F0_F0F0_F0E0u64);
+        assert_eq!(
+            e.current_chunk(),
+            Some(&[0xFF, 0xE0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0x00, 0x00][..])
+        );
+        assert!(!e.advance());
+    }
 }

consensus_encoding/src/lib.rs

@@ -41,6 +41,7 @@ pub use self::encode::encode_to_vec;
 #[cfg(feature = "std")]
 pub use self::encode::encode_to_writer;
 pub use self::encode::encoders::{
-    ArrayEncoder, BytesEncoder, Encoder2, Encoder3, Encoder4, Encoder6, SliceEncoder,
+    ArrayEncoder, BytesEncoder, CompactSizeEncoder, Encoder2, Encoder3, Encoder4, Encoder6,
+    SliceEncoder,
 };
 pub use self::encode::{encode_to_hash_engine, Encodable, Encoder};

primitives/src/witness.rs

@@ -9,10 +9,9 @@ use core::ops::Index;
 
 #[cfg(feature = "arbitrary")]
 use arbitrary::{Arbitrary, Unstructured};
-use encoding::{Encodable, Encoder};
+use encoding::{BytesEncoder, CompactSizeEncoder, Encodable, Encoder, Encoder2};
 #[cfg(feature = "hex")]
 use hex::{error::HexToBytesError, FromHex};
-use internals::array_vec::ArrayVec;
 use internals::compact_size;
 use internals::slice::SliceExt;
 use internals::wrap_debug::WrapDebug;
@@ -262,19 +261,8 @@ fn decode_cursor(bytes: &[u8], start_of_indices: usize, index: usize) -> Option<
     bytes.get_array::<4>(start).map(|index_bytes| u32::from_ne_bytes(*index_bytes) as usize)
 }
 
-/// The maximum length of a compact size encoding.
-const SIZE: usize = compact_size::MAX_ENCODING_SIZE;
-
 /// The encoder for the [`Witness`] type.
-// This is basically an exact copy of the `encoding::BytesEncoder` except we prefix
-// with the number of witness elements not the byte slice length.
-pub struct WitnessEncoder<'a> {
-    /// A slice of all the elements without the initial length prefix
-    /// but with the length prefix on each element.
-    witness_elements: Option<&'a [u8]>,
-    /// Encoding of the number of witness elements.
-    num_elements: Option<ArrayVec<u8, SIZE>>,
-}
+pub struct WitnessEncoder<'a>(Encoder2<CompactSizeEncoder, BytesEncoder<'a>>);
 
 impl Encodable for Witness {
     type Encoder<'a>
@@ -283,33 +271,20 @@ impl Encodable for Witness {
         Self: 'a;
 
     fn encoder(&self) -> Self::Encoder<'_> {
-        let num_elements = Some(compact_size::encode(self.len()));
-        let witness_elements = Some(&self.content[..self.indices_start]);
+        let num_elements = CompactSizeEncoder::new(self.len());
+        let witness_elements =
+            BytesEncoder::without_length_prefix(&self.content[..self.indices_start]);
 
-        WitnessEncoder { witness_elements, num_elements }
+        WitnessEncoder(Encoder2::new(num_elements, witness_elements))
     }
 }
 
 impl<'a> Encoder for WitnessEncoder<'a> {
     #[inline]
-    fn current_chunk(&self) -> Option<&[u8]> {
-        if let Some(num_elements) = self.num_elements.as_ref() {
-            Some(num_elements)
-        } else {
-            self.witness_elements
-        }
-    }
+    fn current_chunk(&self) -> Option<&[u8]> { self.0.current_chunk() }
 
     #[inline]
-    fn advance(&mut self) -> bool {
-        if self.num_elements.is_some() {
-            self.num_elements = None;
-            true
-        } else {
-            self.witness_elements = None;
-            false
-        }
-    }
+    fn advance(&mut self) -> bool { self.0.advance() }
 }
 
 // Note: we use `Borrow` in the following `PartialEq` impls specifically because of its additional