1 files changed, 13 insertions, 2 deletions
diff --git a/src/lib.rs b/src/lib.rs
index 6e411a0..a362946 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -299,8 +299,19 @@ fn encrypt(aez: &Aez, nonce: &[u8], ad: &[&[u8]], tau: u32, buffer: &mut [u8]) {
     // We treat tau as bytes, but according to the spec, tau is actually in bits.
     let tau_block = Block::from_int(tau as u128 * 8);
     let tau_bytes = tau_block.bytes();
-    let mut tweaks = vec![&tau_bytes, nonce];
-    tweaks.extend(ad);
+    let mut tweaks_vec;
+    // We optimize for the common case of having no associated data, or having one item of
+    // associated data (which is all the reference implementation supports anyway). If there's more
+    // associated data, we cave in and allocate a vec.
+    let tweaks = match ad.len() {
+        0 => &[&tau_bytes, nonce] as &[&[u8]],
+        1 => &[&tau_bytes, nonce, ad[0]],
+        _ => {
+            tweaks_vec = vec![&tau_bytes, nonce];
+            tweaks_vec.extend(ad);
+            &tweaks_vec
+        },
+    };
     assert!(buffer.len() >= tau as usize);
     if buffer.len() == tau as usize {
         // As aez_prf only xor's the input in, we have to clear the buffer first