viterbi 0.1.0

// Author: BolivarTech
// Version: 1.0.0
// Date: 2026-07-02
//! Convolutional encoder with zero-tail termination (R1, R2).
use crate::packing::CodedBlock;
use crate::params::{CodeParams, ParamError};
use crate::trellis::{next_state, output_symbol};
use crate::{K, M, MAX_SUPPORTED_INFO_BITS, N};

/// Errors from encoding.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum EncodeError {
    /// `nbits` exceeds the bits available in `data` (`8 * data.len()`).
    NbitsExceedsData {
        /// Requested information-bit count.
        nbits: usize,
        /// Bits actually available in the input slice.
        available: usize,
    },
    /// Payload exceeds `MAX_SUPPORTED_INFO_BITS` (symmetric with the decoder cap; R18).
    PayloadTooLarge {
        /// The hard cap (`MAX_SUPPORTED_INFO_BITS`).
        cap: usize,
        /// The requested information-bit count.
        got: usize,
    },
    /// Output buffer allocation failed (fallible `try_reserve`; never an abort — R11).
    AllocationFailed,
}
impl core::fmt::Display for EncodeError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "{self:?}")
    }
}
impl std::error::Error for EncodeError {}

/// CCSDS K=7 R=1/2 convolutional encoder.
///
/// Encodes an MSB-first bit payload with the CCSDS convolutional code and appends
/// `M = 6` zero tail bits so encoding ends in state 0 (R1, R2). The coded output is
/// byte-packed MSB-first with its exact bit length (R16).
///
/// # Examples
/// ```
/// use viterbi::encoder::ViterbiEncoder;
/// use viterbi::params::CodeParams;
/// let enc = ViterbiEncoder::new(CodeParams::ccsds_r1_2()).unwrap();
/// let coded = enc.encode(&[0xFF]).unwrap();
/// assert_eq!(coded.nbits, (8 + 6) * 2); // 8 info bits + 6 tail, rate 1/2
/// ```
pub struct ViterbiEncoder {
    params: CodeParams,
}

impl ViterbiEncoder {
    /// Build from validated params. Only K=7 rate-1/2 is supported in this increment.
    ///
    /// # Errors
    /// Returns [`ParamError`] if the params fail the full invariant (R19), or if `K`
    /// (`KUnsupported`) or the rate (`RateUnsupported`) is outside this increment's support.
    pub fn new(params: CodeParams) -> Result<Self, ParamError> {
        params.validate()?; // defense-in-depth: re-check the full invariant (release-active, R19)
        if params.k != K {
            return Err(ParamError::KUnsupported { k: params.k });
        }
        if params.generators.len() != N {
            return Err(ParamError::RateUnsupported {
                n: params.generators.len(),
            });
        }
        Ok(Self { params })
    }

    /// Encode every bit of `data` (nbits = 8·len), MSB-first. Fallible only on allocation.
    ///
    /// # Errors
    /// Returns [`EncodeError::PayloadTooLarge`] if the payload exceeds
    /// `MAX_SUPPORTED_INFO_BITS`, or [`EncodeError::AllocationFailed`] on OOM.
    pub fn encode(&self, data: &[u8]) -> Result<CodedBlock, EncodeError> {
        // saturating_mul: on a hypothetical `len > usize::MAX/8` the product saturates instead of
        // wrapping, so an oversized input hits the PayloadTooLarge cap below — never a silent truncation.
        self.encode_bits(data, data.len().saturating_mul(8))
    }

