From 04571021fb5490d0f0008c5b5a968f221de159a0 Mon Sep 17 00:00:00 2001
From: Jack O'Connor <oconnor663@gmail.com>
Date: Thu, 4 Nov 2021 20:37:05 -0400
Subject: [PATCH] add Hasher::count

---
 src/lib.rs              |  5 +++++
 test_vectors/src/lib.rs | 15 +++++++++------
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index a5fbac9..4976c29 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1301,6 +1301,11 @@ impl Hasher {
     pub fn finalize_xof(&self) -> OutputReader {
         OutputReader::new(self.final_output())
     }
+
+    /// Return the total number of bytes hashed so far.
+    pub fn count(&self) -> u64 {
+        self.chunk_state.chunk_counter * CHUNK_LEN as u64 + self.chunk_state.len() as u64
+    }
 }
 
 // Don't derive(Debug), because the state may be secret.
diff --git a/test_vectors/src/lib.rs b/test_vectors/src/lib.rs
index ec769d1..eea5528 100644
--- a/test_vectors/src/lib.rs
+++ b/test_vectors/src/lib.rs
@@ -240,8 +240,9 @@ mod tests {
     ) {
         let mut out = vec![0; expected_hash.len()];
         let mut hasher = blake3::Hasher::new();
-        for &b in input {
-            hasher.update(&[b]);
+        for i in 0..input.len() {
+            hasher.update(&[input[i]]);
+            assert_eq!(i as u64 + 1, hasher.count());
         }
         hasher.finalize_xof().fill(&mut out);
         assert_eq!(expected_hash, &out[..]);
@@ -249,8 +250,9 @@ mod tests {
 
         let mut out = vec![0; expected_keyed_hash.len()];
         let mut hasher = blake3::Hasher::new_keyed(key);
-        for &b in input {
-            hasher.update(&[b]);
+        for i in 0..input.len() {
+            hasher.update(&[input[i]]);
+            assert_eq!(i as u64 + 1, hasher.count());
         }
         hasher.finalize_xof().fill(&mut out);
         assert_eq!(expected_keyed_hash, &out[..]);
@@ -258,8 +260,9 @@ mod tests {
 
         let mut out = vec![0; expected_derive_key.len()];
         let mut hasher = blake3::Hasher::new_derive_key(TEST_CONTEXT);
-        for &b in input {
-            hasher.update(&[b]);
+        for i in 0..input.len() {
+            hasher.update(&[input[i]]);
+            assert_eq!(i as u64 + 1, hasher.count());
         }
         hasher.finalize_xof().fill(&mut out);
         assert_eq!(expected_derive_key, &out[..]);