rework encrypt to take IntoIterator<AsRef<[u8]>>

The trick is to use an intermediate enum which can either contain a
slice (the local tau block/nonce) or the passed in AsRef. This way, we
unify the two types.

1 files changed, 18 insertions, 4 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 421d800..d649dfc 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -135,6 +135,20 @@ enum Mode {
     Decipher,
 }
 
+enum SliceOrAsRef<'a, A> {
+    Slice(&'a [u8]),
+    AsRef(A),
+}
+
+impl<'a, A: AsRef<[u8]>> AsRef<[u8]> for SliceOrAsRef<'a, A> {
+    fn as_ref(&self) -> &[u8] {
+        match self {
+            SliceOrAsRef::Slice(x) => *x,
+            SliceOrAsRef::AsRef(x) => x.as_ref(),
+        }
+    }
+}
+
 /// AEZ encryption scheme.
 ///
 /// See the [module level documentation](index.html) for more information.
@@ -335,7 +349,7 @@ fn append_auth(data_len: usize, buffer: &mut [u8]) {
     }
 }
 
-fn encrypt<'t, A: AsRef<[u8]> + 't, T: IntoIterator<Item = &'t A>>(
+fn encrypt<A: AsRef<[u8]>, T: IntoIterator<Item = A>>(
     aez: &Aez,
     nonce: &[u8],
     ad: T,
@@ -345,9 +359,9 @@ fn encrypt<'t, A: AsRef<[u8]> + 't, T: IntoIterator<Item = &'t A>>(
     // 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 tweaks = iter::once(&tau_bytes as &[_])
-        .chain(iter::once(nonce))
-        .chain(ad.into_iter().map(|r| r.as_ref()));
+    let tweaks = iter::once(SliceOrAsRef::Slice(&tau_bytes))
+        .chain(iter::once(SliceOrAsRef::Slice(nonce)))
+        .chain(ad.into_iter().map(SliceOrAsRef::AsRef));
     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