viterbi 0.2.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::MAX_SUPPORTED_INFO_BITS;

/// 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 convolutional encoder over a runtime rate `1/n` (`n ∈ {2, 3}`).
///
/// Encodes an MSB-first bit payload with the CCSDS convolutional code and appends
/// `m = k - 1` zero tail bits (`6` for the K=7 default, `8` for K=9) so encoding
/// ends in state 0 (R1, R2). The tail length is derived from the code's own `k`,
/// not the crate default `M`, so a generic-K code terminates correctly. Each input
/// bit emits `n = generators().len()` output bits (C1 first, MSB-first), so both the
/// rate-1/2 (`ccsds_r1_2`) and rate-1/3 (`ccsds_r1_3`) mother codes are supported.
/// 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. Supports any `K ≤ 9` at rate `1/n` for `n ∈ {2, 3}`.
    ///
    /// The constraint length is **not** hard-restricted to K=7: `validate()` bounds
    /// `k ∈ 1..=9` (R19) and the zero-tail length is derived from `k` at encode time,
    /// so K=9 (and other `K ≤ 9`) codes encode correctly.
    ///
    /// # Errors
    /// Returns [`ParamError`] if the params fail the full invariant (R19) — including
    /// `k` outside `1..=9` (`KOutOfRange`) — or if the rate (`RateUnsupported`) is not
    /// `1/2` or `1/3` (`n ∉ {2, 3}`).
    pub fn new(params: CodeParams) -> Result<Self, ParamError> {
        params.validate()?; // defense-in-depth: re-check the full invariant (release-active, R19)
        let n = params.generators.len();
        if !(2..=3).contains(&n) {
            return Err(ParamError::RateUnsupported { n });
        }
        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, on the
    /// unreachable overflow arm, if the coded length `(nbits + m) * n` would overflow
    /// `usize`), 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 n = self.params.generators.len(); // runtime rate 1/n (n ∈ {2, 3})
        let m = self.params.k as usize - 1; // memory = k-1 (runtime tail length; 6 for K=7, 8 for K=9)
                                            // Checked `(nbits + m) * n` — uniform with decode/puncture/depuncture (crate
                                            // invariant: no raw length arithmetic). The bound above already keeps this in
                                            // range for supported inputs, so the overflow arm is unreachable; it maps to
                                            // `PayloadTooLarge` (a length-overflow rejection) rather than wrapping.
        let ncoded = nbits
            .checked_add(m)
            .and_then(|stages| stages.checked_mul(n))
            .ok_or(EncodeError::PayloadTooLarge {
                cap: MAX_SUPPORTED_INFO_BITS,
                got: nbits,
            })?;
        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 (m = k-1 stages force the register back to state 0)
            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 { .. })
        ));
    }
}