Fix bitstream implementation

4 files changed, 39 insertions, 35 deletions

diff --git a/lisiblepng/meson.build b/lisiblepng/meson.build
index 22510fe..ae70175 100644
--- a/lisiblepng/meson.build
+++ b/lisiblepng/meson.build
@@ -1,6 +1,10 @@
 cc = meson.get_compiler('c')
 m_dep = cc.find_library('m', required: true)
 
+if get_option('buildtype').startswith('debug')
+  add_project_arguments('-DLPNG_DEBUG_LOG', language: ['c'])
+endif
+
 lisiblepng_incdir = include_directories('src/')
 lisiblepng_lib = library('lisiblepng', 'src/lisiblepng.c', 'src/lisiblepng/deflate.c', 'src/lisiblepng/bitstream.c', dependencies: [m_dep])
 lisiblepng_dep = declare_dependency(include_directories: lisiblepng_incdir, link_with: [lisiblepng_lib], dependencies: [m_dep])
diff --git a/lisiblepng/src/lisiblepng/bitstream.c b/lisiblepng/src/lisiblepng/bitstream.c
index e7ddb79..83e70c6 100644
--- a/lisiblepng/src/lisiblepng/bitstream.c
+++ b/lisiblepng/src/lisiblepng/bitstream.c
@@ -10,44 +10,40 @@ void Bitstream_init(Bitstream *bitstream, const uint8_t *data,
   ASSERT(data != NULL);
   bitstream->data = data;
   bitstream->data_size = data_size;
-  bitstream->current_byte_index = 0;
-  bitstream->current_bit_offset = 0;
+  bitstream->byte_offset = 0;
+  bitstream->bit_index = 0;
 }
 
 void bitstream_advance(Bitstream *bitstream, size_t bit_count) {
   ASSERT(bitstream != NULL);
-  uint8_t current_bit = bitstream->current_bit_offset;
-  bitstream->current_byte_index =
-      bitstream->current_byte_index + (current_bit + bit_count) / 8;
-  bitstream->current_bit_offset = (current_bit + bit_count) % 8;
+  uint8_t current_bit = bitstream->bit_index;
+  bitstream->byte_offset =
+      bitstream->byte_offset + (current_bit + bit_count) / 8;
+  bitstream->bit_index = (current_bit + bit_count) % 8;
 }
 
 void Bitstream_skip(Bitstream *bitstream, size_t bit_count) {
   ASSERT(bitstream != NULL);
   bitstream_advance(bitstream, bit_count);
 }
