ldpc.rs

use arrayvec::ArrayVec;
use labrador_ldpc::LDPCCode;

use crate::{buffer::Buffer, error::EncodeError};

macro_rules! enc_chunk {
    ($d:ident, $i:ident, $t:ident, $n:literal) => {
        ::paste::paste! {
        let mut codeword = [0; $n * 2];
        codeword[0..$n].copy_from_slice(&$d[$i..($i + $n)]);
        labrador_ldpc::LDPCCode::[<$t>].encode(&mut codeword);
        $d.try_extend_from_slice(&codeword[$n..])
            .map_err(|_| EncodeError::CatsOverflow)?;

            $i += $n;
        }
    };
}

macro_rules! dec_chunk {
    ($d:ident, $p:ident, $w:ident, $t:ident, $n:literal) => {
        ::paste::paste! {
            let code_data = &mut $d[..$n];
            let code_parity = &$p.get_mut(..$n)?;

            let mut input = [0; $n * 2];
            input[..$n].copy_from_slice(code_data);
            input[$n..].copy_from_slice(code_parity);
            const CODE: LDPCCode = LDPCCode::[<$t>];
            let mut out = [0; CODE.output_len()];
            CODE.decode_bf(&input, &mut out, &mut $w[..CODE.decode_bf_working_len()], 16);
            code_data.copy_from_slice(&out[..$n]);

            $d = &mut $d[$n..];
            $p = &mut $p[$n..];
        }
    };
}

// On failure this still modifies the data array!
pub(crate) fn encode<const N: usize>(data: &mut Buffer<N>) -> Result<(), EncodeError> {
    let mut i = 0;
    let n = data.len();

    loop {
        match n - i {
            512.. => {
                enc_chunk!(data, i, TM8192, 512);
            }

            128.. => {
                enc_chunk!(data, i, TM2048, 128);
            }

            32.. => {
                enc_chunk!(data, i, TC512, 32);
            }

            16.. => {
                enc_chunk!(data, i, TC256, 16);
            }

            8.. => {
                enc_chunk!(data, i, TC128, 8);
            }

            0 => break,

            _ => {
                let mut codeword = [0xAA; 16];
                codeword[0..(n - i)].copy_from_slice(&data[i..n]);
                LDPCCode::TC128.encode(&mut codeword);
                data.try_extend_from_slice(&codeword[8..])
                    .map_err(|_| EncodeError::CatsOverflow)?;

                i = n;
            }
        }
    }

    data.try_extend_from_slice(&u16::try_from(n).unwrap().to_le_bytes())
        .map_err(|_| EncodeError::CatsOverflow)?;

    Ok(())
}

pub(crate) fn decode<const N: usize>(data_av: &mut ArrayVec<u8, N>) -> Option<()> {
    if data_av.len() < 2 {
        return None;
    }

    let len = [data_av[data_av.len() - 2], data_av[data_av.len() - 1]];
    let len = u16::from_le_bytes(len).into();

    if len >= data_av.len() {
        return None;
    }

    let (mut data, parity) = data_av.split_at_mut(len);
    let mut parity = parity.get_mut(..parity.len().checked_sub(2)?)?;

    let mut working = [0; LDPCCode::TM8192.decode_bf_working_len()];

    loop {
        match data.len() {
            512.. => {
                dec_chunk!(data, parity, working, TM8192, 512);
            }

            128.. => {
                dec_chunk!(data, parity, working, TM2048, 128);
            }

            32.. => {
                dec_chunk!(data, parity, working, TC512, 32);
            }

            16.. => {
                dec_chunk!(data, parity, working, TC256, 16);
            }

            8.. => {
                dec_chunk!(data, parity, working, TC128, 8);
            }

            0 => break,

            _ => {
                let mut code_data = [0xAA; 8];
                code_data[..data.len()].copy_from_slice(data);
                let code_parity = &parity.get_mut(..8)?;

                let mut input = [0; 16];
                input[..8].copy_from_slice(&code_data);
                input[8..].copy_from_slice(code_parity);
                let mut out = [0; LDPCCode::TC128.output_len()];
                LDPCCode::TC128.decode_bf(
                    &input,
                    &mut out,
                    &mut working[..LDPCCode::TC128.decode_bf_working_len()],
                    16,
                );
                data.copy_from_slice(&out[..data.len()]);

                data = &mut data[..0];
                parity = &mut parity[..0];
            }
        }
    }

    // remove the parity data
    data_av.truncate(len);

    Some(())
}

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

    #[test]
    fn len_test() {
        // from the example in the docs
        let mut data: ArrayVec<u8, 8191> = ArrayVec::new();

        for _ in 0..41 {
            data.extend(
                b"Example packet data  wueirpqwerwrywqoeiruy29346129384761"
                    .iter()
                    .cloned(),
            );
        }
        data.extend(
            b"Example packet data  wueirpqwerwrywqoeiru346129384761"
                .iter()
                .cloned(),
        );
        assert_eq!(2349, data.len());

        encode(&mut data).unwrap();

        assert_eq!(4703, data.len());
    }

    #[test]
    fn basic_encode_decode_short() {
        let mut data: ArrayVec<u8, 32> = ArrayVec::new();
        data.try_extend_from_slice(b"Hello world!").unwrap();
        let orig = data.clone();

        encode(&mut data).unwrap();

        decode(&mut data).unwrap();

        assert_eq!(orig, data);
    }

    #[test]
    fn basic_encode_decode() {
        let mut data: ArrayVec<u8, 8191> = ArrayVec::new();
        for _ in 0..50 {
            data.extend(
                b"This is a test packet. jsalksjd093809324JASLD:LKD*#$)(*#@)"
                    .iter()
                    .cloned(),
            );
        }
        let orig = data.clone();

        encode(&mut data).unwrap();
        assert_ne!(orig, data);

        decode(&mut data).unwrap();

        assert_eq!(orig, data);
    }

    #[test]
    fn encode_decode_with_bit_flips() {
        let mut data: ArrayVec<u8, 8191> = ArrayVec::new();
        for _ in 0..50 {
            data.extend(
                b"jsalksjd093809324JASLD:LKD*#$)(*#@) Another test packet"
                    .iter()
                    .cloned(),
            );
        }
        let orig = data.clone();

        encode(&mut data).unwrap();
        assert_ne!(orig, data);

        data[234] ^= 0x55;
        data[0] ^= 0xAA;
        data[999] ^= 0x43;

        decode(&mut data).unwrap();

        assert_eq!(orig, data);
    }
}