    /// Encode the first `nbits` bits of `data`. Payloads over `MAX_SUPPORTED_INFO_BITS` are rejected
    /// with `PayloadTooLarge` (a hard cap symmetric with the decoder). Below that, it writes MSB-first
    /// straight into packed bytes (no 8×-larger intermediate) reserved via `try_reserve` — allocation
    /// **cannot abort** (R11): OOM returns `EncodeError::AllocationFailed`.
    ///
    /// # Errors
    /// Returns [`EncodeError::NbitsExceedsData`] if `nbits > 8 * data.len()`,
    /// [`EncodeError::PayloadTooLarge`] if `nbits > MAX_SUPPORTED_INFO_BITS`, or
    /// [`EncodeError::AllocationFailed`] on OOM.
    pub fn encode_bits(&self, data: &[u8], nbits: usize) -> Result<CodedBlock, EncodeError> {
        let available = data.len().saturating_mul(8); // saturating: no usize overflow on huge slices
        if nbits > available {
            return Err(EncodeError::NbitsExceedsData { nbits, available });
        }
        if nbits > MAX_SUPPORTED_INFO_BITS {
            return Err(EncodeError::PayloadTooLarge {
                cap: MAX_SUPPORTED_INFO_BITS,
                got: nbits,
            });
        }
        let ncoded = (nbits + M as usize) * N;
        let mut bytes: Vec<u8> = Vec::new();
        bytes
            .try_reserve_exact(ncoded.div_ceil(8))
            .map_err(|_| EncodeError::AllocationFailed)?;
        bytes.resize(ncoded.div_ceil(8), 0);
        let mut state = 0usize;
        let mut oi = 0usize; // output bit index (MSB-first, matches `packing`)
        let emit = |bytes: &mut [u8], oi: &mut usize, sym: u32| {
            for j in 0..N {
                if (sym >> (N - 1 - j)) & 1 == 1 {
                    crate::packing::set_bit(bytes, *oi);
                }
                *oi += 1;
            }
        };
        for i in 0..nbits {
            let bit = crate::packing::get_bit(data, i); // MSB-first read via the packing sole-owner
            emit(&mut bytes, &mut oi, output_symbol(state, bit, &self.params));
            state = next_state(state, bit, self.params.k);
        }
        for _ in 0..M {
            // zero-tail flush
            emit(&mut bytes, &mut oi, output_symbol(state, 0, &self.params));
            state = next_state(state, 0, self.params.k);
        }
        Ok(CodedBlock {
            bytes,
            nbits: ncoded,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::params::CodeParams;
    #[test]
    fn impulse_response_matches_traced_ccsds() {
        // input = single bit 1 (MSB of 0x80), 6-bit zero-tail → 7 stages × 2 = 14 coded bits.
        let enc = ViterbiEncoder::new(CodeParams::ccsds_r1_2()).unwrap();
        let cb = enc.encode_bits(&[0x80], 1).unwrap();
        assert_eq!(cb.nbits, 14);
        assert_eq!(cb.bytes, vec![0xBA, 0x48]);
    }
    #[test]
    fn encode_appends_zero_tail_length() {
        let enc = ViterbiEncoder::new(CodeParams::ccsds_r1_2()).unwrap();
        let cb = enc.encode(&[0xFF]).unwrap(); // 8 info bits + 6 tail = 14 stages × 2 = 28 coded bits
        assert_eq!(cb.nbits, (8 + 6) * 2);
    }
    #[test]
    fn encode_bits_rejects_overlong_nbits() {
        let enc = ViterbiEncoder::new(CodeParams::ccsds_r1_2()).unwrap();
        assert_eq!(
            enc.encode_bits(&[0x00], 9),
            Err(EncodeError::NbitsExceedsData {
                nbits: 9,
                available: 8
            })
        );
    }
    #[test]
    fn encode_rejects_oversized_payload() {
        // hard cap symmetric with the decoder: > MAX_SUPPORTED_INFO_BITS ⇒ PayloadTooLarge (no unbounded alloc)
        let enc = ViterbiEncoder::new(CodeParams::ccsds_r1_2()).unwrap();
        let big = vec![0u8; crate::MAX_SUPPORTED_INFO_BITS / 8 + 1]; // 8·(cap/8 + 1) > cap info bits
        assert!(matches!(
            enc.encode(&big),
            Err(EncodeError::PayloadTooLarge { .. })
        ));
    }
}