Skip to main content

zune_jpeg/
bitstream.rs

1/*
2 * Copyright (c) 2023.
3 *
4 * This software is free software;
5 *
6 * You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
7 */
8
9#![allow(
10    clippy::if_not_else,
11    clippy::similar_names,
12    clippy::inline_always,
13    clippy::doc_markdown,
14    clippy::cast_sign_loss,
15    clippy::cast_possible_truncation
16)]
17
18//! This file exposes a single struct that can decode a huffman encoded
19//! Bitstream in a JPEG file
20//!
21//! This code is optimized for speed.
22//! It's meant to be super duper super fast, because everyone else depends on this being fast.
23//! It's (annoyingly) serial hence we cant use parallel bitstreams(it's variable length coding.)
24//!
25//! Furthermore, on the case of refills, we have to do bytewise processing because the standard decided
26//! that we want to support markers in the middle of streams(seriously few people use RST markers).
27//!
28//! So we pull in all optimization steps:
29//! - use `inline[always]`? ✅ ,
30//! - pre-execute most common cases ✅,
31//! - add random comments ✅
32//! -  fast paths ✅.
33//!
34//! Speed-wise: It is probably the fastest JPEG BitStream decoder to ever sail the seven seas because of
35//! a couple of optimization tricks.
36//! 1. Fast refills from libjpeg-turbo
37//! 2. As few as possible branches in decoder fast paths.
38//! 3. Accelerated AC table decoding borrowed from stb_image.h written by Fabian Gissen (@ rygorous),
39//! improved by me to handle more cases.
40//! 4. Safe and extensible routines(e.g. cool ways to eliminate bounds check)
41//! 5. No unsafe here
42//!
43//! Readability comes as a second priority(I tried with variable names this time, and we are wayy better than libjpeg).
44//!
45//! Anyway if you are reading this it means your cool and I hope you get whatever part of the code you are looking for
46//! (or learn something cool)
47//!
48//! Knock yourself out.
49use alloc::format;
50use alloc::string::ToString;
51use zune_core::log::warn;
52use core::cmp::min;
53
54use zune_core::bytestream::{ZByteReaderTrait, ZReader};
55
56use crate::errors::DecodeErrors;
57use crate::huffman::{HuffmanTable, HUFF_LOOKAHEAD};
58use crate::marker::Marker;
59use crate::mcu::DCT_BLOCK;
60use crate::misc::UN_ZIGZAG;
61
62macro_rules! decode_huff {
63    ($stream:tt,$symbol:tt,$table:tt) => {
64        let mut code_length = $symbol >> HUFF_LOOKAHEAD;
65
66        ($symbol) &= (1 << HUFF_LOOKAHEAD) - 1;
67
68        if code_length > i32::from(HUFF_LOOKAHEAD)
69        {
70            // if the symbol cannot be resolved in the first HUFF_LOOKAHEAD bits,
71            // we know it lies somewhere between HUFF_LOOKAHEAD and 16 bits since jpeg imposes 16 bit
72            // limit, we can therefore look 16 bits ahead and try to resolve the symbol
73            // starting from 1+HUFF_LOOKAHEAD bits.
74            $symbol = ($stream).peek_bits::<16>() as i32;
75            // (Credits to Sean T. Barrett stb library for this optimization)
76            // maxcode is pre-shifted 16 bytes long so that it has (16-code_length)
77            // zeroes at the end hence we do not need to shift in the inner loop.
78            while code_length < 17{
79                if $symbol < $table.maxcode[code_length as usize]  {
80                    break;
81                }
82                code_length += 1;
83            }
84
85            if code_length == 17{
86                // symbol could not be decoded.
87                //
88                // We may think, lets fake zeroes, noo
89                // panic, because Huffman codes are sensitive, probably everything
90                // after this will be corrupt, so no need to continue.
91                // panic!("Bad Huffman code length");
92                return Err(DecodeErrors::Format(format!("Bad Huffman Code 0x{:X}, corrupt JPEG",$symbol)))
93            }
94
95            $symbol >>= (16-code_length);
96            ($symbol) = i32::from(
97                ($table).values
98                    [(($symbol + ($table).offset[code_length as usize]) & 0xFF) as usize],
99            );
100        }
101        if code_length> i32::from(($stream).bits_left){
102            return Err(DecodeErrors::Format(format!("Code length {code_length}  more than bits left {}",($stream).bits_left)))
103        }
104        // drop bits read
105        ($stream).drop_bits(code_length as u8);
106    };
107}
108
109/// A `BitStream` struct, a bit by bit reader with super powers
110///
111#[rustfmt::skip]
112pub(crate) struct BitStream {
113    /// A MSB type buffer that is used for some certain operations
114    pub buffer:              u64,
115    /// A TOP  aligned MSB type buffer that is used to accelerate some operations like
116    /// peek_bits and get_bits.
117    ///
118    /// By top aligned, I mean the top bit (63) represents the top bit in the buffer.
119    aligned_buffer:          u64,
120    /// Tell us the bits left the two buffer
121    pub(crate) bits_left:    u8,
122    /// Did we find a marker(RST/EOF) during decoding?
123    pub marker:              Option<Marker>,
124    /// An i16 with the bit corresponding to successive_low set to 1, others 0.
125    pub successive_low_mask: i16,
126    spec_start:              u8,
127    spec_end:                u8,
128    pub eob_run:             i32,
129    pub overread_by:         usize,
130    /// True if we have seen end of image marker.
131    /// Don't read anything after that.
132    pub seen_eoi:            bool,
133}
134
135impl BitStream {
136    /// Create a new BitStream
137    #[rustfmt::skip]
138    pub(crate) const fn new() -> BitStream {
139        BitStream {
140            buffer:              0,
141            aligned_buffer:      0,
142            bits_left:           0,
143            marker:              None,
144            successive_low_mask: 1,
145            spec_start:          0,
146            spec_end:            0,
147            eob_run:             0,
148            overread_by:         0,
149            seen_eoi:            false,
150        }
151    }
152
153    /// Create a new Bitstream for progressive decoding
154    #[allow(clippy::redundant_field_names)]
155    #[rustfmt::skip]
156    pub(crate) fn new_progressive(al: u8, spec_start: u8, spec_end: u8) -> BitStream {
157        BitStream {
158            buffer:              0,
159            aligned_buffer:      0,
160            bits_left:           0,
161            marker:              None,
162            successive_low_mask: 1i16 << al,
163            spec_start:          spec_start,
164            spec_end:            spec_end,
165            eob_run:             0,
166            overread_by:         0,
167            seen_eoi:            false,
168        }
169    }
170
171    /// Refill the bit buffer by (a maximum of) 32 bits
172    ///
173    /// # Arguments
174    ///  - `reader`:`&mut BufReader<R>`: A mutable reference to an underlying
175    ///    File/Memory buffer containing a valid JPEG stream
176    ///
177    /// This function will only refill if `self.count` is less than 32
178    #[inline(always)] // to many call sites? ( perf improvement by 4%)
179    pub fn refill<T>(&mut self, reader: &mut ZReader<T>) -> Result<bool, DecodeErrors>
180    where
181        T: ZByteReaderTrait
182    {
183        /// Macro version of a single byte refill.
184        /// Arguments
185        /// buffer-> our io buffer, because rust macros cannot get values from
186        /// the surrounding environment bits_left-> number of bits left
187        /// to full refill
188        macro_rules! refill {
189            ($buffer:expr,$byte:expr,$bits_left:expr) => {
190                // read a byte from the stream
191                $byte = u64::from(reader.read_u8());
192                self.overread_by += usize::from(reader.eof()?);
193                // append to the buffer
194                // JPEG is a MSB type buffer so that means we append this
195                // to the lower end (0..8) of the buffer and push the rest bits above..
196                $buffer = ($buffer << 8) | $byte;
197                // Increment bits left
198                $bits_left += 8;
199                // Check for special case  of OxFF, to see if it's a stream or a marker
200                if $byte == 0xff {
201                    // read next byte
202                    let mut next_byte = u64::from(reader.read_u8());
203                    // Byte snuffing, if we encounter byte snuff, we skip the byte
204                    if next_byte != 0x00 {
205                        // skip that byte we read
206                        while next_byte == 0xFF {
207                            next_byte = u64::from(reader.read_u8());
208                        }
209
210                        if next_byte != 0x00 {
211                            // Undo the byte append and return
212                            $buffer >>= 8;
213                            $bits_left -= 8;
214
215                            if $bits_left != 0 {
216                                self.aligned_buffer = $buffer << (64 - $bits_left);
217                            }
218
219                            let marker = Marker::from_u8(next_byte as u8);
220                            self.marker = marker;
221
222                            if let Some(Marker::UNKNOWN(_)) = marker{
223                                return Err(DecodeErrors::Format("Unknown marker in bit stream".to_string()));
224                            }
225                            if next_byte == 0xD9 {
226                                // special handling for eoi, fill some bytes,even if its zero,
227                                // removes some panics
228                                self.buffer <<= 8;
229                                self.bits_left += 8;
230                                self.aligned_buffer = self.buffer << (64 - self.bits_left);
231                            }
232
233                            return Ok(false);
234                        }
235                    }
236                }
237            };
238        }
239
240        // 32 bits is enough for a decode(16 bits) and receive_extend(max 16 bits)
241        if self.bits_left < 32 {
242            if self.marker.is_some() || self.seen_eoi {
243                // found a marker, or we are in EOI
244                // also we are in over-reading mode, where we fill it with zeroes
245
246                // fill with zeroes
247                self.buffer <<= 32;
248                self.bits_left += 32;
249                self.aligned_buffer = self.buffer << (64 - self.bits_left);
250                return Ok(true);
251            }
252
253            if self.overread_by > 0 {
254                if self.bits_left == 0 {
255                    return Err(DecodeErrors::ExhaustedData);
256                }
257                // We already hit EOF while refilling. Continue consuming the buffered bits
258                // but don't synthesize additional bytes from zero-fill.
259                return Ok(true);
260            }
261
262            // we optimize for the case where we don't have 255 in the stream and have 4 bytes left
263            // as it is the common case
264            //
265            // so we always read 4 bytes, if read_fixed_bytes errors out, the cursor is
266            // guaranteed not to advance in case of failure (is this true), so
267            // we revert the read later on (if we have 255), if this fails, we use the normal
268            // byte at a time read
269
270            if let Ok(bytes) = reader.read_fixed_bytes_or_error::<4>() {
271                // we have 4 bytes to spare, read the 4 bytes into a temporary buffer
272                // create buffer
273                let msb_buf = u32::from_be_bytes(bytes);
274                // check if we have 0xff
275                if !has_byte(msb_buf, 255) {
276                    self.bits_left += 32;
277                    self.buffer <<= 32;
278                    self.buffer |= u64::from(msb_buf);
279                    self.aligned_buffer = self.buffer << (64 - self.bits_left);
280                    return Ok(true);
281                }
282
283                reader.rewind(4)?;
284            }
285            // This serves two reasons,
286            // 1: Make clippy shut up
287            // 2: Favour register reuse
288            let mut byte;
289            // 4 refills, if all succeed the stream should contain enough bits to decode a
290            // value
291            refill!(self.buffer, byte, self.bits_left);
292            refill!(self.buffer, byte, self.bits_left);
293            refill!(self.buffer, byte, self.bits_left);
294            refill!(self.buffer, byte, self.bits_left);
295            // Construct an MSB buffer whose top bits are the bitstream we are currently holding.
296            self.aligned_buffer = self.buffer << (64 - self.bits_left);
297        }
298        return Ok(true);
299    }
300    /// Decode the DC coefficient in a MCU block.
301    ///
302    /// The decoded coefficient is written to `dc_prediction`
303    ///
304    #[allow(
305        clippy::cast_possible_truncation,
306        clippy::cast_sign_loss,
307        clippy::unwrap_used
308    )]
309    #[inline(always)]
310    fn decode_dc<T>(
311        &mut self, reader: &mut ZReader<T>, dc_table: &HuffmanTable, dc_prediction: &mut i32
312    ) -> Result<bool, DecodeErrors>
313    where
314        T: ZByteReaderTrait
315    {
316        let (mut symbol, r);
317
318        if self.bits_left < 32 {
319            self.refill(reader)?;
320        };
321        // look a head HUFF_LOOKAHEAD bits into the bitstream
322        symbol = self.peek_bits::<HUFF_LOOKAHEAD>();
323        symbol = dc_table.lookup[symbol as usize];
324
325        decode_huff!(self, symbol, dc_table);
326
327        if symbol != 0 {
328            r = self.get_bits(symbol as u8);
329            symbol = huff_extend(r, symbol);
330        }
331        // Update DC prediction
332        *dc_prediction = dc_prediction.wrapping_add(symbol);
333
334        return Ok(true);
335    }
336
337    /// Like `decode_dc` but we do not need the result of the component, we only want to remove it
338    /// from the bitstream of the MCU.
339    fn discard_dc<T>(
340        &mut self, reader: &mut ZReader<T>, dc_table: &HuffmanTable
341    ) -> Result<bool, DecodeErrors>
342    where
343        T: ZByteReaderTrait
344    {
345        let mut symbol;
346
347        if self.bits_left < 32 {
348            self.refill(reader)?;
349        };
350        // look a head HUFF_LOOKAHEAD bits into the bitstream
351        symbol = self.peek_bits::<HUFF_LOOKAHEAD>();
352        symbol = dc_table.lookup[symbol as usize];
353
354        decode_huff!(self, symbol, dc_table);
355
356        if symbol != 0 {
357            let _ = self.get_bits(symbol as u8);
358        }
359
360        return Ok(true);
361    }
362
363    /// Decode a Minimum Code Unit(MCU) as quickly as possible
364    ///
365    /// # Arguments
366    /// - reader: The bitstream from where we read more bits.
367    /// - dc_table: The Huffman table used to decode the DC coefficient
368    /// - ac_table: The Huffman table used to decode AC values
369    /// - block: A memory region where we will write out the decoded values
370    /// - DC prediction: Last DC value for this component
371    ///
372    #[allow(
373        clippy::many_single_char_names,
374        clippy::cast_possible_truncation,
375        clippy::cast_sign_loss
376    )]
377    #[inline(never)]
378    pub fn decode_mcu_block<T>(
379        &mut self, reader: &mut ZReader<T>, dc_table: &HuffmanTable, ac_table: &HuffmanTable,
380        qt_table: &[i32; DCT_BLOCK], block: &mut [i32; 64], dc_prediction: &mut i32
381    ) -> Result<u16, DecodeErrors>
382    where
383        T: ZByteReaderTrait
384    {
385        // Get fast AC table as a reference before we enter the hot path
386        let ac_lookup = ac_table.ac_lookup.as_ref().unwrap();
387
388        let (mut symbol, mut r, mut fast_ac);
389        // Decode AC coefficients
390        let mut pos: usize = 1;
391        if  self.bits_left < 1 && self.marker.is_some() {
392            return Err(DecodeErrors::Format(
393                "No more bytes left in stream before marker".to_string()
394            ));
395        }
396        // decode DC, dc prediction will contain the value
397        self.decode_dc(reader, dc_table, dc_prediction)?;
398
399        // set dc to be the dc prediction.
400        block[0] = *dc_prediction * qt_table[0];
401
402        while pos < 64 {
403            self.refill(reader)?;
404            symbol = self.peek_bits::<HUFF_LOOKAHEAD>();
405            fast_ac = ac_lookup[symbol as usize];
406            symbol = ac_table.lookup[symbol as usize];
407
408            if fast_ac != 0 {
409                //  FAST AC path
410                pos += ((fast_ac >> 4) & 15) as usize; // run
411                let t_pos = UN_ZIGZAG[min(pos, 63)] & 63;
412
413                block[t_pos] = i32::from(fast_ac >> 8) * (qt_table[t_pos]); // Value
414                self.drop_bits((fast_ac & 15) as u8);
415                pos += 1;
416            } else {
417                decode_huff!(self, symbol, ac_table);
418
419                r = symbol >> 4;
420                symbol &= 15;
421
422                if symbol != 0 {
423                    pos += r as usize;
424                    r = self.get_bits(symbol as u8);
425                    symbol = huff_extend(r, symbol);
426                    let t_pos = UN_ZIGZAG[pos & 63] & 63;
427
428                    block[t_pos] = symbol * qt_table[t_pos];
429
430                    pos += 1;
431                } else if r != 15 {
432                    return Ok(pos as u16);
433                } else {
434                    pos += 16;
435                }
436            }
437        }
438
439        return Ok(64);
440    }
441
442    /// Advance the bitstream over a block but ignore the data contained.
443    ///
444    /// This updates DC prediction but we never dequantize and we never do any Zig-Zag translation
445    /// either. Still returns the index of the last component read.
446    pub fn discard_mcu_block<T>(
447        &mut self, reader: &mut ZReader<T>, dc_table: &HuffmanTable, ac_table: &HuffmanTable
448    ) -> Result<u16, DecodeErrors>
449    where
450        T: ZByteReaderTrait
451    {
452        // Get fast AC table as a reference before we enter the hot path
453        let ac_lookup = ac_table.ac_lookup.as_ref().unwrap();
454
455        let (mut symbol, mut r, mut fast_ac);
456        // Decode AC coefficients
457        let mut pos: usize = 1;
458
459        // decode DC, dc prediction will contain the value
460        self.discard_dc(reader, dc_table)?;
461
462        while pos < 64 {
463            self.refill(reader)?;
464            symbol = self.peek_bits::<HUFF_LOOKAHEAD>();
465            fast_ac = ac_lookup[symbol as usize];
466            symbol = ac_table.lookup[symbol as usize];
467
468            if fast_ac != 0 {
469                //  FAST AC path
470                pos += ((fast_ac >> 4) & 15) as usize; // run
471
472                self.drop_bits((fast_ac & 15) as u8);
473                pos += 1;
474            } else {
475                decode_huff!(self, symbol, ac_table);
476
477                r = symbol >> 4;
478                symbol &= 15;
479
480                if symbol != 0 {
481                    pos += r as usize;
482                    // Advance over bits but ignore.
483                    let _ = self.get_bits(symbol as u8);
484
485                    pos += 1;
486                } else if r != 15 {
487                    return Ok(pos as u16);
488                } else {
489                    pos += 16;
490                }
491            }
492        }
493
494        return Ok(64);
495    }
496
497    /// Peek `look_ahead` bits ahead without discarding them from the buffer
498    #[inline(always)]
499    #[allow(clippy::cast_possible_truncation)]
500    const fn peek_bits<const LOOKAHEAD: u8>(&self) -> i32 {
501        (self.aligned_buffer >> (64 - LOOKAHEAD)) as i32
502    }
503
504    /// Discard the next `N` bits without checking
505    #[inline]
506    fn drop_bits(&mut self, n: u8) {
507        // PS: Its a good check, but triggers fuzzer and a lot of false positives
508        //debug_assert!(self.bits_left >= n);
509        //self.bits_left -= n;
510        self.bits_left = self.bits_left.saturating_sub(n);
511        self.aligned_buffer <<= n;
512    }
513
514    /// Read `n_bits` from the buffer  and discard them
515    #[inline(always)]
516    #[allow(clippy::cast_possible_truncation)]
517    fn get_bits(&mut self, n_bits: u8) -> i32 {
518        let mask = (1_u64 << n_bits) - 1;
519
520        self.aligned_buffer = self.aligned_buffer.rotate_left(u32::from(n_bits));
521        let bits = (self.aligned_buffer & mask) as i32;
522        self.bits_left = self.bits_left.wrapping_sub(n_bits);
523        bits
524    }
525
526    /// Decode a DC block
527    #[allow(clippy::cast_possible_truncation)]
528    #[inline]
529    pub(crate) fn decode_prog_dc_first<T>(
530        &mut self, reader: &mut ZReader<T>, dc_table: &HuffmanTable, block: &mut i16,
531        dc_prediction: &mut i32
532    ) -> Result<(), DecodeErrors>
533    where
534        T: ZByteReaderTrait
535    {
536        self.decode_dc(reader, dc_table, dc_prediction)?;
537        *block = (*dc_prediction as i16).wrapping_mul(self.successive_low_mask);
538        return Ok(());
539    }
540    #[inline]
541    pub(crate) fn decode_prog_dc_refine<T>(
542        &mut self, reader: &mut ZReader<T>, block: &mut i16
543    ) -> Result<(), DecodeErrors>
544    where
545        T: ZByteReaderTrait
546    {
547        // refinement scan
548        if self.bits_left < 1 {
549            self.refill(reader)?;
550            // if we find a marker, it may happens we don't refill.
551            // So let's confirm again that refill worked
552            if self.bits_left < 1 {
553                return Err(DecodeErrors::Format(
554                    "Marker found where not expected in refine bit".to_string()
555                ));
556            }
557        }
558
559        if self.get_bit() == 1 {
560            *block = block.wrapping_add(self.successive_low_mask);
561        }
562
563        Ok(())
564    }
565
566    /// Get a single bit from the bitstream
567    fn get_bit(&mut self) -> u8 {
568        let k = (self.aligned_buffer >> 63) as u8;
569        // discard a bit
570        self.drop_bits(1);
571        return k;
572    }
573    pub(crate) fn decode_mcu_ac_first<T>(
574        &mut self, reader: &mut ZReader<T>, ac_table: &HuffmanTable, block: &mut [i16; 64]
575    ) -> Result<bool, DecodeErrors>
576    where
577        T: ZByteReaderTrait
578    {
579        let fast_ac = ac_table.ac_lookup.as_ref().unwrap();
580        let bit = self.successive_low_mask;
581
582        let mut k = self.spec_start as usize;
583        let (mut symbol, mut r, mut fac);
584
585        // EOB runs are handled in mcu_prog.rs
586        'block: loop {
587            self.refill(reader)?;
588            // Check for marker in the stream
589
590            symbol = self.peek_bits::<HUFF_LOOKAHEAD>();
591            fac = fast_ac[symbol as usize];
592            symbol = ac_table.lookup[symbol as usize];
593
594            if fac != 0 {
595                // fast ac path
596                k += ((fac >> 4) & 15) as usize; // run
597                block[UN_ZIGZAG[min(k, 63)] & 63] = (fac >> 8).wrapping_mul(bit); // value
598                self.drop_bits((fac & 15) as u8);
599                k += 1;
600            } else {
601                decode_huff!(self, symbol, ac_table);
602
603                r = symbol >> 4;
604                symbol &= 15;
605
606                if symbol != 0 {
607                    k += r as usize;
608                    r = self.get_bits(symbol as u8);
609                    symbol = huff_extend(r, symbol);
610                    block[UN_ZIGZAG[k & 63] & 63] = (symbol as i16).wrapping_mul(bit);
611                    k += 1;
612                } else {
613                    if r != 15 {
614                        self.eob_run = 1 << r;
615                        self.eob_run += self.get_bits(r as u8);
616                        self.eob_run -= 1;
617                        break;
618                    }
619
620                    k += 16;
621                }
622            }
623
624            if k > self.spec_end as usize {
625                break 'block;
626            }
627        }
628        return Ok(true);
629    }
630    #[allow(clippy::too_many_lines, clippy::op_ref)]
631    pub(crate) fn decode_mcu_ac_refine<T>(
632        &mut self, reader: &mut ZReader<T>, table: &HuffmanTable, block: &mut [i16; 64]
633    ) -> Result<bool, DecodeErrors>
634    where
635        T: ZByteReaderTrait
636    {
637        let bit = self.successive_low_mask;
638
639        let mut k = self.spec_start;
640        let (mut symbol, mut r);
641
642        if self.eob_run == 0 {
643            'no_eob: loop {
644                // Decode a coefficient from the bit stream
645                self.refill(reader)?;
646
647                symbol = self.peek_bits::<HUFF_LOOKAHEAD>();
648                symbol = table.lookup[symbol as usize];
649
650                decode_huff!(self, symbol, table);
651
652                r = symbol >> 4;
653                symbol &= 15;
654
655                if symbol == 0 {
656                    if r != 15 {
657                        // EOB run is 2^r + bits
658                        self.eob_run = 1 << r;
659                        self.eob_run += self.get_bits(r as u8);
660                        // EOB runs are handled by the eob logic
661                        break 'no_eob;
662                    }
663                } else {
664                    // libjpeg-turbo also doesn't return an error here, so let's also warn.
665                    if symbol != 1 {
666                        warn!("Bad Huffman code, corrupt JPEG?");
667                    }
668                    // get sign bit
669                    // We assume we have enough bits, which should be correct for sane images
670                    // since we refill by 32 above
671                    if self.get_bit() == 1 {
672                        symbol = i32::from(bit);
673                    } else {
674                        symbol = i32::from(-bit);
675                    }
676                }
677
678                // Advance over already nonzero coefficients  appending
679                // correction bits to the non-zeroes.
680                // A correction bit is 1 if the absolute value of the coefficient must be increased
681
682                if k <= self.spec_end {
683                    'advance_nonzero: loop {
684                        let coefficient = &mut block[UN_ZIGZAG[k as usize & 63] & 63];
685
686                        if *coefficient != 0 {
687                            if self.bits_left < 1 {
688                                self.refill(reader)?;
689                                if self.bits_left < 1 && self.marker.is_some() {
690                                    return Err(DecodeErrors::Format(
691                                        "Marker found where not expected in refine bit".to_string()
692                                    ));
693                                }
694                            }
695                            if self.get_bit() == 1 && (*coefficient & bit) == 0 {
696                                if *coefficient > 0 {
697                                    *coefficient = coefficient.wrapping_add(bit);
698                                } else {
699                                    *coefficient = coefficient.wrapping_sub(bit);
700                                }
701                            }
702                        } else {
703                            r -= 1;
704
705                            if r < 0 {
706                                // reached target zero coefficient.
707                                break 'advance_nonzero;
708                            }
709                        };
710
711                        if k == self.spec_end {
712                            break 'advance_nonzero;
713                        }
714
715                        k += 1;
716                    }
717                }
718
719                if symbol != 0 {
720                    let pos = UN_ZIGZAG[k as usize & 63];
721                    // output new non-zero coefficient.
722                    block[pos & 63] = symbol as i16;
723                }
724
725                k += 1;
726
727                if k > self.spec_end {
728                    break 'no_eob;
729                }
730            }
731        }
732        if self.eob_run > 0 {
733            // only run if block does not consists of purely zeroes
734            if &block[1..] != &[0; 63] {
735                self.refill(reader)?;
736
737                while k <= self.spec_end {
738                    let coefficient = &mut block[UN_ZIGZAG[k as usize & 63] & 63];
739
740                    if *coefficient != 0 && self.get_bit() == 1 {
741                        // check if we already modified it, if so do nothing, otherwise
742                        // append the correction bit.
743                        if (*coefficient & bit) == 0 {
744                            if *coefficient >= 0 {
745                                *coefficient = coefficient.wrapping_add(bit);
746                            } else {
747                                *coefficient = coefficient.wrapping_sub(bit);
748                            }
749                        }
750                    }
751                    if self.bits_left < 1 {
752                        // refill at the last possible moment
753                        self.refill(reader)?;
754                    }
755                    k += 1;
756                }
757            }
758            // count a block completed in EOB run
759            self.eob_run -= 1;
760        }
761        return Ok(true);
762    }
763
764    pub fn update_progressive_params(&mut self, _ah: u8, al: u8, spec_start: u8, spec_end: u8) {
765        self.successive_low_mask = 1i16 << al;
766        self.spec_start = spec_start;
767        self.spec_end = spec_end;
768    }
769
770    /// Reset the stream if we have a restart marker
771    ///
772    /// Restart markers indicate drop those bits in the stream and zero out
773    /// everything
774    #[cold]
775    pub fn reset(&mut self) {
776        self.bits_left = 0;
777        self.marker = None;
778        self.buffer = 0;
779        self.aligned_buffer = 0;
780        self.eob_run = 0;
781    }
782}
783
784/// Do the equivalent of JPEG HUFF_EXTEND
785#[inline(always)]
786fn huff_extend(x: i32, s: i32) -> i32 {
787    // if x<s return x else return x+offset[s] where offset[s] = ( (-1<<s)+1)
788    (x) + ((((x) - (1 << ((s) - 1))) >> 31) & (((-1) << (s)) + 1))
789}
790
791const fn has_zero(v: u32) -> bool {
792    // Retrieved from Stanford bithacks
793    // @ https://graphics.stanford.edu/~seander/bithacks.html#ZeroInWord
794    return !((((v & 0x7F7F_7F7F) + 0x7F7F_7F7F) | v) | 0x7F7F_7F7F) != 0;
795}
796
797const fn has_byte(b: u32, val: u8) -> bool {
798    // Retrieved from Stanford bithacks
799    // @ https://graphics.stanford.edu/~seander/bithacks.html#ZeroInWord
800    has_zero(b ^ ((!0_u32 / 255) * (val as u32)))
801}
802
803// mod tests {
804//     use zune_core::bytestream::ZCursor;
805//     use zune_core::colorspace::ColorSpace;
806//     use zune_core::options::DecoderOptions;
807//
808//     use crate::JpegDecoder;
809//
810//     #[test]
811//     fn test_image() {
812//         let img = "/Users/etemesi/Downloads/test_IDX_45_RAND_168601280367171438891916_minimized_837.jpg";
813//         let data = std::fs::read(img).unwrap();
814//         let options = DecoderOptions::new_cmd().jpeg_set_out_colorspace(ColorSpace::RGB);
815//         let mut decoder = JpegDecoder::new_with_options(ZCursor::new(&data[..]), options);
816//
817//         decoder.decode().unwrap();
818//         println!("{:?}", decoder.options.jpeg_get_out_colorspace())
819//     }
820// }