tx2-iff 0.1.0

//! Simple compression utilities for integer data
//!
//! Implements Run-Length Encoding (RLE) combined with VByte and ZigZag encoding
//! for efficient storage of sparse integer data (like wavelet coefficients).

use std::io::{Cursor, Read, Write};
use crate::error::{IffError, Result};

/// ZigZag encode a signed integer to unsigned
/// Maps small signed integers to small unsigned integers
/// 0 -> 0, -1 -> 1, 1 -> 2, -2 -> 3, ...
pub fn zigzag_encode(n: i32) -> u32 {
    ((n << 1) ^ (n >> 31)) as u32
}

/// ZigZag decode an unsigned integer to signed
pub fn zigzag_decode(n: u32) -> i32 {
    ((n >> 1) as i32) ^ -((n & 1) as i32)
}

/// Encode a u32 using VByte (variable length) encoding
/// Returns the number of bytes written
pub fn encode_vbyte<W: Write>(writer: &mut W, mut n: u32) -> Result<usize> {
    let mut bytes_written = 0;
    loop {
        let mut byte = (n & 0x7F) as u8;
        n >>= 7;
        if n != 0 {
            byte |= 0x80;
            writer.write_all(&[byte])?;
            bytes_written += 1;
        } else {
            writer.write_all(&[byte])?;
            bytes_written += 1;
            break;
        }
    }
    Ok(bytes_written)
}

/// Decode a u32 using VByte encoding
pub fn decode_vbyte<R: Read>(reader: &mut R) -> Result<u32> {
    let mut n = 0u32;
    let mut shift = 0;
    let mut byte = [0u8; 1];

    loop {
        reader.read_exact(&mut byte).map_err(|_| IffError::InsufficientData { expected: 1, got: 0 })?;
        let b = byte[0];
        n |= ((b & 0x7F) as u32) << shift;
        if (b & 0x80) == 0 {
            break;
        }
        shift += 7;
        if shift >= 35 {
            return Err(IffError::Other("VByte overflow".to_string()));
        }
    }
    Ok(n)
}

/// Compress a slice of i32 using RLE + ZigZag + VByte
/// Format:
/// - 0 (VByte 0x00) -> Run of zeros. Next VByte is count.
/// - >0 (VByte) -> Literal value (ZigZag encoded).
pub fn compress_rle(data: &[i32]) -> Result<Vec<u8>> {
    let mut output = Vec::new();
    let mut i = 0;
    while i < data.len() {
        if data[i] == 0 {
            // Run of zeros
            let mut count = 0;
            while i < data.len() && data[i] == 0 {
                count += 1;
                i += 1;
            }
            // Write 0 marker
            encode_vbyte(&mut output, 0)?;
            // Write count
            encode_vbyte(&mut output, count)?;
        } else {
            // Literal value
            let zz = zigzag_encode(data[i]);
            encode_vbyte(&mut output, zz)?;
            i += 1;
        }
    }
    Ok(output)
}

/// Decompress RLE data to Vec<i32>
/// `expected_len` is the expected number of integers (optional check)
pub fn decompress_rle(data: &[u8], expected_len: Option<usize>) -> Result<Vec<i32>> {
    let mut reader = Cursor::new(data);
    let mut output = Vec::with_capacity(expected_len.unwrap_or(1024));

    while reader.position() < data.len() as u64 {
        let val = decode_vbyte(&mut reader)?;
        if val == 0 {
            // Run of zeros
            let count = decode_vbyte(&mut reader)?;
            for _ in 0..count {
                output.push(0);
            }
        } else {
            // Literal value
            let n = zigzag_decode(val);
            output.push(n);
        }
    }

    if let Some(len) = expected_len {
        if output.len() != len {
             // It's possible the RLE stream ended early or went over if corrupted
             // But for now we just return what we have, or maybe pad with zeros if short?
             // Let's be strict for now.
             if output.len() < len {
                 // Pad with zeros if we are short (implicit zeros at end)
                 output.resize(len, 0);
             } else if output.len() > len {
                 // Truncate? Or error?
                 // Error is safer
                 return Err(IffError::Other(format!(
                     "Decompression length mismatch: expected {}, got {}",
                     len, output.len()
                 )));
             }
        }
    }

    Ok(output)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_zigzag() {
        assert_eq!(zigzag_encode(0), 0);
        assert_eq!(zigzag_encode(-1), 1);
        assert_eq!(zigzag_encode(1), 2);
        assert_eq!(zigzag_encode(-2), 3);
        assert_eq!(zigzag_encode(i32::MAX), 0xFFFFFFFE);
        assert_eq!(zigzag_encode(i32::MIN), 0xFFFFFFFF);

        assert_eq!(zigzag_decode(0), 0);
        assert_eq!(zigzag_decode(1), -1);
        assert_eq!(zigzag_decode(2), 1);
        assert_eq!(zigzag_decode(3), -2);
    }

    #[test]
    fn test_vbyte() {
        let mut buf = Vec::new();
        encode_vbyte(&mut buf, 0).unwrap();
        encode_vbyte(&mut buf, 127).unwrap();
        encode_vbyte(&mut buf, 128).unwrap();
        encode_vbyte(&mut buf, 300).unwrap();

        let mut reader = Cursor::new(buf);
        assert_eq!(decode_vbyte(&mut reader).unwrap(), 0);
        assert_eq!(decode_vbyte(&mut reader).unwrap(), 127);
        assert_eq!(decode_vbyte(&mut reader).unwrap(), 128);
        assert_eq!(decode_vbyte(&mut reader).unwrap(), 300);
    }

    #[test]
    fn test_rle() {
        let data = vec![0, 0, 0, 5, -3, 0, 0, 10, 0];
        let compressed = compress_rle(&data).unwrap();
        let decompressed = decompress_rle(&compressed, Some(data.len())).unwrap();
        assert_eq!(data, decompressed);
    }
}