melpe-rs 0.1.2

/// Bitstream packing and unpacking for 600 bps MELPe superframes
///
/// Packs a SuperFrameQuantized into a 6-byte (41-bit) wire format
/// and unpacks it back. Bit-level operations with a clear allocation table.

use crate::core_types::{
    NUM_LSF, SUPERFRAME_BITS_600, SUPERFRAME_BYTES_600,
    LSF_BITS_600, PITCH_BITS_600, GAIN_BITS_600,
    VOICING_BITS_600, JITTER_BITS_600, SPARE_BITS_600,
    SuperFrameQuantized,
};

/// Bit allocation table for the 600 bps superframe.
/// Each entry is (field_name, bit_count) in transmission order.
/// Total: 41 bits → 6 bytes (7 spare bits in last byte).
pub const BIT_ALLOC_600: &[(&str, usize)] = &[
    ("lsf0", 3), ("lsf1", 3), ("lsf2", 3), ("lsf3", 3),
    ("lsf4", 2), ("lsf5", 2), ("lsf6", 2), ("lsf7", 2),
    ("lsf8", 2), ("lsf9", 2),
    ("pitch", 6),
    ("gain", 5),
    ("voicing", 4),
    ("jitter", 1),
    ("spare", 1),
];

// ── Bitstream writer ──

/// Writes bits MSB-first into a byte buffer.
struct BitWriter<'a> {
    buf: &'a mut [u8],
    byte_pos: usize,
    bit_pos: u8, // 0..7, counts from MSB
}

impl<'a> BitWriter<'a> {
    fn new(buf: &'a mut [u8]) -> Self {
        for b in buf.iter_mut() {
            *b = 0;
        }
        Self {
            buf,
            byte_pos: 0,
            bit_pos: 0,
        }
    }

    /// Write `n_bits` (1..=8) from the low bits of `value`.
    fn write(&mut self, value: u8, n_bits: usize) {
        debug_assert!(n_bits >= 1 && n_bits <= 8);
        for i in (0..n_bits).rev() {
            let bit = (value >> i) & 1;
            if bit == 1 {
                self.buf[self.byte_pos] |= 1 << (7 - self.bit_pos);
            }
            self.bit_pos += 1;
            if self.bit_pos >= 8 {
                self.bit_pos = 0;
                self.byte_pos += 1;
            }
        }
    }

    /// Total bits written so far.
    fn bits_written(&self) -> usize {
        self.byte_pos * 8 + self.bit_pos as usize
    }
}

// ── Bitstream reader ──

/// Reads bits MSB-first from a byte buffer.
struct BitReader<'a> {
    buf: &'a [u8],
    byte_pos: usize,
    bit_pos: u8,
}

impl<'a> BitReader<'a> {
    fn new(buf: &'a [u8]) -> Self {
        Self {
            buf,
            byte_pos: 0,
            bit_pos: 0,
        }
    }

    /// Read `n_bits` (1..=8) and return them in the low bits.
    fn read(&mut self, n_bits: usize) -> u8 {
        debug_assert!(n_bits >= 1 && n_bits <= 8);
        let mut value: u8 = 0;
        for _ in 0..n_bits {
            value <<= 1;
            if self.byte_pos < self.buf.len() {
                let bit = (self.buf[self.byte_pos] >> (7 - self.bit_pos)) & 1;
                value |= bit;
            }
            self.bit_pos += 1;
            if self.bit_pos >= 8 {
                self.bit_pos = 0;
                self.byte_pos += 1;
            }
        }
        value
    }

    /// Total bits read so far.
    fn bits_read(&self) -> usize {
        self.byte_pos * 8 + self.bit_pos as usize
    }
}

// ── Pack / Unpack ──

/// Pack a SuperFrameQuantized into a byte buffer (6 bytes for 600 bps).
/// Returns the number of bits written (should be SUPERFRAME_BITS_600).
pub fn pack_superframe(sq: &SuperFrameQuantized, out: &mut [u8; SUPERFRAME_BYTES_600]) -> usize {
    let mut w = BitWriter::new(out);

    // LSFs: 10 coefficients with variable bit widths
    for i in 0..NUM_LSF {
        w.write(sq.lsf_indices[i], LSF_BITS_600[i]);
    }

    // Pitch: 6 bits
    w.write(sq.pitch_index, PITCH_BITS_600);

    // Gain: 5 bits
    w.write(sq.gain_index, GAIN_BITS_600);

    // Voicing: 4 bits
    w.write(sq.voicing_bits, VOICING_BITS_600);

    // Jitter: 1 bit
    w.write(sq.jitter_bit, JITTER_BITS_600);

    // Spare: 1 bit (write 0)
    w.write(0, SPARE_BITS_600);

    w.bits_written()
}

