packbits-rle 0.2.0

#![doc = include_str!("../README.md")]
#![feature(seek_stream_len)]
#![feature(iter_array_chunks)]
use std::{
    io::{self, ErrorKind},
    iter,
};

mod reader;

pub use reader::Reader;

#[derive(Debug, thiserror::Error)]
/// General error used in packbit extraction
pub enum Error {
    /// Input does not provide enough data to fill the requested size
    #[error("Input does not contain enough packbits data to fill the buffer")]
    NotEnoughInputData,

    /// The output buffer is too small or the decoded size does not fit into the requested size
    #[error("Output buffer is too small to hold decoded data")]
    TooMuchOutputData,

    /// Input has additional unread data left after decoding
    #[error("Input has unread data left after filling the output buffer")]
    TooMuchInputData,

    /// An operation (like read or seek) on the underlying reader has failed
    #[error(transparent)]
    Io(#[from] io::Error),
}

impl From<Error> for io::Error {
    fn from(value: Error) -> Self {
        match value {
            Error::Io(error) => error,
            me => io::Error::other(me),
        }
    }
}
#[derive(Debug)]
/// A packbits command, also called the flag-counter byte, specifying how the next few bytes in a stream will be treated
pub enum Command {
    /// Escape code, next byte will be emitted literally
    Escape,
    /// Copy the next `self.0` bytes verbatim to output
    Literal(u8),
    /// Repeat the next byte `self.0` times
    Repeat(u16),
}

impl Command {
    #[inline]
    /// Produces a byte and returns the next decoder state ([`Operation`]) at the same time
    pub fn execute<T: io::Read>(self, reader: &mut T) -> Result<(u8, Operation), OperationError> {
        match self {
            Command::Escape => {
                let byte = Self::read_byte(reader)?;
                Ok((byte, Operation::Literal(0)))
            }
            Command::Literal(count) => {
                let literal = Self::read_byte(reader)?;
                Ok((literal, Operation::Literal(count - 1)))
            }
            Command::Repeat(count) => {
                let literal = Self::read_byte(reader)?;
                Ok((literal, Operation::Repeat(literal, count - 1)))
            }
        }
    }

    #[inline]
    fn read_byte<T: io::Read>(reader: &mut T) -> Result<u8, OperationError> {
        let mut buf = [0u8];
        match reader.read(&mut buf) {
            Ok(0) => Err(OperationError::UnexpectedEof),
            Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => {
                Err(OperationError::UnexpectedEof)
            }
            Err(e) => Err(e)?,
            Ok(_) => Ok(buf[0]),
        }
    }
}

impl From<i8> for Command {
    #[inline]
    fn from(val: i8) -> Self {
        match val {
            -128 => Command::Escape,
            n if n < 0 => Command::Repeat((1 - n) as u16),
            n => Command::Literal((1 + n) as u8),
        }
    }
}

impl From<u8> for Command {
    #[inline]
    fn from(val: u8) -> Self {
        match val {
            n if n < 128 => Command::Literal(1 + n),
            128 => Command::Escape,
            n => Command::Repeat(257 - n as u16),
        }
    }
}

#[derive(Debug)]
/// Describes an unpacking operation, that can be fed additional data to recieved an output
/// byte and the next operation
pub enum Operation {
    /// We're in the middle of returning literal bytes
    Literal(u8),
    /// The unpacker will emit `.1` more repetitions of byte `.0`
    Repeat(u8, u16),
}

#[derive(Debug)]
/// Describes an unpacking operation, that can be fed additional data to recieved an output
/// byte and the next operation
pub enum Operation16 {
    /// We're in the middle of returning literal bytes
    Literal(u8),
    /// The unpacker will emit `.1` more repetitions of byte `.0`
    Repeat(u16, u16),
}

#[derive(Debug, thiserror::Error)]
/// Error for executing [`Operation`]s, on insufficient input it provides a [`Command`] to continue the
/// operation when more data becomes available
pub enum OperationError {
    /// The input stream ended, before the operation could complete.
    ///
    /// If more data becomes available, the provided command can be used to continue unpacking
    /// where we left off.
    #[error("More input data is required to executed command")]
    InsufficientInput(Command),

