small cleanup in fill_one_block

1 files changed, 12 insertions, 9 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 7239edb..9309249 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1634,19 +1634,21 @@ impl OutputReader {
         }
     }
 
-    // It would be more natural if this helper took &mut &mut [u8], but that doesn't seem to be allowed.
-    fn fill_one_block<'a>(&mut self, mut buf: &'a mut [u8]) -> &'a mut [u8] {
-        let block: [u8; BLOCK_LEN] = self.inner.root_output_block();
-        let output_bytes = &block[self.position_within_block as usize..];
+    // This helper function handles both the case where the output buffer is
+    // shorter than one block, and the case where our position_within_block is
+    // non-zero.
+    fn fill_one_block(&mut self, buf: &mut &mut [u8]) {
+        let output_block: [u8; BLOCK_LEN] = self.inner.root_output_block();
+        let output_bytes = &output_block[self.position_within_block as usize..];
         let take = cmp::min(buf.len(), output_bytes.len());
         buf[..take].copy_from_slice(&output_bytes[..take]);
-        buf = &mut buf[take..];
         self.position_within_block += take as u8;
         if self.position_within_block == BLOCK_LEN as u8 {
             self.inner.counter += 1;
             self.position_within_block = 0;
         }
-        buf
+        // Advance the dest buffer. mem::take() is a borrowck workaround.
+        *buf = &mut core::mem::take(buf)[take..];
     }
 
     /// Fill a buffer with output bytes and advance the position of the
@@ -1671,10 +1673,11 @@ impl OutputReader {
 
         // If we're partway through a block, try to get to a block boundary.
         if self.position_within_block != 0 {
-            buf = self.fill_one_block(buf);
+            self.fill_one_block(&mut buf);
         }
 
         let full_blocks = buf.len() / BLOCK_LEN;
+        let full_blocks_len = full_blocks * BLOCK_LEN;
         if full_blocks > 0 {
             debug_assert_eq!(0, self.position_within_block);
             self.inner.platform.xof_many(
@@ -1683,7 +1686,7 @@ impl OutputReader {
                 self.inner.block_len,
                 self.inner.counter,
                 self.inner.flags | ROOT,
-                buf,
+                &mut buf[..full_blocks_len],
             );
             self.inner.counter += full_blocks as u64;
             buf = &mut buf[full_blocks * BLOCK_LEN..];
@@ -1691,7 +1694,7 @@ impl OutputReader {
 
         if !buf.is_empty() {
             debug_assert!(buf.len() < BLOCK_LEN);
-            buf = self.fill_one_block(buf);
+            self.fill_one_block(&mut buf);
             debug_assert!(buf.is_empty());
         }
     }