From 895b5cd866578a7052a5374ffe0918e01f29b51f Mon Sep 17 00:00:00 2001
From: Riccardo Casatta <riccardo@casatta.it>
Date: Mon, 13 Jun 2022 15:39:34 +0200
Subject: [PATCH] Implement iter::size_hint and ExactSizeIterator for Witness
 Iter

// with size_hint
// test blockdata::witness::benches::bench_big_witness_to_vec              ... bench:         313 ns/iter (+/- 13)
// test blockdata::witness::benches::bench_witness_to_vec                  ... bench:         204 ns/iter (+/- 11)

// without
// test blockdata::witness::benches::bench_big_witness_to_vec              ... bench:         489 ns/iter (+/- 28)
// test blockdata::witness::benches::bench_witness_to_vec                  ... bench:         221 ns/iter (+/- 102)
---
 src/blockdata/witness.rs | 38 ++++++++++++++++++++++++++++++++++++--
 1 file changed, 36 insertions(+), 2 deletions(-)

diff --git a/src/blockdata/witness.rs b/src/blockdata/witness.rs
index e40fe8aab1..5bb1c0dab0 100644
--- a/src/blockdata/witness.rs
+++ b/src/blockdata/witness.rs
@@ -43,7 +43,7 @@ pub struct Witness {
 }
 
 /// Support structure to allow efficient and convenient iteration over the Witness elements
-pub struct Iter<'a>(::core::slice::Iter<'a, u8>);
+pub struct Iter<'a>(::core::slice::Iter<'a, u8>, usize);
 
 impl Decodable for Witness {
     fn consensus_decode<D: Read>(mut d: D) -> Result<Self, Error> {
@@ -172,7 +172,7 @@ impl Witness {
 
     /// Returns a struct implementing [`Iterator`]
     pub fn iter(&self) -> Iter {
-        Iter(self.content.iter())
+        Iter(self.content.iter(), self.witness_elements)
     }
 
     /// Returns the number of elements this witness holds
@@ -263,8 +263,14 @@ impl<'a> Iterator for Iter<'a> {
         }
         Some(slice)
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (self.1, Some(self.1))
+    }
 }
 
+impl<'a> ExactSizeIterator for Iter<'a> {}
+
 // Serde keep backward compatibility with old Vec<Vec<u8>> format
 #[cfg(feature = "serde")]
 impl serde::Serialize for Witness {
@@ -408,3 +414,31 @@ mod test {
         assert_eq!(new_witness_format, back);
     }
 }
+
+
+#[cfg(all(test, feature = "unstable"))]
+mod benches {
+    use test::{Bencher, black_box};
+    use super::Witness;
+
+    #[bench]
+    pub fn bench_big_witness_to_vec(bh: &mut Bencher) {
+        let raw_witness = vec![vec![1u8]; 5];
+        let witness = Witness::from_vec(raw_witness);
+
+        bh.iter(|| {
+            black_box(witness.to_vec());
+        });
+    }
+
+    #[bench]
+    pub fn bench_witness_to_vec(bh: &mut Bencher) {
+        let raw_witness = vec![vec![1u8]; 3];
+        let witness = Witness::from_vec(raw_witness);
+
+        bh.iter(|| {
+            black_box(witness.to_vec());
+        });
+    }
+
+}