fix: 修正 BitStream byte stuffing 和 RST marker 處理錯誤

  - 修正 0xFF 0x00 byte stuffing 未跳過 0x00 的嚴重 bug
  - 修正 RST marker (0xFF 0xD0-0xD7) 未正確跳過的問題
  - 重構 BitStream：引入 32-bit 緩衝區機制提升效能
  - 新增 precision 驗證和清理未使用變數

Author: 5000user5000

src/cpp/bitstream.cpp

@@ -5,114 +5,117 @@
 namespace jpeg {
 
 BitStream::BitStream(const uint8_t* data, size_t size)
-    : data_(data), size_(size), byte_pos_(0), bit_pos_(0) {
+    : data_(data), size_(size), byte_pos_(0), bit_buffer_(0), bits_in_buffer_(0) {
     if (!data || size == 0) {
         throw std::invalid_argument("BitStream: invalid data or size");
     }
 }
 
-uint16_t BitStream::readBits(int n_bits) {
-    if (n_bits <= 0 || n_bits > 16) {
-        throw std::invalid_argument("BitStream: n_bits must be between 1 and 16");
-    }
-    
-    uint16_t result = 0;
-    
-    for (int i = 0; i < n_bits; ++i) {
-        // update byte and bit positions
-        if (bit_pos_ == 8) {
-            byte_pos_++;
-            bit_pos_ = 0;
-        }
-        
-        // 檢查是否超出範圍
-        if (byte_pos_ >= size_) {
-            throw std::runtime_error("BitStream: unexpected end of data");
-        }
-        
-        // 處理 byte stuffing: 如果當前 byte 是 0xFF
-        uint8_t current_byte = data_[byte_pos_];
-        if (current_byte == 0xFF && bit_pos_ == 0) {
-            // 檢查下一個 byte
-            if (byte_pos_ + 1 < size_) {
-                uint8_t next_byte = data_[byte_pos_ + 1];
+void BitStream::fillBuffer() {
+    while (bits_in_buffer_ <= 24 && byte_pos_ < size_) {
+        uint8_t byte = data_[byte_pos_++];
+
+        // 處理 byte stuffing: 如果是 0xFF，檢查下一個 byte
+        if (byte == 0xFF) {
+            if (byte_pos_ < size_) {
+                uint8_t next_byte = data_[byte_pos_];
                 if (next_byte == 0x00) {
-                    // 這是 byte stuffing，跳過 0x00
-                    // 繼續使用 0xFF
+                    // 是 byte stuffing (FF 00)，跳過 00
+                    byte_pos_++;
                 } else if (next_byte >= 0xD0 && next_byte <= 0xD7) {
-                    // RST marker，這裡簡單處理，繼續讀取
-                    // 更完整的實作應該處理 restart interval
+                    // 是 RST marker (FF Dx)，跳過整個 marker
+                    byte_pos_++;
+                    continue;  // 不要將 0xFF 放入緩衝區，繼續讀取下一個 byte
                 } else {
-                    // 其他 marker，可能是資料結束
-                    // 目前先簡單拋出異常
-                    throw std::runtime_error("BitStream: unexpected marker");
+                    // 其他 marker，可能是錯誤或資料結束
+                    // 根據 JPEG 規範，這裡應該是資料結束，但我們保守地放入 0xFF
                 }
             }
         }
-        
-        // 讀取一個 bit (MSB first)
-        int bit = (current_byte >> (7 - bit_pos_)) & 1;
-        result = (result << 1) | bit;
-        bit_pos_++;
+
+        bit_buffer_ = (bit_buffer_ << 8) | byte;
+        bits_in_buffer_ += 8;
     }
-    
-    return result;
 }
 
-uint16_t BitStream::peekBits(int n_bits) {
-    // 保存當前狀態
-    size_t saved_byte_pos = byte_pos_;
-    int saved_bit_pos = bit_pos_;
+uint16_t BitStream::readBits(int n_bits) {
+    if (n_bits <= 0 || n_bits > 16) {
+        throw std::invalid_argument("BitStream: n_bits must be between 1 and 16");
+    }
+    
+    if (bits_in_buffer_ < n_bits) {
+        fillBuffer();
+    }
+    
+    if (bits_in_buffer_ < n_bits) {
+        throw std::runtime_error("BitStream: unexpected end of data");
+    }
 
-    uint16_t result = readBits(n_bits);
+    uint16_t result = bit_buffer_ >> (bits_in_buffer_ - n_bits);
 
-    // 恢復狀態
-    byte_pos_ = saved_byte_pos;
-    bit_pos_ = saved_bit_pos;
+    bits_in_buffer_ -= n_bits;
+    bit_buffer_ &= (1U << bits_in_buffer_) - 1;
 
     return result;
 }
 
+uint16_t BitStream::peekBits(int n_bits) {
+    if (n_bits <= 0 || n_bits > 16) {
+        throw std::invalid_argument("BitStream: n_bits must be between 1 and 16");
+    }
+
+    if (bits_in_buffer_ < n_bits) {
+        fillBuffer();
+    }
+
+    if (bits_in_buffer_ < n_bits) {
+        throw std::runtime_error("BitStream: unexpected end of data");
+    }
+
+    return bit_buffer_ >> (bits_in_buffer_ - n_bits);
+}
+
 void BitStream::skipBits(int n_bits) {
-    readBits(n_bits);
+    if (bits_in_buffer_ >= n_bits) {
+        bits_in_buffer_ -= n_bits;
+        bit_buffer_ &= (1U << bits_in_buffer_) - 1;
+    } else {
+        n_bits -= bits_in_buffer_;
+        bits_in_buffer_ = 0;
+        bit_buffer_ = 0;
+        
+        // 這裡可以優化，但為了簡單起見，先這樣
+        size_t bytes_to_skip = n_bits / 8;
+        byte_pos_ += bytes_to_skip;
+        n_bits %= 8;
+        
+        if (n_bits > 0) {
+            readBits(n_bits);
+        }
+    }
 }
 
 bool BitStream::hasMoreData() const {
-    return byte_pos_ < size_;
+    return bits_in_buffer_ > 0 || byte_pos_ < size_;
 }
 
 size_t BitStream::getBitPosition() const {
-    return byte_pos_ * 8 + bit_pos_;
+    return byte_pos_ * 8 - bits_in_buffer_;
 }
 
 void BitStream::reset(size_t byte_pos, int bit_pos) {
-    if (byte_pos >= size_) {
+    if (byte_pos > size_) { // allow reset to end
         throw std::out_of_range("BitStream: byte_pos out of range");
     }
     if (bit_pos < 0 || bit_pos >= 8) {
         throw std::out_of_range("BitStream: bit_pos must be 0-7");
     }
     byte_pos_ = byte_pos;
-    bit_pos_ = bit_pos;
-}
-
-uint8_t BitStream::getNextByte() {
-    if (byte_pos_ >= size_) {
-        throw std::runtime_error("BitStream: unexpected end of data");
-    }
-    
-    uint8_t byte = data_[byte_pos_++];
-    
-    // 處理 byte stuffing
-    if (byte == 0xFF && byte_pos_ < size_) {
-        uint8_t next = data_[byte_pos_];
-        if (next == 0x00) {
-            // Byte stuffing: 跳過 0x00
-            byte_pos_++;
-        }
+    bits_in_buffer_ = 0;
+    bit_buffer_ = 0;
+    if (bit_pos > 0) {
+        readBits(bit_pos); // 讀取並丟棄，以對齊 bit 位置
     }
-    
-    return byte;
 }
 
 } // namespace jpeg

src/cpp/bitstream.h

@@ -60,13 +60,14 @@ class BitStream {
     const uint8_t* data_;      // 原始資料指標
     size_t size_;              // 資料大小
     size_t byte_pos_;          // 當前位元組位置
-    int bit_pos_;              // 當前 bit 位置 (0-7, 0 = MSB)
+
+    uint32_t bit_buffer_;      // 位元緩衝區
+    int bits_in_buffer_;       // 緩衝區中剩餘的 bit 數量
 
     /**
-     * 讀取下一個 byte，處理 byte stuffing
-     * @return 下一個有效的 byte
+     * 填充緩衝區，處理 byte stuffing
      */
-    uint8_t getNextByte();
+    void fillBuffer();
 };
 
 } // namespace jpeg

src/cpp/decoder.cpp

@@ -153,9 +153,13 @@ bool JPEGDecoder::processDQT() {
 }
 
 bool JPEGDecoder::processSOF0() {
-    uint16_t length = readWord();
-    
-    uint8_t precision = readByte();  // 通常是 8
+    readWord(); // length
+
+    uint8_t precision = readByte();  // precision (通常是 8)
+    if (precision != 8) {
+        throw std::runtime_error("Only 8-bit precision is supported");
+    }
+
     height_ = readWord();
     width_ = readWord();
     num_components_ = readByte();
@@ -206,13 +210,13 @@ bool JPEGDecoder::processDHT() {
 }
 
 bool JPEGDecoder::processDRI() {
-    uint16_t length = readWord();
+    readWord(); // length
     restart_interval_ = readWord();
     return true;
 }
 
 bool JPEGDecoder::processSOS() {
-    uint16_t length = readWord();
+    readWord(); // length
 
     int num_components = readByte();
 
@@ -412,4 +416,4 @@ void JPEGDecoder::upsample(const std::vector<std::vector<uint8_t>>& y_blocks,
     }
 }
 
-} // namespace jpeg
+} // namespace jpeg