Add decode_in_place and benchmarks to b64ct (#206)

Author: newpavlov

b64ct/benches/mod.rs

@@ -0,0 +1,58 @@
+#![feature(test)]
+extern crate test;
+use test::Bencher;
+
+const B64_LEN: usize = 100_002;
+const RAW_LEN: usize = (3 * B64_LEN) / 4;
+
+#[inline(never)]
+fn get_raw_data() -> Vec<u8> {
+    (0..RAW_LEN).map(|i| i as u8).collect()
+}
+
+#[inline(never)]
+fn get_b64_data() -> String {
+    (0..B64_LEN)
+        .map(|i| match (i % 64) as u8 {
+            v @ 0..=25 => (v + 'A' as u8) as char,
+            v @ 26..=51 => (v - 26 + 'a' as u8) as char,
+            v @ 52..=61 => (v - 52 + '0' as u8) as char,
+            62 => '+',
+            _ => '/',
+        })
+        .collect()
+}
+
+#[bench]
+fn decode_bench(b: &mut Bencher) {
+    let b64_data = get_b64_data();
+    let mut buf = get_raw_data();
+    b.iter(|| {
+        let out = b64ct::decode(&b64_data, &mut buf).unwrap();
+        test::black_box(out);
+    });
+    b.bytes = RAW_LEN as u64;
+}
+
+#[bench]
+fn decode_in_place_bench(b: &mut Bencher) {
+    let mut b64_data = get_b64_data().into_bytes();
+    b.iter(|| {
+        // since it works on the same buffer over and over,
+        // almost always `out` will be an error
+        let out = b64ct::decode_in_place(&mut b64_data);
+        let _ = test::black_box(out);
+    });
+    b.bytes = RAW_LEN as u64;
+}
+
+#[bench]
+fn encode_bench(b: &mut Bencher) {
+    let mut buf = get_b64_data().into_bytes();
+    let raw_data = get_raw_data();
+    b.iter(|| {
+        let out = b64ct::encode(&raw_data, &mut buf).unwrap();
+        test::black_box(out);
+    });
+    b.bytes = RAW_LEN as u64;
+}

b64ct/src/lib.rs

@@ -123,6 +123,40 @@ pub fn decode<'a>(src: &str, dst: &'a mut [u8]) -> Result<&'a [u8], Error> {
     }
 }
 
+/// Decode B64-encoded string in-place.
+pub fn decode_in_place(buf: &mut [u8]) -> Result<&[u8], InvalidEncodingError> {
+    // TODO: make panic-free
+    let mut err: isize = 0;
+    let mut tmp_in = [0u8; 4];
+    let mut tmp_out = [0u8; 3];
+
+    let full_chunks = buf.len() / 4;
+
+    for chunk in 0..full_chunks {
+        tmp_in.copy_from_slice(&buf[4 * chunk..][..4]);
+        err |= decode_3bytes(&tmp_in, &mut tmp_out);
+        buf[3 * chunk..][..3].copy_from_slice(&tmp_out);
+    }
+
+    let dlen = decoded_len_inner(buf.len());
+    let src_rem_pos = 4 * full_chunks;
+    let src_rem_len = buf.len() - src_rem_pos;
+    let dst_rem_pos = 3 * full_chunks;
+    let dst_rem_len = dlen - dst_rem_pos;
+
+    err |= !(src_rem_len == 0 || src_rem_len >= 2) as isize;
+    tmp_in = [b'A'; 4];
+    tmp_in[..src_rem_len].copy_from_slice(&buf[src_rem_pos..]);
+    err |= decode_3bytes(&tmp_in, &mut tmp_out);
+    buf[dst_rem_pos..][..dst_rem_len].copy_from_slice(&tmp_out[..dst_rem_len]);
+
+    if err == 0 {
+        Ok(&buf[..dlen])
+    } else {
+        Err(InvalidEncodingError)
+    }
+}
+
 /// Decode a "B64"-encoded string into a byte vector.
 #[cfg(feature = "alloc")]
 #[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
@@ -139,7 +173,11 @@ pub fn decode_vec(input: &str) -> Result<Vec<u8>, InvalidEncodingError> {
 
 /// Get the length of the output from decoding the provided "B64"-encoded input.
 pub const fn decoded_len(bytes: &str) -> usize {
-    (bytes.len() * 3) / 4
+    decoded_len_inner(bytes.len())
+}
+
+const fn decoded_len_inner(n: usize) -> usize {
+    (n * 3) / 4
 }
 
 // B64 character set:

b64ct/tests/mod.rs

@@ -1,4 +1,4 @@
-use b64ct::{decode, decoded_len, encode, encoded_len, Error};
+use b64ct::{decode, decode_in_place, decoded_len, encode, encoded_len, Error};
 
 /// "B64" test vector
 struct TestVector {
@@ -53,10 +53,14 @@ fn decode_test_vectors() {
     let mut buf = [0u8; 1024];
 
     for vector in TEST_VECTORS {
-        println!("{:?}", vector.raw);
         let out = decode(vector.b64, &mut buf).unwrap();
         assert_eq!(decoded_len(vector.b64), out.len());
         assert_eq!(vector.raw, &out[..]);
+
+        let n = vector.b64.len();
+        buf[..n].copy_from_slice(vector.b64.as_bytes());
+        let out = decode_in_place(&mut buf[..n]).unwrap();
+        assert_eq!(vector.raw, out);
     }
 }
 
@@ -72,6 +76,11 @@ fn encode_and_decode_various_lengths() {
         // Make sure it round trips
         let decoded = decode(encoded, &mut outbuf).unwrap();
         assert_eq!(decoded, &data[..i]);
+
+        let elen = encode(&data[..i], &mut inbuf).unwrap().len();
+        let buf = &mut inbuf[..elen];
+        let decoded = decode_in_place(buf).unwrap();
+        assert_eq!(decoded, &data[..i]);
     }
 }