    /// The input stream ended with no unfinished command, this could be a good place to stop
    /// unpacking if enough input bytes have been read or enough output bytes provided.
    #[error("More input data is required to continue")]
    UnexpectedEof,

    /// An operation like _read_ or _seek_ on the underlying underlying reader has failed.
    #[error(transparent)]
    Io(#[from] io::Error),
}

impl From<OperationError> for io::Error {
    fn from(value: OperationError) -> Self {
        match value {
            OperationError::Io(error) => error,
            other => io::Error::other(other),
        }
    }
}

impl Default for Operation {
    fn default() -> Self {
        Self::new()
    }
}

impl Default for Operation16 {
    fn default() -> Self {
        Self::new()
    }
}
impl Operation16 {
    /// Create a new empty operation, advancing ([`Operation::advance`]) it will read an entire new operation from the
    /// underlying reader.
    pub fn new() -> Self {
        Self::Literal(0)
    }

    #[inline]
    fn read_byte<T: io::Read>(reader: &mut T) -> Result<u8, OperationError> {
        let mut buf = [0u8];
        match reader.read(&mut buf) {
            Ok(0) => Err(OperationError::UnexpectedEof),
            Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => {
                Err(OperationError::UnexpectedEof)
            }
            Err(e) => Err(e)?,
            Ok(_) => Ok(buf[0]),
        }
    }

    #[inline]
    fn read_word<T: io::Read>(reader: &mut T) -> Result<u16, OperationError> {
        let mut buf = [0u8, 0u8];
        match reader.read(&mut buf) {
            Ok(0) => Err(OperationError::UnexpectedEof),
            Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => {
                Err(OperationError::UnexpectedEof)
            }
            Err(e) => Err(e)?,
            Ok(_) => Ok(((buf[0] as u16) << 8) | buf[1] as u16),
        }
    }

    #[inline]
    /// Produce a byte and return next state of the decoder at the same time
    ///
    /// If `self` was already completed, this function will try to fetch the next command from the underlying
    /// reader. Otherwise it will just perform the repetition, or read the next literal.
    pub fn advance<T: io::Read>(&self, reader: &mut T) -> Result<(u16, Self), OperationError> {
        match self {
            op if op.is_completed() => match Self::read_byte(reader)? {
                128 => match Self::read_byte(reader) {
                    Ok(byte) => Ok((
                        (byte as u16) << 8 | (Self::read_byte(reader)? as u16),
                        Operation16::Literal(0),
                    )),
                    Err(OperationError::UnexpectedEof) => {
                        Err(OperationError::InsufficientInput(Command::Escape))
                    }
                    Err(e) => Err(e),
                },
                n if n < 128 => match Self::read_word(reader) {
                    Ok(byte) => Ok((byte, Operation16::Literal(n))),
                    Err(OperationError::UnexpectedEof) => {
                        Err(OperationError::InsufficientInput(Command::Literal(n + 1)))
                    }
                    Err(e) => Err(e),
                },
                n => match Self::read_word(reader) {
                    Err(OperationError::UnexpectedEof) => Err(OperationError::InsufficientInput(
                        Command::Repeat(256 - n as u16 + 1),
                    )),
                    Ok(byte) => Ok((byte, Operation16::Repeat(byte, 256 - n as u16))),
                    Err(e) => Err(e),
                },
            },
            Operation16::Repeat(value, count) => {
                Ok((*value, Operation16::Repeat(*value, count - 1)))
            }
            Operation16::Literal(count) => match Self::read_word(reader) {
                Ok(byte) => Ok((byte, Self::Literal(count - 1))),
                Err(OperationError::UnexpectedEof) => {
                    Err(OperationError::InsufficientInput(Command::Literal(*count)))
                }
                Err(e) => Err(e),
            },
        }
    }

