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// }