-
 uint16_t Bitstream_next_bits(Bitstream *bitstream, int bit_count) {
   ASSERT(bitstream != NULL);
   ASSERT(bit_count <= 16);
-  ASSERT((bitstream->current_bit_offset + (size_t)bit_count) / 8 <=
+  ASSERT((bitstream->bit_index + bit_count - 1) / 8 + bitstream->byte_offset <
          bitstream->data_size);
 
-  int bit_to_read = bit_count;
-  uint16_t val = bitstream->data[bitstream->current_byte_index] >>
-                 bitstream->current_bit_offset;
-
-  size_t bit_read = MIN(bit_count, 8 - bitstream->current_bit_offset);
-  bitstream_advance(bitstream, bit_read);
-  bit_to_read -= bit_read;
-  while (bit_to_read > 0) {
-    val |= (bitstream->data[bitstream->current_byte_index] >>
-            bitstream->current_bit_offset)
-           << bit_read;
-    bit_read = MIN(bit_to_read, 8 - bitstream->current_bit_offset);
-    bitstream_advance(bitstream, bit_read);
-    bit_to_read -= bit_read;
+  uint16_t val = 0;
+  int bits_collected = 0;
+  while (bits_collected < bit_count) {
+    int bits_in_current_byte = 8 - bitstream->bit_index;
+    int bits_to_extract = MIN(bit_count - bits_collected, bits_in_current_byte);
+    uint16_t extracted_bits =
+        (bitstream->data[bitstream->byte_offset] >> bitstream->bit_index) &
+        ((1 << bits_to_extract) - 1);
+    val |= extracted_bits << bits_collected;
+    bits_collected += bits_to_extract;
+    bitstream_advance(bitstream, bits_to_extract);
   }
 
-  return (val & ((1 << bit_count) - 1));
+  return val;
 }
diff --git a/lisiblepng/src/lisiblepng/bitstream.h b/lisiblepng/src/lisiblepng/bitstream.h
index 72d8295..2cb4fdc 100644
--- a/lisiblepng/src/lisiblepng/bitstream.h
+++ b/lisiblepng/src/lisiblepng/bitstream.h
@@ -8,8 +8,8 @@ typedef struct {
   const uint8_t *data;
   size_t data_size;
 
-  size_t current_byte_index;
-  uint8_t current_bit_offset;
+  size_t byte_offset;
+  uint8_t bit_index;
 } Bitstream;
 
 void Bitstream_init(Bitstream *bitstream, const uint8_t *data,
diff --git a/lisiblepng/src/lisiblepng/deflate.c b/lisiblepng/src/lisiblepng/deflate.c
index 5b33b1d..36073c6 100644
--- a/lisiblepng/src/lisiblepng/deflate.c
+++ b/lisiblepng/src/lisiblepng/deflate.c
@@ -71,7 +71,7 @@ void OutputBuffer_expand(OutputBuffer *output_buffer) {
   output_buffer->cap = new_cap;
 }
 
-void OutputBuffer_push(OutputBuffer *output_buffer, char byte) {
+void OutputBuffer_push(OutputBuffer *output_buffer, uint8_t byte) {
   ASSERT(output_buffer != NULL);
   if (output_buffer->len == output_buffer->cap) {
     OutputBuffer_expand(output_buffer);
@@ -196,9 +196,13 @@ bool deflate_decompress_(Bitstream *bitstream,
       return false;
     }
     if (symbol < 256) {
-      OutputBuffer_push(output, (char)symbol);
+      OutputBuffer_push(output, symbol);
     } else if (symbol > 256) {
       symbol -= 257;
+      if (symbol >= 29) {
+        return false;
+      }
+
       int length = length_size_base[symbol] +
                    Bitstream_next_bits(bitstream, length_extra_bits[symbol]);
       symbol = huffman_table_decode(distance_table->symbols,
@@ -207,14 +211,13 @@ bool deflate_decompress_(Bitstream *bitstream,
         return false;
       }
 
-      int distance =
-          distance_offset_base[symbol] +
+      int distance_increment =
           Bitstream_next_bits(bitstream, distance_extra_bits[symbol]);
-      while (length > 0) {
+      unsigned distance = distance_offset_base[symbol] + distance_increment;
+      while (length--) {
         size_t output_buffer_length = OutputBuffer_length(output);
-        OutputBuffer_push(
-            output, (char)output->buffer[output_buffer_length - distance]);
-        length--;
+        OutputBuffer_push(output,
+                          output->buffer[output_buffer_length - distance]);
       }
     }
   }
@@ -230,6 +233,7 @@ bool deflate_decompress(Bitstream *bitstream, OutputBuffer *output) {
   while (!b_final) {
     LPNG_LOG_DBG0("Parse deflate block");
     b_final = Bitstream_next_bits(bitstream, BFINAL_LENGTH_BITS);
+    LPNG_LOG_DBG("Final block: %d", b_final);
     const uint8_t b_type = Bitstream_next_bits(bitstream, BTYPE_LENGTH_BITS);
     if (b_type == DeflateBlockType_NoCompression) {
       LPNG_LOG_ERR0("Uncompressed deflate blocks aren't supported");
@@ -391,8 +395,8 @@ uint8_t *zlib_decompress(const uint8_t *compressed_data_buffer,
     return NULL;
   }
 
-  size_t adler32_offset = bitstream.current_byte_index;
-  if (bitstream.current_bit_offset % 8 != 0) {
+  size_t adler32_offset = bitstream.byte_offset;
+  if (bitstream.bit_index % 8 != 0) {
     adler32_offset++;
   }