inline pad_to_blocks

I like the way with chunks_exact more, as we get the remainder for free
and we don't have to do the chunk.len() dance

2 files changed, 19 insertions, 17 deletions

diff --git a/src/block.rs b/src/block.rs
index e1140e7..2b7907a 100644
--- a/src/block.rs
+++ b/src/block.rs
@@ -94,6 +94,16 @@ impl Block {
         result
     }
 
+    /// Pad the block to full length.
+    ///
+    /// Unlike [`pad`], this function takes the length in bytes.
+    pub fn pad_bytes(&self, length: u8) -> Block {
+        assert!(length <= 15);
+        let mut result = *self;
+        result[length as usize] = 0x80;
+        result
+    }
+
     /// Clip the block by setting all bits beyond the given length to 0.
     pub fn clip(&self, length: usize) -> Block {
         match length {
diff --git a/src/lib.rs b/src/lib.rs
index 95c1622..41b8d72 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -688,16 +688,6 @@ fn cipher_aez_core<A: AsRef<[u8]>, T: IntoIterator<Item = A>>(
     blocks.set_m_y(c_y);
 }
 
-fn pad_to_blocks(value: &[u8]) -> impl Iterator<Item = Block> {
-    value.chunks(16).map(|chunk| {
-        if chunk.len() == 16 {
-            Block::from_slice(chunk)
-        } else {
-            Block::from_slice(chunk).pad(chunk.len() * 8)
-        }
-    })
-}
-
 fn aez_hash<A: AsRef<[u8]>, T: IntoIterator<Item = A>>(aez: &Aez, tweaks: T) -> Block {
     let mut hash = Block::null();
     for (i, tweak) in tweaks.into_iter().enumerate() {
@@ -716,15 +706,17 @@ fn aez_hash<A: AsRef<[u8]>, T: IntoIterator<Item = A>>(aez: &Aez, tweaks: T) ->
                 hash = hash ^ ej.eval(Block::from_slice(chunk));
             }
         } else {
-            let blocks = pad_to_blocks(tweak);
-            for (l, chunk) in blocks.enumerate() {
+            let blocks = tweak.chunks_exact(16);
+            let remainder = blocks.remainder();
+
+            for chunk in blocks {
                 ej.advance();
-                if l == tweak.len() / 16 {
-                    hash = hash ^ e(j.try_into().unwrap(), 0, aez, chunk);
-                } else {
-                    hash = hash ^ ej.eval(chunk);
-                }
+                hash = hash ^ ej.eval(Block::from_slice(chunk));
             }
+
+            ej.advance();
+            let chunk = Block::from_slice(remainder).pad_bytes(remainder.len() as u8);
+            hash = hash ^ e(j.try_into().unwrap(), 0, aez, chunk);
         }
     }
     hash