zune_jpeg/

huffman.rs

/*
 * Copyright (c) 2023.
 *
 * This software is free software;
 *
 * You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
 */

//! This file contains a single struct `HuffmanTable` that
//! stores Huffman tables needed during `BitStream` decoding.
#![allow(clippy::similar_names, clippy::module_name_repetitions)]

use alloc::string::ToString;

use crate::errors::DecodeErrors;

/// Determines how many bits of lookahead we have for our bitstream decoder.

pub const HUFF_LOOKAHEAD: u8 = 9;

/// A struct which contains necessary tables for decoding a JPEG
/// huffman encoded bitstream

pub struct HuffmanTable {
    // element `[0]` of each array is unused
    /// largest code of length k
    pub(crate) maxcode: [i32; 18],
    /// offset for codes of length k
    /// Answers the question, where do code-lengths of length k end
    /// Element 0 is unused
    pub(crate) offset:  [i32; 18],
    /// lookup table for fast decoding
    ///
    /// top  bits above HUFF_LOOKAHEAD contain the code length.
    ///
    /// Lower (8) bits contain the symbol in order of increasing code length.
    pub(crate) lookup:  [i32; 1 << HUFF_LOOKAHEAD],

    /// A table which can be used to decode small AC coefficients and
    /// do an equivalent of receive_extend
    pub(crate) ac_lookup: Option<[i16; 1 << HUFF_LOOKAHEAD]>,

    /// Directly represent contents of a JPEG DHT marker
    ///
    /// \# number of symbols with codes of length `k` bits
    // bits[0] is unused
    /// Symbols in order of increasing code length
    pub(crate) values: [u8; 256]
}

impl HuffmanTable {
    pub fn new(
        codes: &[u8; 17], values: [u8; 256], is_dc: bool, is_progressive: bool
    ) -> Result<HuffmanTable, DecodeErrors> {
        let too_long_code = (i32::from(HUFF_LOOKAHEAD) + 1) << HUFF_LOOKAHEAD;
        let mut p = HuffmanTable {
            maxcode: [0; 18],
            offset: [0; 18],
            lookup: [too_long_code; 1 << HUFF_LOOKAHEAD],
            values,
            ac_lookup: None
        };

        p.make_derived_table(is_dc, is_progressive, codes)?;

        Ok(p)
    }

    /// Create a new huffman tables with values that aren't fixed
    /// used by fill_mjpeg_tables
    pub fn new_unfilled(
        codes: &[u8; 17], values: &[u8], is_dc: bool, is_progressive: bool
    ) -> Result<HuffmanTable, DecodeErrors> {
        let mut buf = [0; 256];
        buf[..values.len()].copy_from_slice(values);
        HuffmanTable::new(codes, buf, is_dc, is_progressive)
    }

