From 9687b036c37c15de7405597e9ad52a713d2b9704 Mon Sep 17 00:00:00 2001
From: Marshall Pierce <575695+marshallpierce@users.noreply.github.com>
Date: Sun, 5 Feb 2023 09:35:53 -0700
Subject: [PATCH] Switch add_padding() to use encoded len

Simpler arithmetic offers a small speedup on short inputs.
---
 src/chunked_encoder.rs |  3 ++-
 src/encode.rs          | 28 ++++++++++++++--------------
 src/engine/tests.rs    | 13 ++++++++-----
 3 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/src/chunked_encoder.rs b/src/chunked_encoder.rs
index 0457259..bc3810a 100644
--- a/src/chunked_encoder.rs
+++ b/src/chunked_encoder.rs
@@ -49,7 +49,8 @@ impl<'e, E: Engine + ?Sized> ChunkedEncoder<'e, E> {
             if self.engine.config().encode_padding() && !more_input_left {
                 // no more input, add padding if needed. Buffer will have room because
                 // max_input_length leaves room for it.
-                b64_bytes_written += add_padding(bytes.len(), &mut encode_buf[b64_bytes_written..]);
+                b64_bytes_written +=
+                    add_padding(b64_bytes_written, &mut encode_buf[b64_bytes_written..]);
             }
 
             sink.write_encoded_bytes(&encode_buf[0..b64_bytes_written])?;
diff --git a/src/encode.rs b/src/encode.rs
index cb17650..15b903d 100644
--- a/src/encode.rs
+++ b/src/encode.rs
@@ -77,7 +77,7 @@ pub(crate) fn encode_with_padding<E: Engine + ?Sized>(
     let b64_bytes_written = engine.internal_encode(input, output);
 
     let padding_bytes = if engine.config().encode_padding() {
-        add_padding(input.len(), &mut output[b64_bytes_written..])
+        add_padding(b64_bytes_written, &mut output[b64_bytes_written..])
     } else {
         0
     };
@@ -117,20 +117,20 @@ pub fn encoded_len(bytes_len: usize, padding: bool) -> Option<usize> {
 }
 
 /// Write padding characters.
-/// `input_len` is the size of the original, not encoded, input.
+/// `unpadded_output_len` is the size of the unpadded but base64 encoded data.
 /// `output` is the slice where padding should be written, of length at least 2.
 ///
 /// Returns the number of padding bytes written.
-pub(crate) fn add_padding(input_len: usize, output: &mut [u8]) -> usize {
-    // TODO base on encoded len to use cheaper mod by 4 (aka & 7)
-    let rem = input_len % 3;
-    let mut bytes_written = 0;
-    for _ in 0..((3 - rem) % 3) {
-        output[bytes_written] = PAD_BYTE;
-        bytes_written += 1;
+pub(crate) fn add_padding(unpadded_output_len: usize, output: &mut [u8]) -> usize {
+    let pad_bytes = (4 - (unpadded_output_len % 4)) % 4;
+    // for just a couple bytes, this has better performance than using
+    // .fill(), or iterating over mutable refs, which call memset()
+    #[allow(clippy::needless_range_loop)]
+    for i in 0..pad_bytes {
+        output[i] = PAD_BYTE;
     }
 
-    bytes_written
+    pad_bytes
 }
 
 /// Errors that can occur while encoding into a slice.
@@ -434,18 +434,18 @@ mod tests {
 
         let mut rng = rand::rngs::SmallRng::from_entropy();
 
-        // cover our bases for length % 3
-        for input_len in 0..10 {
+        // cover our bases for length % 4
+        for unpadded_output_len in 0..20 {
             output.clear();
 
             // fill output with random
-            for _ in 0..10 {
+            for _ in 0..100 {
                 output.push(rng.gen());
             }
 
             let orig_output_buf = output.clone();
 
-            let bytes_written = add_padding(input_len, &mut output);
+            let bytes_written = add_padding(unpadded_output_len, &mut output);
 
             // make sure the part beyond bytes_written is the same garbage it was before
             assert_eq!(orig_output_buf[bytes_written..], output[bytes_written..]);
diff --git a/src/engine/tests.rs b/src/engine/tests.rs
index 2b4d2e5..d2851ec 100644
--- a/src/engine/tests.rs
+++ b/src/engine/tests.rs
@@ -86,7 +86,7 @@ fn rfc_test_vectors_std_alphabet<E: EngineWrapper>(engine_wrapper: E) {
                 &encoded_without_padding,
                 &std::str::from_utf8(&encode_buf[0..encode_len]).unwrap()
             );
-            let pad_len = add_padding(orig.len(), &mut encode_buf[encode_len..]);
+            let pad_len = add_padding(encode_len, &mut encode_buf[encode_len..]);
             assert_eq!(encoded.as_bytes(), &encode_buf[..encode_len + pad_len]);
 
             let decode_len = engine
@@ -195,7 +195,10 @@ fn encode_doesnt_write_extra_bytes<E: EngineWrapper>(engine_wrapper: E) {
 
         // pad so we can decode it in case our random engine requires padding
         let pad_len = if padded {
-            add_padding(orig_len, &mut encode_buf[prefix_len + encoded_len_no_pad..])
+            add_padding(
+                encoded_len_no_pad,
+                &mut encode_buf[prefix_len + encoded_len_no_pad..],
+            )
         } else {
             0
         };
@@ -382,7 +385,7 @@ fn decode_detect_invalid_last_symbol_every_possible_two_symbols<E: EngineWrapper
     for b in 0_u8..=255 {
         let mut b64 = vec![0_u8; 4];
         assert_eq!(2, engine.internal_encode(&[b], &mut b64[..]));
-        let _ = add_padding(1, &mut b64[2..]);
+        let _ = add_padding(2, &mut b64[2..]);
 
         assert!(base64_to_bytes.insert(b64, vec![b]).is_none());
     }
@@ -442,7 +445,7 @@ fn decode_detect_invalid_last_symbol_every_possible_three_symbols<E: EngineWrapp
             bytes[1] = b2;
             let mut b64 = vec![0_u8; 4];
             assert_eq!(3, engine.internal_encode(&bytes, &mut b64[..]));
-            let _ = add_padding(2, &mut b64[3..]);
+            let _ = add_padding(3, &mut b64[3..]);
 
             let mut v = Vec::with_capacity(2);
             v.extend_from_slice(&bytes[..]);
@@ -1245,7 +1248,7 @@ fn generate_random_encoded_data<E: Engine, R: rand::Rng, D: distributions::Distr
     let base_encoded_len = engine.internal_encode(&orig_data[..], &mut encode_buf[..]);
 
     let enc_len_with_padding = if padding {
-        base_encoded_len + add_padding(orig_len, &mut encode_buf[base_encoded_len..])
+        base_encoded_len + add_padding(base_encoded_len, &mut encode_buf[base_encoded_len..])
     } else {
         base_encoded_len
     };