zenjxl-decoder 0.3.8

// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

use crate::bit_reader::BitReader;
use crate::error::Error;

use crate::util::CeilLog2;

#[derive(Debug)]
pub struct HybridUint {
    split_token: u32,
    split_exponent: u32,
    msb_in_token: u32,
    lsb_in_token: u32,
}

impl HybridUint {
    pub(super) fn is_split_exponent_zero(&self) -> bool {
        self.split_exponent == 0
    }

    pub fn decode(log_alpha_size: usize, br: &mut BitReader) -> Result<HybridUint, Error> {
        let split_exponent = br.read((log_alpha_size + 1).ceil_log2())? as u32;
        let split_token = 1u32 << split_exponent;
        let msb_in_token;
        let lsb_in_token;
        if split_exponent != log_alpha_size as u32 {
            let nbits = (split_exponent + 1).ceil_log2() as usize;
            msb_in_token = br.read(nbits)? as u32;
            if msb_in_token > split_exponent {
                return Err(Error::InvalidUintConfig(split_exponent, msb_in_token, None));
            }
            let nbits = (split_exponent - msb_in_token + 1).ceil_log2() as usize;
            lsb_in_token = br.read(nbits)? as u32;
        } else {
            msb_in_token = 0;
            lsb_in_token = 0;
        }
        if lsb_in_token + msb_in_token > split_exponent {
            return Err(Error::InvalidUintConfig(
                split_exponent,
                msb_in_token,
                Some(lsb_in_token),
            ));
        }
        Ok(HybridUint {
            split_token,
            split_exponent,
            msb_in_token,
            lsb_in_token,
        })
    }

    /// Returns true if this config matches the 420 pattern (common in e3 images):
    /// split_exponent=4, msb_in_token=2, lsb_in_token=0
    #[inline(always)]
    pub fn is_config_420(&self) -> bool {
        self.split_exponent == 4
            && self.split_token == 16
            && self.msb_in_token == 2
            && self.lsb_in_token == 0
    }

    /// Specialized fast path for 420 config:
    /// split_exponent=4, msb_in_token=2, lsb_in_token=0
    ///
    /// `nbits_acc` accumulates raw nbits values via OR. If any call produces
    /// nbits >= 32 (invalid bitstream), bit 5+ will be set in the accumulator.
    /// Check `nbits_acc >= 32` after a batch of reads.
    #[inline(always)]
    pub fn read_config_420(symbol: u32, br: &mut BitReader, nbits_acc: &mut u32) -> u32 {
        if symbol < 16 {
            return symbol;
        }

        // Equivalent to: 2 + ((symbol - 16) >> 2)
        let nbits_raw = (symbol >> 2) - 2;
        *nbits_acc |= nbits_raw;
        let nbits = nbits_raw & 31;
        let bits = br.read_optimistic(nbits as usize) as u32;
        let hi = (symbol & 3) | 4;

        (hi << nbits) | bits
    }

    /// Reads a hybrid uint value from the bitstream.
    ///
    /// `nbits_acc` accumulates raw nbits values via OR. If any call produces
    /// nbits >= 32 (invalid bitstream), bit 5+ will be set in the accumulator.
    /// Check `nbits_acc >= 32` after a batch of reads.
    #[inline(always)]
    pub fn read(&self, symbol: u32, br: &mut BitReader, nbits_acc: &mut u32) -> u32 {
        if symbol < self.split_token {
            return symbol;
        }
        // Fast path: when msb_in_token == 0 && lsb_in_token == 0, hi is always 1
        // and the low/shift computation is a no-op.
        if self.msb_in_token == 0 && self.lsb_in_token == 0 {
            let nbits_raw = self.split_exponent + symbol - self.split_token;
            *nbits_acc |= nbits_raw;
            let nbits = nbits_raw & 31;
            let bits = br.read_optimistic(nbits as usize) as u32;
            return (1 << nbits) | bits;
        }
        let bits_in_token = self.lsb_in_token + self.msb_in_token;
        let nbits_raw =
            self.split_exponent - bits_in_token + ((symbol - self.split_token) >> bits_in_token);
        *nbits_acc |= nbits_raw;
        let nbits = nbits_raw & 31;
        let low = symbol & ((1 << self.lsb_in_token) - 1);
        let symbol_nolow = symbol >> self.lsb_in_token;
        let bits = br.read_optimistic(nbits as usize) as u32;
        let hi = (symbol_nolow & ((1 << self.msb_in_token) - 1)) | (1 << self.msb_in_token);
        (((hi << nbits) | bits) << self.lsb_in_token) | low
    }
}

#[cfg(test)]
impl HybridUint {
    pub fn new(split_exponent: u32, msb_in_token: u32, lsb_in_token: u32) -> Self {
        Self {
            split_token: 1 << split_exponent,
            split_exponent,
            msb_in_token,
            lsb_in_token,
        }
    }
}

#[cfg(test)]
mod test {
    #[test]
    fn test_hybrid_uint_decode_invalid() {
        use super::*;
        let mut br = BitReader::new(&[10, 75, 10, 75, 168, 139, 132, 255, 244]);
        br.skip_bits(1).unwrap();
        if let Ok(uint) = HybridUint::decode(15, &mut br) {
            let mut acc = 0u32;
            uint.read(1022, &mut br, &mut acc);
        }
    }
}