     return (r_bit >> r_bit_count) & ((1 << nbits) - 1);
 }
 
+/**
+ * Gets bits without removing them from the buffer.
+ */
+__device__ inline unsigned int
+gpujpeg_huffman_gpu_decoder_peek_bits(
+                const unsigned int nbits, unsigned int & r_bit, unsigned int & r_bit_count, uint4 * const s_byte, 
+                unsigned int & s_byte_idx, const uint4 * & d_byte, unsigned int & d_byte_chunk_count)
+{
+    // load bits into the register if haven't got enough
+    gpujpeg_huffman_gpu_decoder_load_bits(nbits, r_bit, r_bit_count, s_byte, s_byte_idx, d_byte, d_byte_chunk_count);
+    
+    // return bits 
+    return (r_bit >> (r_bit_count - nbits)) & ((1 << nbits) - 1);
+}
+
+/**
+ * Removes some bits from the buffer (assumes that they are there)
+ */
+__device__ inline void
+gpujpeg_huffman_gpu_decoder_discard_bits(const unsigned int nb, unsigned int, unsigned int & r_bit_count) {
+    r_bit_count -= nb;
+}
 
 /**
  * Special Huffman decode:

{
    // Section F.2.2.1: decode the DC coefficient difference
    // get dc category number, s
//     int s = gpujpeg_huffman_gpu_decoder_get_category(get_bits, get_buff, data, data_size, table_dc);
    int s = gpujpeg_huffman_gpu_decoder_decode_special_decode(table_dc, 1, r_bit, r_bit_count, s_byte, s_byte_idx, d_byte, d_byte_chunk_count);
    if ( s ) {
        // Get offset in this dc category
        int r = gpujpeg_huffman_gpu_decoder_get_bits(s, r_bit, r_bit_count, s_byte, s_byte_idx, d_byte, d_byte_chunk_count);

    // Since zeroes are skipped, output area must be cleared beforehand
    for ( int k = 1; k < 64; k++ ) {
        // s: (run, category)
        int s = gpujpeg_huffman_gpu_decoder_get_category(r_bit, r_bit_count, s_byte, s_byte_idx, d_byte, d_byte_chunk_count, table_ac);
        // r: run length for ac zero, 0 <= r < 16
        int r = s >> 4;
        // s: category for this non-zero ac