    #[inline]
    /// Determines if the current operation has been advanced to completion.
    pub fn is_completed(&self) -> bool {
        matches!(self, Self::Literal(0) | Self::Repeat(_, 0))
    }
}

impl Operation {
    /// Create a new empty operation, advancing ([`Operation::advance`]) it will read an entire new operation from the
    /// underlying reader.
    pub fn new() -> Self {
        Self::Literal(0)
    }

    #[inline]
    fn read_byte<T: io::Read>(reader: &mut T) -> Result<u8, OperationError> {
        let mut buf = [0u8];
        match reader.read(&mut buf) {
            Ok(0) => Err(OperationError::UnexpectedEof),
            Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => {
                Err(OperationError::UnexpectedEof)
            }
            Err(e) => Err(e)?,
            Ok(_) => Ok(buf[0]),
        }
    }

    #[inline]
    /// Produce a byte and return next state of the decoder at the same time
    ///
    /// If `self` was already completed, this function will try to fetch the next command from the underlying
    /// reader. Otherwise it will just perform the repetition, or read the next literal.
    pub fn advance<T: io::Read>(&self, reader: &mut T) -> Result<(u8, Self), OperationError> {
        match self {
            op if op.is_completed() => match Self::read_byte(reader)? {
                128 => match Self::read_byte(reader) {
                    Ok(byte) => Ok((byte, Operation::Literal(0))),
                    Err(OperationError::UnexpectedEof) => {
                        Err(OperationError::InsufficientInput(Command::Escape))
                    }
                    Err(e) => Err(e),
                },
                n if n < 128 => match Self::read_byte(reader) {
                    Ok(byte) => Ok((byte, Operation::Literal(n))),
                    Err(OperationError::UnexpectedEof) => {
                        Err(OperationError::InsufficientInput(Command::Literal(n + 1)))
                    }
                    Err(e) => Err(e),
                },
                n => match Self::read_byte(reader) {
                    Err(OperationError::UnexpectedEof) => Err(OperationError::InsufficientInput(
                        Command::Repeat(256 - n as u16 + 1),
                    )),
                    Ok(byte) => Ok((byte, Operation::Repeat(byte, 256 - n as u16))),
                    Err(e) => Err(e),
                },
            },
            Operation::Repeat(value, count) => Ok((*value, Operation::Repeat(*value, count - 1))),
            Operation::Literal(count) => match Self::read_byte(reader) {
                Ok(byte) => Ok((byte, Self::Literal(count - 1))),
                Err(OperationError::UnexpectedEof) => {
                    Err(OperationError::InsufficientInput(Command::Literal(*count)))
                }
                Err(e) => Err(e),
            },
        }
    }

    #[inline]
    /// Determines if the current operation has been advanced to completion.
    pub fn is_completed(&self) -> bool {
        matches!(self, Self::Literal(0) | Self::Repeat(_, 0))
    }
}

#[cfg(test)]
mod operation {
    use crate::{Command, Operation, OperationError};
    use std::io;

    #[test]
    fn read_literal() {
        let mut reader = io::Cursor::new(b"\x01\xAB\xBC");
        assert!(matches!(
            Operation::default().advance(&mut reader),
            Ok((0xab, Operation::Literal(1)))
        ));
    }

    #[test]
    fn read_repeat() {
        let mut reader = io::Cursor::new(b"\xFF\xAB");
        assert!(matches!(
            Operation::default().advance(&mut reader),
            Ok((0xab, Operation::Repeat(0xab, 1)))
        ));
    }

    #[test]
    fn read_escape() {
        let mut reader = io::Cursor::new(b"\x80\xAB");
        assert!(matches!(
            Operation::default().advance(&mut reader),
            Ok((0xAB, Operation::Repeat(_, 0) | Operation::Literal(0)))
        ));
    }

    #[test]
    fn read_partial_literal() {
        let mut reader = io::Cursor::new(b"\x01");
        assert!(matches!(
            Operation::default().advance(&mut reader),
            Err(OperationError::InsufficientInput(Command::Literal(2)))
        ));
    }

    #[test]
    fn read_partial_literal_to_completion() {
        let mut reader = io::Cursor::new(b"\xAB");
        assert!(matches!(
            Command::Literal(2).execute(&mut reader),
            Ok((0xab, Operation::Literal(1)))
        ));
    }

