ljpeg 0.1.2

pub(crate) trait BitPump {
    fn peek_bits(&mut self, num: u32) -> u32;
    fn consume_bits(&mut self, num: u32);

    #[inline]
    fn get_bits(&mut self, num: u32) -> u32 {
        if num == 0 {
            return 0;
        }

        let val = self.peek_bits(num);
        self.consume_bits(num);

        val
    }
}

#[derive(Debug, Copy, Clone)]
pub(crate) struct BitPumpMSB32<'a> {
    buffer: &'a [u8],
    pos: usize,
    bits: u64,
    nbits: u32,
}

impl BitPumpMSB32<'_> {
    pub(crate) fn new(src: &[u8]) -> BitPumpMSB32 {
        BitPumpMSB32 {
            buffer: src,
            pos: 0,
            bits: 0,
            nbits: 0,
        }
    }
}

impl BitPump for BitPumpMSB32<'_> {
    #[inline]
    fn peek_bits(&mut self, num: u32) -> u32 {
        if num > self.nbits {
            let inbits =
                u32::from_le_bytes(self.buffer[self.pos..self.pos + 4].try_into().unwrap()) as u64;
            self.bits = (self.bits << 32) | inbits;
            self.pos += 4;
            self.nbits += 32;
        }
        (self.bits >> (self.nbits - num)) as u32
    }

    #[inline]
    fn consume_bits(&mut self, num: u32) {
        self.nbits -= num;
        self.bits &= (1 << self.nbits) - 1;
    }
}

#[derive(Debug, Copy, Clone)]
pub(crate) struct BitPumpJPEG<'a> {
    buffer: &'a [u8],
    pos: usize,
    bits: u64,
    nbits: u32,
    finished: bool,
}

impl BitPumpJPEG<'_> {
    pub(crate) fn new(src: &[u8]) -> BitPumpJPEG {
        BitPumpJPEG {
            buffer: src,
            pos: 0,
            bits: 0,
            nbits: 0,
            finished: false,
        }
    }
}

impl BitPump for BitPumpJPEG<'_> {
    #[inline]
    fn peek_bits(&mut self, num: u32) -> u32 {
        if num > self.nbits && !self.finished {
            if (self.buffer.len() >= 4)
                && self.pos < self.buffer.len() - 4
                && self.buffer[self.pos] != 0xff
                && self.buffer[self.pos + 1] != 0xff
                && self.buffer[self.pos + 2] != 0xff
                && self.buffer[self.pos + 3] != 0xff
            {
                let inbits =
                    u32::from_be_bytes(self.buffer[self.pos..self.pos + 4].try_into().unwrap())
                        as u64;
                self.bits = (self.bits << 32) | inbits;
                self.pos += 4;
                self.nbits += 32;
            } else {
                // Read 32 bits the hard way
                let mut read_bytes = 0;
                while read_bytes < 4 && !self.finished {
                    let byte = {
                        if self.pos >= self.buffer.len() {
                            self.finished = true;
                            0
                        } else {
                            let nextbyte = self.buffer[self.pos];
                            if nextbyte != 0xff {
                                nextbyte
                            } else if self.buffer[self.pos + 1] == 0x00 {
                                self.pos += 1; // Skip the extra byte used to mark 255
                                nextbyte
                            } else {
                                self.finished = true;
                                0
                            }
                        }
                    };
                    self.bits = (self.bits << 8) | (byte as u64);
                    self.pos += 1;
                    self.nbits += 8;
                    read_bytes += 1;
                }
            }
        }
        if num > self.nbits && self.finished {
            // Stuff with zeroes to not fail to read
            self.bits <<= 32;
            self.nbits += 32;
        }

        (self.bits >> (self.nbits - num)) as u32
    }

    #[inline]
    fn consume_bits(&mut self, num: u32) {
        debug_assert!(num <= self.nbits);
        self.nbits -= num;
        self.bits &= (1 << self.nbits) - 1;
    }
}

#[derive(Debug, Copy, Clone)]
pub(crate) struct ByteStream<'a> {
    buffer: &'a [u8],
    pos: usize,
}

impl ByteStream<'_> {
    #[inline]
    pub(crate) fn new(buffer: &[u8]) -> ByteStream {
        ByteStream { buffer, pos: 0 }
    }

    #[inline]
    pub(crate) fn get_pos(&self) -> usize {
        self.pos
    }

    #[inline]
    pub(crate) fn peek_u8(&self) -> u8 {
        self.buffer[self.pos]
    }
    #[inline]
    pub(crate) fn get_u8(&mut self) -> u8 {
        let val = self.peek_u8();
        self.pos += 1;
        val
    }

    #[inline]
    pub(crate) fn peek_u16(&self) -> u16 {
        u16::from_be_bytes(self.buffer[self.pos..self.pos + 2].try_into().unwrap())
    }
    #[inline]
    pub(crate) fn get_u16(&mut self) -> u16 {
        let val = self.peek_u16();
        self.pos += 2;
        val
    }

    #[inline]
    pub(crate) fn skip_to_marker(&mut self) -> Option<usize> {
        let mut skip_count = 0;
        while !(self.buffer[self.pos] == 0xFF
            && self.buffer[self.pos + 1] != 0
            && self.buffer[self.pos + 1] != 0xFF)
        {
            self.pos += 1;
            skip_count += 1;
            if self.pos >= self.buffer.len() {
                return None;
            }
        }
        self.pos += 1; // Make the next byte the marker
        Some(skip_count + 1)
    }
}