From c12a9a86ae0e0b712ef726dd9d8957b11e58e4f2 Mon Sep 17 00:00:00 2001
From: Carl Lerche <me@carllerche.com>
Date: Tue, 13 Jun 2017 14:39:39 -0700
Subject: [PATCH] Don't allocate on each call to huffman decode

---
 src/hpack/decoder.rs     | 64 +++++++++++++++++++++-------------------
 src/hpack/encoder.rs     |  3 +-
 src/hpack/huffman/mod.rs | 17 ++++++-----
 3 files changed, 45 insertions(+), 39 deletions(-)

diff --git a/src/hpack/decoder.rs b/src/hpack/decoder.rs
index bbe0884..36207fb 100644
--- a/src/hpack/decoder.rs
+++ b/src/hpack/decoder.rs
@@ -2,7 +2,7 @@ use super::{huffman, header as h2_header, Header};
 use util::byte_str::FromUtf8Error;
 
 use http::{method, header, status, StatusCode, Method};
-use bytes::{Buf, Bytes};
+use bytes::{Buf, Bytes, BytesMut};
 
 use std::cmp;
 use std::io::Cursor;
@@ -14,6 +14,7 @@ pub struct Decoder {
     max_size_update: Option<usize>,
     last_max_update: usize,
     table: Table,
+    buffer: BytesMut,
 }
 
 /// Represents all errors that can be encountered while performing the decoding
@@ -139,6 +140,7 @@ impl Decoder {
             max_size_update: None,
             last_max_update: size,
             table: Table::new(size),
+            buffer: BytesMut::with_capacity(4096),
         }
     }
 
@@ -234,7 +236,7 @@ impl Decoder {
         self.table.get(index)
     }
 
-    fn decode_literal(&self, buf: &mut Cursor<&Bytes>, index: bool)
+    fn decode_literal(&mut self, buf: &mut Cursor<&Bytes>, index: bool)
         -> Result<Header, DecoderError>
     {
         let prefix = if index {
@@ -249,17 +251,43 @@ impl Decoder {
         // First, read the header name
         if table_idx == 0 {
             // Read the name as a literal
-            let name = try!(decode_string(buf));
-            let value = try!(decode_string(buf));
+            let name = try!(self.decode_string(buf));
+            let value = try!(self.decode_string(buf));
 
             Header::new(name, value)
         } else {
             let e = try!(self.table.get(table_idx));
-            let value = try!(decode_string(buf));
+            let value = try!(self.decode_string(buf));
 
             e.name().into_entry(value)
         }
     }
+
+    fn decode_string(&mut self, buf: &mut Cursor<&Bytes>) -> Result<Bytes, DecoderError> {
+        const HUFF_FLAG: u8 = 0b10000000;
+
+        // The first bit in the first byte contains the huffman encoded flag.
+        let huff = peek_u8(buf) & HUFF_FLAG == HUFF_FLAG;
+
+        // Decode the string length using 7 bit prefix
+        let len = try!(decode_int(buf, 7));
+
+        if len > buf.remaining() {
+            return Err(DecoderError::StringUnderflow);
+        }
+
+        if huff {
+            let ret = {
+                let raw = &buf.bytes()[..len];
+                huffman::decode(raw, &mut self.buffer).map(Into::into)
+            };
+
+            buf.advance(len);
+            return ret;
+        } else {
+            Ok(take(buf, len))
+        }
+    }
 }
 
 impl Default for Decoder {
@@ -356,32 +384,6 @@ fn decode_int<B: Buf>(buf: &mut B, prefix_size: u8) -> Result<usize, DecoderErro
     Err(DecoderError::IntegerUnderflow)
 }
 
-fn decode_string(buf: &mut Cursor<&Bytes>) -> Result<Bytes, DecoderError> {
-    const HUFF_FLAG: u8 = 0b10000000;
-
-    // The first bit in the first byte contains the huffman encoded flag.
-    let huff = peek_u8(buf) & HUFF_FLAG == HUFF_FLAG;
-
-    // Decode the string length using 7 bit prefix
-    let len = try!(decode_int(buf, 7));
-
-    if len > buf.remaining() {
-        return Err(DecoderError::StringUnderflow);
-    }
-
-    if huff {
-        let ret = {
-            let raw = &buf.bytes()[..len];
-            huffman::decode(raw).map(Into::into)
-        };
-
-        buf.advance(len);
-        return ret;
-    } else {
-        Ok(take(buf, len))
-    }
-}
-
 fn peek_u8<B: Buf>(buf: &mut B) -> u8 {
     buf.bytes()[0]
 }
diff --git a/src/hpack/encoder.rs b/src/hpack/encoder.rs
index b8d0b3e..a689452 100644
--- a/src/hpack/encoder.rs
+++ b/src/hpack/encoder.rs
@@ -664,6 +664,7 @@ mod test {
     }
 
     fn huff_decode(src: &[u8]) -> BytesMut {
-        huffman::decode(src).unwrap()
+        let mut buf = BytesMut::new();
+        huffman::decode(src, &mut buf).unwrap()
     }
 }
diff --git a/src/hpack/huffman/mod.rs b/src/hpack/huffman/mod.rs
index 9c172f0..fee3163 100644
--- a/src/hpack/huffman/mod.rs
+++ b/src/hpack/huffman/mod.rs
@@ -17,21 +17,19 @@ const MAYBE_EOS: u8 = 1;
 const DECODED: u8 = 2;
 const ERROR: u8 = 4;
 
-pub fn decode(src: &[u8]) -> Result<BytesMut, DecoderError> {
-    // TODO: This should not allocate and instead take a dst
-
+pub fn decode(src: &[u8], buf: &mut BytesMut) -> Result<BytesMut, DecoderError> {
     let mut decoder = Decoder::new();
 
     // Max compression ratio is >= 0.5
-    let mut dst = BytesMut::with_capacity(src.len() << 1);
+    buf.reserve(src.len() << 1);
 
     for b in src {
         if let Some(b) = try!(decoder.decode4(b >> 4)) {
-            dst.put_u8(b);
+            buf.put_u8(b);
         }
 
         if let Some(b) = try!(decoder.decode4(b & 0xf)) {
-            dst.put_u8(b);
+            buf.put_u8(b);
         }
     }
 
@@ -39,7 +37,7 @@ pub fn decode(src: &[u8]) -> Result<BytesMut, DecoderError> {
         return Err(DecoderError::InvalidHuffmanCode);
     }
 
-    Ok(dst)
+    Ok(buf.take())
 }
 
 // TODO: return error when there is not enough room to encode the value
@@ -119,6 +117,11 @@ impl Decoder {
 mod test {
     use super::*;
 
+    fn decode(src: &[u8]) -> Result<BytesMut, DecoderError> {
+        let mut buf = BytesMut::new();
+        super::decode(src, &mut buf)
+    }
+
     #[test]
     fn decode_single_byte() {
         assert_eq!("o", decode(&[0b00111111]).unwrap());