    #[test]
    fn read_partial_repeat() {
        let mut reader = io::Cursor::new(b"\xFF");
        assert!(matches!(
            Operation::default().advance(&mut reader),
            Err(OperationError::InsufficientInput(Command::Repeat(2)))
        ));
    }

    #[test]
    fn read_partial_repeat_to_completion() {
        let mut reader = io::Cursor::new(b"\xAB");
        assert!(matches!(
            Command::Repeat(2).execute(&mut reader),
            Ok((0xab, Operation::Repeat(0xab, 1)))
        ));
    }

    #[test]
    fn read_partial_escape() {
        let mut reader = io::Cursor::new(b"\x80");
        assert!(matches!(
            Operation::default().advance(&mut reader),
            Err(OperationError::InsufficientInput(Command::Escape))
        ));
    }

    #[test]
    fn read_partial_escape_to_completion() {
        let mut reader = io::Cursor::new(b"\x80");
        assert!(matches!(
            Command::Escape.execute(&mut reader),
            Ok((0x80, Operation::Literal(0) | Operation::Repeat(_, 0)))
        ));
    }
}

pub fn unpack_words_exact(input: &[u8], count: usize) -> Result<Vec<u8>, Error> {
    let mut output = Vec::with_capacity(count);
    let mut remaining_output = count;
    let mut i = 0;

    loop {
        if remaining_output == 0 {
            return if i == input.len() {
                Ok(output)
            } else {
                Err(Error::TooMuchInputData)
            };
        }

        if i == input.len() {
            return Err(Error::NotEnoughInputData);
        }

        i = match input[i].into() {
            Command::Escape => i + 1,
            Command::Literal(count) => {
                if input.len() < i + count as usize {
                    return Err(Error::NotEnoughInputData);
                }

                if count as usize > remaining_output {
                    return Err(Error::TooMuchOutputData);
                }

                output.extend(&input[i + 1..(i + 1 + (count as usize * 2))]);
                remaining_output -= count as usize * 2;
                i + count as usize * 2 + 1
            }
            Command::Repeat(count) => {
                if input.len() < i + 2 {
                    return Err(Error::NotEnoughInputData);
                }

                if count as usize * 2 > remaining_output {
                    return Err(Error::TooMuchOutputData);
                }

                for _ in 0..(count as usize) {
                    output.extend(&input[(i + 1)..=(i + 2)]);
                }

                remaining_output -= count as usize * 2;
                i + 3
            }
        }
    }
}

/// Unpack `input` into exactly `count` bytes
pub fn unpack_exact(input: &[u8], count: usize) -> Result<Vec<u8>, Error> {
    let mut output = Vec::with_capacity(count);
    let mut remaining_output = count;
    let mut i = 0;

    loop {
        if remaining_output == 0 {
            return if i == input.len() {
                Ok(output)
            } else {
                Err(Error::TooMuchInputData)
            };
        }

        if i == input.len() {
            return Err(Error::NotEnoughInputData);
        }

        i = match input[i].into() {
            Command::Escape => i + 1,
            Command::Literal(count) => {
                if input.len() < i + count as usize {
                    return Err(Error::NotEnoughInputData);
                }

                if count as usize > remaining_output {
                    return Err(Error::TooMuchOutputData);
                }

                output.extend(&input[i + 1..(i + 1 + (count as usize))]);
                remaining_output -= count as usize;
                i + count as usize + 1
            }
            Command::Repeat(count) => {
                if input.len() < i + 1 {
                    return Err(Error::NotEnoughInputData);
                }

                if count as usize > remaining_output {
                    return Err(Error::TooMuchOutputData);
                }

                output.extend(iter::repeat_n(input[i + 1], count as usize));
                remaining_output -= count as usize;
                i + 2
            }
        }
    }
}

/// Unpack all data from `buffer` into a newly allocated vector
pub fn unpack_buf(buffer: &[u8]) -> io::Result<Vec<u8>> {
    unpack(io::Cursor::new(buffer))
}

/// Unpack all data from `reader` into a newly allocated vector
pub fn unpack<R: io::Read>(mut reader: R) -> io::Result<Vec<u8>> {
    let mut output = Vec::new();
    let mut buf = [0u8];
    loop {
        let command = match reader.read_exact(&mut buf) {
            Ok(()) => buf[0].into(),
            Err(e) if e.kind() == ErrorKind::UnexpectedEof => return Ok(output),
            Err(e) => return Err(e)?,
        };

        match command {
            Command::Escape => {
                reader.read_exact(&mut buf)?;
                output.push(buf[0]);
            }
            Command::Literal(count) => {
                let mut buffer = vec![0u8; count as usize];
                reader.read_exact(&mut buffer)?;
                output.append(&mut buffer);
            }
            Command::Repeat(count) => {
                reader.read_exact(&mut buf)?;
                output.reserve(count as usize);
                for _ in 0..count {
                    output.push(buf[0]);
                }
            }
        }
    }
}

/// Extension for [`std::io::Read`]ers to unpack _PackBits_ data
pub trait PackBitsReaderExt {
    /// Unpack PackBit data into the buffer, returning the number of bytes written to the buffer
    fn read_packbits(&mut self, target: &mut [u8]) -> io::Result<usize>;
    fn read_packbits_words(&mut self, target: &mut [u8]) -> io::Result<usize>;
}

impl<T: io::Read> PackBitsReaderExt for T {
    fn read_packbits(&mut self, target: &mut [u8]) -> io::Result<usize> {
        let mut op = Operation::default();
        for (idx, byte) in target.iter_mut().enumerate() {
            match op.advance(self) {
                Err(OperationError::InsufficientInput(_)) => {
                    // TODO: We've read a command that can't be completed, this should be
                    // communicated to users
                    return Ok(idx);
                }
                Err(OperationError::UnexpectedEof) => return Ok(idx),
                Err(other) => return Err(other)?,
                Ok((value, next)) => {
                    *byte = value;
                    op = next;
                }
            }
        }

        Ok(target.len())
    }

