Fixed handling of the stream end in Huffmann decompression

2 files changed, 32 insertions, 29 deletions

diff --git a/src/huffman/huff.cpp b/src/huffman/huff.cpp
index 6930354..1b81017 100644
--- a/src/huffman/huff.cpp
+++ b/src/huffman/huff.cpp
@@ -269,24 +269,6 @@ TInputStream::TInputStream(void * pvInBuffer, size_t cbInBuffer)
     BitCount = 0;
 }
 
-// Gets 7 bits from the stream. DOES NOT remove the bits from input stream
-unsigned int TInputStream::Peek7Bits()
-{
-    unsigned int dwReloadByte = 0;
-
-    // If there is not enough bits to get the value,
-    // we have to add 8 more bits from the input buffer
-    if(BitCount < 7)
-    {
-        dwReloadByte = *pbInBuffer++;
-        BitBuffer |= dwReloadByte << BitCount;
-        BitCount += 8;
-    }
-
-    // Return the first available 7 bits. DO NOT remove them from the input stream
-    return (BitBuffer & 0x7F);
-}
-
 // Gets one bit from input stream
 unsigned int TInputStream::Get1Bit()
 {
@@ -330,6 +312,30 @@ unsigned int TInputStream::Get8Bits()
     return dwOneByte;
 }
 
+// Gets 7 bits from the stream. DOES NOT remove the bits from input stream
+bool TInputStream::Peek7Bits(unsigned int & Value)
+{
+    unsigned int Value8Bits = 0;
+
+    // If there is not enough bits to get the value,
+    // we have to add 8 more bits from the input buffer
+    if(BitCount < 7)
+    {
+        // Load additional 8 bits. Be careful if we have no more data
+        if(pbInBuffer >= pbInBufferEnd)
+            return false;
+        Value8Bits = *pbInBuffer++;
+
+        // Add these 8 bits to the bit buffer
+        BitBuffer |= Value8Bits << BitCount;
+        BitCount += 8;
+    }
+
+    // Return 7 bits of data. DO NOT remove them from the input stream
+    Value = (BitBuffer & 0x7F);
+    return true;
+}
+
 void TInputStream::SkipBits(unsigned int dwBitsToSkip)
 {
     unsigned int dwReloadByte = 0;
@@ -696,16 +702,13 @@ unsigned int THuffmannTree::DecodeOneByte(TInputStream * is)
     THTreeItem * pItem;
     unsigned int ItemLinkIndex;
     unsigned int BitCount = 0;
+    bool bHasItemLinkIndex;
 
-    // Check for the end of the input stream
-    if(is->pbInBuffer >= is->pbInBufferEnd && is->BitCount < 7)
-        return 0x1FF;
-
-    // Get the eventual quick-link index
-    ItemLinkIndex = is->Peek7Bits();
+    // Try to retrieve quick link index
+    bHasItemLinkIndex = is->Peek7Bits(ItemLinkIndex);
 
     // Is the quick-link item valid?
-    if(QuickLinks[ItemLinkIndex].ValidValue > MinValidValue)
+    if(bHasItemLinkIndex && QuickLinks[ItemLinkIndex].ValidValue > MinValidValue)
     {
         // If that item needs less than 7 bits, we can get decompressed value directly
         if(QuickLinks[ItemLinkIndex].ValidBits <= 7)
@@ -745,7 +748,7 @@ unsigned int THuffmannTree::DecodeOneByte(TInputStream * is)
 
     // If we didn't get the item from the quick-link array,
     // set the entry in it
-    if(QuickLinks[ItemLinkIndex].ValidValue < MinValidValue)
+    if(bHasItemLinkIndex && QuickLinks[ItemLinkIndex].ValidValue < MinValidValue)
     {
         // If the current compressed byte was more than 7 bits,
         // set a quick-link item with pointer to tree item
@@ -876,10 +879,10 @@ unsigned int THuffmannTree::Decompress(void * pvOutBuffer, unsigned int cbOutLen
                 IncWeightsAndRebalance(ItemsByByte[DecompressedValue]);
         }
 
-        // A byte successfully decoded - store it in the output stream
-        *pbOutBuffer++ = (unsigned char)DecompressedValue;
+        // Store the byte to the output stream
         if(pbOutBuffer >= pbOutBufferEnd)
             break;
+        *pbOutBuffer++ = (unsigned char)DecompressedValue;
 
         if(bIsCmp0)
         {
diff --git a/src/huffman/huff.h b/src/huffman/huff.h
index b75acbb..cf1ca4c 100644
--- a/src/huffman/huff.h
+++ b/src/huffman/huff.h
@@ -29,8 +29,8 @@ class TInputStream
 
     TInputStream(void * pvInBuffer, size_t cbInBuffer);
     unsigned int Get1Bit();
-    unsigned int Peek7Bits();
     unsigned int Get8Bits();
+    bool Peek7Bits(unsigned int & Value);
     void SkipBits(unsigned int BitCount);
 
     unsigned char * pbInBufferEnd;      // End position in the the input buffer