    /// Compute derived values for a Huffman table
    ///
    /// This routine performs some validation checks on the table
    #[allow(
        clippy::cast_possible_truncation,
        clippy::cast_possible_wrap,
        clippy::cast_sign_loss,
        clippy::too_many_lines,
        clippy::needless_range_loop
    )]
    fn make_derived_table(
        &mut self, is_dc: bool, _is_progressive: bool, bits: &[u8; 17]
    ) -> Result<(), DecodeErrors> {
        // build a list of code size
        let mut huff_size = [0; 257];
        // Huffman code lengths
        let mut huff_code: [u32; 257] = [0; 257];
        // figure C.1 make table of Huffman code length for each symbol
        let mut p = 0;

        for l in 1..=16 {
            let mut i = i32::from(bits[l]);
            // table overrun is checked before ,so we dont need to check
            while i != 0 {
                huff_size[p] = l as u8;
                p += 1;
                i -= 1;
            }
        }

        huff_size[p] = 0;

        let num_symbols = p;
        // Generate the codes themselves
        // We also validate that the counts represent a legal Huffman code tree
        let mut code = 0;
        let mut si = i32::from(huff_size[0]);

        p = 0;

        while huff_size[p] != 0 {
            while i32::from(huff_size[p]) == si {
                huff_code[p] = code;
                code += 1;
                p += 1;
            }
            // maximum code of length si, pre-shifted by 16-k bits
            self.maxcode[si as usize] = (code << (16 - si)) as i32;
            // code is now 1 more than the last code used for code-length si; but
            // it must still fit in si bits, since no code is allowed to be all ones.
            if (code as i32) >= (1 << si) {
                return Err(DecodeErrors::HuffmanDecode("Bad Huffman Table".to_string()));
            }

            code <<= 1;
            si += 1;
        }

        // Figure F.15 generate decoding tables for bit-sequential decoding
        p = 0;

        for l in 0..=16 {
            if bits[l] == 0 {
                // -1 if no codes of this length
                self.maxcode[l] = -1;
            } else {
                // offset[l]=codes[index of 1st symbol of code length l
                // minus minimum code of length l]
                self.offset[l] = (p as i32) - (huff_code[p]) as i32;
                p += usize::from(bits[l]);
            }
        }

        self.offset[17] = 0;
        // we ensure that decode terminates
        self.maxcode[17] = 0x000F_FFFF;

        /*
         * Compute lookahead tables to speed up decoding.
         * First we set all the table entries to 0(left justified), indicating "too long";
         * (Note too long was set during initialization)
         * then we iterate through the Huffman codes that are short enough and
         * fill in all the entries that correspond to bit sequences starting
         * with that code.
         */

        p = 0;

        for l in 1..=HUFF_LOOKAHEAD {
            for _ in 1..=i32::from(bits[usize::from(l)]) {
                // l -> Current code length,
                // p => Its index in self.code and self.values
                // Generate left justified code followed by all possible bit sequences
                let mut look_bits = (huff_code[p] as usize) << (HUFF_LOOKAHEAD - l);

                for _ in 0..1 << (HUFF_LOOKAHEAD - l) {
                    self.lookup[look_bits] =
                        (i32::from(l) << HUFF_LOOKAHEAD) | i32::from(self.values[p]);
                    look_bits += 1;
                }

                p += 1;
            }
        }
        // build an ac table that does an equivalent of decode and receive_extend
        if !is_dc {
            let mut fast = [255; 1 << HUFF_LOOKAHEAD];
            // Iterate over number of symbols
            for i in 0..num_symbols {
                // get code size for an item
                let s = huff_size[i];

                if s <= HUFF_LOOKAHEAD {
                    // if it's lower than what we need for our lookup table create the table
                    let c = (huff_code[i] << (HUFF_LOOKAHEAD - s)) as usize;
                    let m = (1 << (HUFF_LOOKAHEAD - s)) as usize;

                    for j in 0..m {
                        fast[c + j] = i as i16;
                    }
                }
            }

            // build a table that decodes both magnitude and value of small ACs in
            // one go.
            let mut fast_ac = [0; 1 << HUFF_LOOKAHEAD];

            for i in 0..(1 << HUFF_LOOKAHEAD) {
                let fast_v = fast[i];

                if fast_v < 255 {
                    // get symbol value from AC table
                    let rs = self.values[fast_v as usize];
                    // shift by 4 to get run length
                    let run = i16::from((rs >> 4) & 15);
                    // get magnitude bits stored at the lower 3 bits
                    let mag_bits = i16::from(rs & 15);
                    // length of the bit we've read
                    let len = i16::from(huff_size[fast_v as usize]);

                    if mag_bits != 0 && (len + mag_bits) <= i16::from(HUFF_LOOKAHEAD) {
                        // magnitude code followed by receive_extend code
                        let mut k = (((i as i16) << len) & ((1 << HUFF_LOOKAHEAD) - 1))
                            >> (i16::from(HUFF_LOOKAHEAD) - mag_bits);
                        let m = 1 << (mag_bits - 1);

                        if k < m {
                            k += (!0_i16 << mag_bits) + 1;
                        };

                        // if result is small enough fit into fast ac table
                        if (-128..=127).contains(&k) {
                            fast_ac[i] = (k << 8) + (run << 4) + (len + mag_bits);
                        }
                    }
                }
            }
            self.ac_lookup = Some(fast_ac);
        }

        // Validate symbols as being reasonable
        // For AC tables, we make no check, but accept all byte values 0..255
        // For DC tables, we require symbols to be in range 0..15
        if is_dc {
            for i in 0..num_symbols {
                let sym = self.values[i];

                if sym > 15 {
                    return Err(DecodeErrors::HuffmanDecode("Bad Huffman Table".to_string()));
                }
            }
        }

        Ok(())
    }
}