/// Unpack a byte buffer into a SuperFrameQuantized.
/// Returns the struct and the number of bits consumed.
pub fn unpack_superframe(data: &[u8; SUPERFRAME_BYTES_600]) -> (SuperFrameQuantized, usize) {
    let mut r = BitReader::new(data);
    let mut sq = SuperFrameQuantized::default();

    // LSFs
    for i in 0..NUM_LSF {
        sq.lsf_indices[i] = r.read(LSF_BITS_600[i]);
    }

    // Pitch
    sq.pitch_index = r.read(PITCH_BITS_600);

    // Gain
    sq.gain_index = r.read(GAIN_BITS_600);

    // Voicing
    sq.voicing_bits = r.read(VOICING_BITS_600);

    // Jitter
    sq.jitter_bit = r.read(JITTER_BITS_600);

    // Skip spare
    let _spare = r.read(SPARE_BITS_600);

    (sq, r.bits_read())
}

/// Verify that the bit allocation table sums correctly.
pub fn verify_bit_budget() -> bool {
    let total: usize = BIT_ALLOC_600.iter().map(|(_, bits)| *bits).sum();
    total == SUPERFRAME_BITS_600
}

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

    // ── BitWriter / BitReader unit tests ──

    #[test]
    fn test_bitwriter_single_bits() {
        let mut buf = [0u8; 2];
        let bits_written;
        {
            let mut w = BitWriter::new(&mut buf);
            w.write(1, 1); // bit 7
            w.write(0, 1); // bit 6
            w.write(1, 1); // bit 5
            w.write(1, 1); // bit 4
            w.write(0, 1); // bit 3
            w.write(0, 1); // bit 2
            w.write(1, 1); // bit 1
            w.write(0, 1); // bit 0
            bits_written = w.bits_written();
        }
        // Should be 0b10110010 = 0xB2
        assert_eq!(buf[0], 0xB2);
        assert_eq!(bits_written, 8);
    }

    #[test]
    fn test_bitwriter_multi_bit() {
        let mut buf = [0u8; 2];
        let bits_written;
        {
            let mut w = BitWriter::new(&mut buf);
            w.write(0b101, 3); // bits 7-5
            w.write(0b11010, 5); // bits 4-0
            bits_written = w.bits_written();
        }
        // byte 0: 10111010 = 0xBA
        assert_eq!(buf[0], 0xBA);
        assert_eq!(bits_written, 8);
    }

    #[test]
    fn test_bitwriter_crosses_byte() {
        let mut buf = [0u8; 2];
        let bits_written;
        {
            let mut w = BitWriter::new(&mut buf);
            w.write(0b111, 3);    // byte 0 bits 7-5
            w.write(0b101010, 6); // byte 0 bits 4-0 + byte 1 bit 7
            bits_written = w.bits_written();
        }
        // byte 0: 111_10101 = 0xF5
        // byte 1: 0_0000000 = 0x00
        assert_eq!(buf[0], 0xF5);
        assert_eq!(buf[1], 0x00);
        assert_eq!(bits_written, 9);
    }

    #[test]
    fn test_bitreader_single_bits() {
        let buf = [0xB2u8, 0x00]; // 10110010
        let mut r = BitReader::new(&buf);
        assert_eq!(r.read(1), 1);
        assert_eq!(r.read(1), 0);
        assert_eq!(r.read(1), 1);
        assert_eq!(r.read(1), 1);
        assert_eq!(r.read(1), 0);
        assert_eq!(r.read(1), 0);
        assert_eq!(r.read(1), 1);
        assert_eq!(r.read(1), 0);
        assert_eq!(r.bits_read(), 8);
    }

    #[test]
    fn test_bitreader_multi_bit() {
        let buf = [0xBA]; // 10111010
        let mut r = BitReader::new(&buf);
        assert_eq!(r.read(3), 0b101); // 101
        assert_eq!(r.read(5), 0b11010); // 11010
    }

    #[test]
    fn test_writer_reader_roundtrip() {
        let mut buf = [0u8; 4];
        let mut w = BitWriter::new(&mut buf);
        w.write(0b110, 3);
        w.write(0b01011, 5);
        w.write(0b1111, 4);
        w.write(0b00, 2);
        w.write(0b101101, 6);

        let mut r = BitReader::new(&buf);
        assert_eq!(r.read(3), 0b110);
        assert_eq!(r.read(5), 0b01011);
        assert_eq!(r.read(4), 0b1111);
        assert_eq!(r.read(2), 0b00);
        assert_eq!(r.read(6), 0b101101);
    }

    // ── Bit budget verification ──

    #[test]
    fn test_bit_budget_table() {
        assert!(verify_bit_budget(), "Bit allocation table doesn't sum to {}", SUPERFRAME_BITS_600);
    }

    #[test]
    fn test_bit_alloc_field_count() {
        // 10 LSFs + pitch + gain + voicing + jitter + spare = 15 entries
        assert_eq!(BIT_ALLOC_600.len(), 15);
    }

    // ── Pack/Unpack roundtrips ──

    #[test]
    fn test_pack_unpack_zeros() {
        let sq = SuperFrameQuantized::default();
        let mut buf = [0u8; SUPERFRAME_BYTES_600];
        let bits = pack_superframe(&sq, &mut buf);
        assert_eq!(bits, SUPERFRAME_BITS_600);

        let (recovered, bits_read) = unpack_superframe(&buf);
        assert_eq!(bits_read, SUPERFRAME_BITS_600);
        assert_eq!(recovered, sq);
    }

    #[test]
    fn test_pack_unpack_max_values() {
        let sq = SuperFrameQuantized {
            lsf_indices: [
                (1 << LSF_BITS_600[0]) - 1,
                (1 << LSF_BITS_600[1]) - 1,
                (1 << LSF_BITS_600[2]) - 1,
                (1 << LSF_BITS_600[3]) - 1,
                (1 << LSF_BITS_600[4]) - 1,
                (1 << LSF_BITS_600[5]) - 1,
                (1 << LSF_BITS_600[6]) - 1,
                (1 << LSF_BITS_600[7]) - 1,
                (1 << LSF_BITS_600[8]) - 1,
                (1 << LSF_BITS_600[9]) - 1,
            ],
            pitch_index: (1 << PITCH_BITS_600) - 1,
            gain_index: (1 << GAIN_BITS_600) - 1,
            voicing_bits: (1 << VOICING_BITS_600) - 1,
            jitter_bit: 1,
        };
        let mut buf = [0u8; SUPERFRAME_BYTES_600];
        pack_superframe(&sq, &mut buf);

        let (recovered, _) = unpack_superframe(&buf);
        assert_eq!(recovered, sq);
    }

    #[test]
    fn test_pack_unpack_mixed_values() {
        let sq = SuperFrameQuantized {
            lsf_indices: [5, 3, 7, 1, 2, 0, 3, 1, 2, 3],
            pitch_index: 42,
            gain_index: 17,
            voicing_bits: 0b1011,
            jitter_bit: 1,
        };
        let mut buf = [0u8; SUPERFRAME_BYTES_600];
        pack_superframe(&sq, &mut buf);

        let (recovered, _) = unpack_superframe(&buf);
        assert_eq!(recovered.lsf_indices, sq.lsf_indices);
        assert_eq!(recovered.pitch_index, sq.pitch_index);
        assert_eq!(recovered.gain_index, sq.gain_index);
        assert_eq!(recovered.voicing_bits, sq.voicing_bits);
        assert_eq!(recovered.jitter_bit, sq.jitter_bit);
    }

    #[test]
    fn test_pack_unpack_exhaustive_lsf() {
        // Test every valid LSF index combination for the first 4 LSFs (3-bit each)
        for i0 in 0..8u8 {
            for i1 in 0..8u8 {
                let sq = SuperFrameQuantized {
                    lsf_indices: [i0, i1, 4, 2, 1, 3, 0, 2, 1, 3],
                    pitch_index: 30,
                    gain_index: 15,
                    voicing_bits: 5,
                    jitter_bit: 0,
                };
                let mut buf = [0u8; SUPERFRAME_BYTES_600];
                pack_superframe(&sq, &mut buf);
                let (recovered, _) = unpack_superframe(&buf);
                assert_eq!(recovered, sq, "Failed for lsf[0]={}, lsf[1]={}", i0, i1);
            }
        }
    }

    #[test]
    fn test_pack_bit_count() {
        let sq = SuperFrameQuantized::default();
        let mut buf = [0u8; SUPERFRAME_BYTES_600];
        let bits = pack_superframe(&sq, &mut buf);
        assert_eq!(bits, 41);
    }

    #[test]
    fn test_buffer_size() {
        // 41 bits needs exactly 6 bytes
        assert_eq!(SUPERFRAME_BYTES_600, 6);
        assert_eq!((SUPERFRAME_BITS_600 + 7) / 8, 6);
    }

    #[test]
    fn test_spare_bits_zero() {
        // Spare bit should be 0 after packing
        let sq = SuperFrameQuantized {
            lsf_indices: [7, 7, 7, 7, 3, 3, 3, 3, 3, 3],
            pitch_index: 63,
            gain_index: 31,
            voicing_bits: 15,
            jitter_bit: 1,
        };
        let mut buf = [0u8; SUPERFRAME_BYTES_600];
        pack_superframe(&sq, &mut buf);

        // Read back and verify the spare bit is 0
        let mut r = BitReader::new(&buf);
        // Skip all fields: 24 + 6 + 5 + 4 + 1 = 40 bits
        for _ in 0..40 {
            r.read(1);
        }
        let spare = r.read(1);
        assert_eq!(spare, 0, "Spare bit should be 0");
    }

    #[test]
    fn test_pack_deterministic() {
        let sq = SuperFrameQuantized {
            lsf_indices: [2, 5, 1, 6, 3, 0, 2, 1, 3, 1],
            pitch_index: 33,
            gain_index: 20,
            voicing_bits: 0b0110,
            jitter_bit: 0,
        };

        let mut buf1 = [0u8; SUPERFRAME_BYTES_600];
        let mut buf2 = [0u8; SUPERFRAME_BYTES_600];
        pack_superframe(&sq, &mut buf1);
        pack_superframe(&sq, &mut buf2);
        assert_eq!(buf1, buf2, "Packing must be deterministic");
    }
}