    fn read_packbits_words(&mut self, target: &mut [u8]) -> io::Result<usize> {
        let mut op = Operation16::default();
        for (idx, [byte1, byte2]) in target.iter_mut().array_chunks().enumerate() {
            match op.advance(self) {
                Err(OperationError::InsufficientInput(_)) => {
                    // TODO: We've read a command that can't be completed, this should be
                    // communicated to users
                    return Ok(idx * 2);
                }
                Err(OperationError::UnexpectedEof) => return Ok(idx * 2),
                Err(other) => return Err(other)?,
                Ok((value, next)) => {
                    *byte1 = (value >> 8) as u8;
                    *byte2 = (value & 0xFF) as u8;
                    op = next;
                }
            }
        }

        Ok(target.len())
    }
}

#[cfg(test)]
mod test {
    use crate::{Error, unpack_buf};

    use super::{unpack, unpack_exact};

    #[test]
    fn canonical_example() {
        let input = b"\xFE\xAA\x02\x80\x00\x2A\xFD\xAA\x03\x80\x00\x2A\x22\xF7\xAA";
        let unpacked = b"\xAA\xAA\xAA\x80\x00\x2A\xAA\xAA\xAA\xAA\x80\x00\x2A\x22\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA";

        assert_eq!(unpack(input.as_slice()).unwrap(), unpacked);
    }

    #[test]
    fn test_simple_buffer_expansion() {
        let input = b"\xFE\xAA\x02\x80\x00\x2A\xFD\xAA\x03\x80\x00\x2A\x22\xF7\xAA";
        let expectation = b"\xAA\xAA\xAA\x80\x00\x2A\xAA\xAA\xAA\xAA\x80\x00\x2A\x22\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA";
        let result = unpack_buf(input).unwrap();
        assert_eq!(expectation.to_vec(), result);

        let result = unpack_exact(input, 24);
        assert_eq!(expectation.to_vec(), result.unwrap());
    }

    //#[test]
    //fn test_clean_maximum_output() {
    //let input = b"\xFE\xAA\x02\x80\x00\x2A\xFD\xAA\x03\x80\x00\x2A\x22\xF7\xAA\xF7\xAA";
    //let expectation = b"\xAA\xAA\xAA\x80\x00\x2A\xAA\xAA\xAA\xAA\x80\x00\x2A\x22\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA";

    //let result = unpack_until(input, 24);
    //assert_eq!((expectation.to_vec(), 15), result.unwrap());

    //let input = b"";
    //let result = unpack_until(input, 0);
    //assert_eq!((Vec::<u8>::new(), 0), result.unwrap());
    //}

    #[test]
    fn test_too_much_output() {
        let input = b"\xFE\xAA\x02\x80\x00\x2A\xFD\xAA\x03\x80\x00\x2A\x22\xF7\xAA";

        let result = unpack_exact(input, 20);
        assert!(result.is_err_and(|e| matches!(e, Error::TooMuchOutputData)));
    }

    #[test]
    fn test_not_enough_input() {
        let input = b"\xFE\xAA\x02\x80\x00\x2A\xFD\xAA\x03\x80\x00\x2A\x22\xF7\xAA";

        let result = unpack_exact(input, 26);
        assert!(result.is_err_and(|e| matches!(e, Error::NotEnoughInputData)));

        let input = b"";
        let result = unpack_exact(input, 1);
        assert!(result.is_err_and(|e| matches!(e, Error::NotEnoughInputData)));
    }

    mod operation {
        use crate::{Command, Operation, OperationError};
        use std::io;

        #[test]
        fn read_literal() {
            let mut reader = io::Cursor::new(b"\x01\xAB\xBC");
            assert!(matches!(
                Operation::default().advance(&mut reader),
                Ok((0xab, Operation::Literal(1)))
            ));
        }

        #[test]
        fn read_repeat() {
            let mut reader = io::Cursor::new(b"\xFF\xAB");
            assert!(matches!(
                Operation::default().advance(&mut reader),
                Ok((0xab, Operation::Repeat(0xab, 1)))
            ));
        }

        #[test]
        fn read_escape() {
            let mut reader = io::Cursor::new(b"\x80\xAB");
            assert!(matches!(
                Operation::default().advance(&mut reader),
                Ok((0xAB, Operation::Repeat(_, 0) | Operation::Literal(0)))
            ));
        }

        #[test]
        fn read_partial_literal() {
            let mut reader = io::Cursor::new(b"\x01");
            assert!(matches!(
                Operation::default().advance(&mut reader),
                Err(OperationError::InsufficientInput(Command::Literal(2)))
            ));
        }

        #[test]
        fn read_partial_literal_to_completion() {
            let mut reader = io::Cursor::new(b"\xAB");
            assert!(matches!(
                Command::Literal(2).execute(&mut reader),
                Ok((0xab, Operation::Literal(1)))
            ));
        }

        #[test]
        fn read_partial_repeat() {
            let mut reader = io::Cursor::new(b"\xFF");
            assert!(matches!(
                Operation::default().advance(&mut reader),
                Err(OperationError::InsufficientInput(Command::Repeat(2)))
            ));
        }

        #[test]
        fn read_partial_repeat_to_completion() {
            let mut reader = io::Cursor::new(b"\xAB");
            assert!(matches!(
                Command::Repeat(2).execute(&mut reader),
                Ok((0xab, Operation::Repeat(0xab, 1)))
            ));
        }

        #[test]
        fn read_partial_escape() {
            let mut reader = io::Cursor::new(b"\x80");
            assert!(matches!(
                Operation::default().advance(&mut reader),
                Err(OperationError::InsufficientInput(Command::Escape))
            ));
        }

        #[test]
        fn read_partial_escape_to_completion() {
            let mut reader = io::Cursor::new(b"\x80");
            assert!(matches!(
                Command::Escape.execute(&mut reader),
                Ok((0x80, Operation::Literal(0) | Operation::Repeat(_, 0)))
            ));
        }
    }
}