viterbi 0.2.0

// Author: BolivarTech
// Version: 1.0.0
// Date: 2026-07-02
//! Hard-decision block Viterbi decoder (ACS + traceback).
//!
//! Orchestrates ACS + traceback and enforces the memory/length limits (R3, R11, R18).
pub(crate) mod acs;
pub(crate) mod traceback;

use crate::metric::{BranchMetric, HardHamming, SoftLlr};
use crate::packing::{pack_into, unpack_into, CodedBlock, DecodedBlock};
use crate::params::{CodeParams, ParamError};
use crate::trellis::Trellis;
use crate::MAX_SUPPORTED_INFO_BITS;
use acs::{acs_stage, SENTINEL};

/// Errors from [`ViterbiDecoder::new`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ConfigError {
    /// A [`ParamError`] surfaced while validating the code parameters or the state count `S`.
    Param(ParamError),
    /// `max_info_bits` was zero (a decoder must admit at least one info bit).
    MaxBlockZero,
    /// `max_info_bits` exceeded the hard cap [`MAX_SUPPORTED_INFO_BITS`] (R18).
    MaxBlockTooLarge {
        /// The hard cap.
        cap: usize,
        /// The requested value.
        got: usize,
    },
    /// A fallible scratch preallocation failed (out of memory) — reported, never a panic (R11).
    AllocationFailed,
}
/// Errors from decoding (boundary/allocation checks only — the ACS correction loop is infallible, R11).
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DecodeError {
    /// The implied payload length exceeded the decoder's configured `max_info_bits` (R18).
    InputTooLong {
        /// The configured cap.
        max_bits: usize,
        /// The rejected payload length in bits.
        got_bits: usize,
    },
    /// The coded input length was inconsistent (not a multiple of the code's `n`, outside the byte
    /// buffer, too short for the `m = k-1` tail, or `received.len()` inconsistent with `nbits`).
    LengthMismatch,
    /// The recovered-bytes output allocation failed (out of memory) — reported, never a panic (R11).
    AllocationFailed,
}
impl core::fmt::Display for ConfigError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "{self:?}")
    }
}
impl core::fmt::Display for DecodeError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "{self:?}")
    }
}
impl std::error::Error for ConfigError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            ConfigError::Param(e) => Some(e),
            _ => None,
        }
    }
}
impl std::error::Error for DecodeError {}

/// Hard-decision block Viterbi decoder over `S` states, generic over the branch metric `Me` (R7, R12).
///
/// All scratch (survivor, traceback, sample buffers) is preallocated in [`new`](Self::new), so the
/// ACS correction loop allocates nothing and cannot OOM-panic (R11, R18).
pub struct ViterbiDecoder<const S: usize, Me: BranchMetric> {
    trellis: Trellis<S>,
    max_info_bits: usize,
    survivors: Vec<u64>, // preallocated (max_info_bits + m) · ⌈S/64⌉ decision words (m = k-1)
    inputs: Vec<u8>,     // preallocated traceback scratch (max_info_bits + m)
    samples: Vec<Me::Sample>, // preallocated unpack scratch ((max_info_bits + m) * n, n = rate 1/n)
    wps: usize,          // decision words per stage = ⌈S/64⌉ (cached; computed once in new)
    terminal_metric: u32, // metric of state 0 after the last decode (diagnostic)
    _metric: core::marker::PhantomData<Me>,
}

/// First-class CCSDS hard-decision profile.
pub type CcsdsViterbiDecoder = ViterbiDecoder<64, HardHamming>;

/// First-class CCSDS soft-decision profile (K=7 R=1/2, LLR metric).
pub type CcsdsSoftDecoder = ViterbiDecoder<64, SoftLlr>;

impl<const S: usize, Me: BranchMetric> ViterbiDecoder<S, Me> {
    /// Build the decoder; validate params and `S`, preallocate all scratch via `try_reserve` (R11, R18).
    ///
    /// Never panics: `S != 2^(k-1)` is returned as `ConfigError::Param(StateCountMismatch)`.
    ///
    /// # Memory budget (generic K)
    /// The survivor buffer is `(max_info_bits + m)·⌈S/64⌉·8` bytes (`m = k-1`), so it scales with the
    /// constraint length: `8·(max_info_bits + 6)` at K=7 (`S=64`, 1 word/stage) and `32·(max_info_bits + 8)`
    /// at K=9 (`S=256`, 4 words/stage — 4× the K=7 cost). This footprint is bounded by `max_info_bits`
    /// (R18); a K=9 decoder with `max_info_bits = 1_000_000` reserves ≈ 32 MB. Every scratch reservation
    /// goes through `try_reserve`, so an over-large `S`/`max_info_bits` fails with
    /// [`ConfigError::AllocationFailed`] rather than an OOM abort.
    ///
    /// # Errors
    /// - [`ConfigError::Param`] — invalid/unsupported/catastrophic params, or `S != 2^(k-1)`.
    /// - [`ConfigError::MaxBlockZero`] / [`ConfigError::MaxBlockTooLarge`] — `max_info_bits` out of range.
    /// - [`ConfigError::AllocationFailed`] — a scratch preallocation failed.
    pub fn new(params: CodeParams, max_info_bits: usize) -> Result<Self, ConfigError> {
        params.validate().map_err(ConfigError::Param)?; // defense-in-depth: re-check the invariant (R19)
                                                        // Rate guard: only 1/2 and 1/3 (n ∈ {2, 3}) are supported by the engine. The constraint length
                                                        // `k` is NOT hard-rejected here — `Trellis::from_params` validates `S == 2^(k-1)`, so a K≠7 code
                                                        // is admitted iff its const state count `S` matches (K=7 stays the default; K=9 uses S=256).
        let n = params.generators.len();
        if !(2..=3).contains(&n) {
            return Err(ConfigError::Param(ParamError::RateUnsupported { n }));
        }
        if max_info_bits == 0 {
            return Err(ConfigError::MaxBlockZero);
        }
        if max_info_bits > MAX_SUPPORTED_INFO_BITS {
            return Err(ConfigError::MaxBlockTooLarge {
                cap: MAX_SUPPORTED_INFO_BITS,
                got: max_info_bits,
            });
        }
        let trellis = Trellis::from_params(&params).map_err(ConfigError::Param)?; // validates S; no panic
        let m = params.k as usize - 1; // memory = k-1 (runtime tail length; 6 for K=7, 8 for K=9)
        let stages = max_info_bits + m;
        let wps = S.div_ceil(64); // decision words per stage (scales with S — R12)
        let mut survivors = Vec::new();
        let mut inputs = Vec::new();
        let mut samples = Vec::new();
        // saturating_mul on the size products for defense-in-depth consistency with encode_bits (these
        // are already bounded: stages ≤ max_info_bits + m ≤ 1M+8, wps ≤ 4, n ≤ 3 → no overflow). The
        // sample scratch is sized `stages · n` (runtime rate 1/n), so rate-1/3 gets its 3-per-stage room.
        survivors
            .try_reserve_exact(stages.saturating_mul(wps))
            .map_err(|_| ConfigError::AllocationFailed)?;
        inputs
            .try_reserve_exact(stages)
            .map_err(|_| ConfigError::AllocationFailed)?;
        samples
            .try_reserve_exact(stages.saturating_mul(n))
            .map_err(|_| ConfigError::AllocationFailed)?;
        // `survivors` is indexed by stage, so it is resized (zero-filled) up front. `inputs` and
        // `samples` are reused push-buffers: kept at len 0 here and refilled within their reserved
        // capacity each decode (traceback / unpack), so they are intentionally not resized.
        survivors.resize(stages.saturating_mul(wps), 0);
        Ok(Self {
            trellis,
            max_info_bits,
            survivors,
            inputs,
            samples,
            wps,
            terminal_metric: 0,
            _metric: core::marker::PhantomData,
        })
    }

    /// Metric of the terminal state 0 after the last `decode` — a soft-confidence diagnostic
    /// (higher ⇒ the received block was farther from any valid terminated codeword). Never a pass/fail.
    /// Reflects the last `decode` that reached traceback; **unspecified** after a call that returned `Err`.
    #[must_use]
    pub fn last_terminal_metric(&self) -> u32 {
        self.terminal_metric
    }

    /// Decode `(nbits + m)` stages of `received` (`m = k-1` tail stages, `n` samples per stage, `n` = the
    /// code's rate 1/n) into the payload bits.
    ///
    /// The ACS correction loop allocates nothing (preallocated survivor buffer + stack arrays);
    /// only the recovered-bytes output uses a fallible `try_reserve` (R11).
    ///
    /// # Errors
    /// - [`DecodeError::InputTooLong`] — `nbits` exceeds the configured `max_info_bits`.
    /// - [`DecodeError::LengthMismatch`] — `received.len()` is inconsistent with `nbits`.
    /// - [`DecodeError::AllocationFailed`] — the recovered-bytes output allocation failed.
    pub fn decode(
        &mut self,
        received: &[Me::Sample],
        nbits: usize,
    ) -> Result<DecodedBlock, DecodeError> {
        if nbits > self.max_info_bits {
            return Err(DecodeError::InputTooLong {
                max_bits: self.max_info_bits,
                got_bits: nbits,
            });
        }
        let n = self.trellis.n; // runtime rate 1/n (2 or 3) — never the CCSDS-default const N
        let m = self.trellis.k as usize - 1; // memory = k-1 (runtime tail length; never the const M)
                                             // Explicit received-length validation with CHECKED arithmetic (never trust `received.len()`
                                             // implicitly): `(nbits + m) * n`. Both the `+m` add and the `*n` mul are checked so an
                                             // adversarial near-`usize::MAX` `nbits` yields `None` → `LengthMismatch`, never an overflow
                                             // panic. Done BEFORE the ACS loop so every per-stage slice is exactly `n` samples.
        let expected = nbits.checked_add(m).and_then(|s| s.checked_mul(n));
        if expected != Some(received.len()) {
            return Err(DecodeError::LengthMismatch);
        }
        let stages = nbits + m;
        let wps = self.wps; // cached ⌈S/64⌉ (computed once in `new`)
        let mut cur = [SENTINEL; S];
        cur[0] = 0;
        let mut next = [0u32; S];
        for stage in 0..stages {
            let slice = &received[stage * n..stage * n + n];
            acs_stage::<S, Me>(
                &cur,
                &mut next,
                &mut self.survivors[stage * wps..(stage + 1) * wps],
                &self.trellis,
                slice,
            );
            core::mem::swap(&mut cur, &mut next);
        }
        self.terminal_metric = cur[0]; // diagnostic: metric of the assumed terminal state 0
                                       // traceback into the preallocated scratch (disjoint field borrows; no allocation)
        traceback::traceback::<S>(
            &self.trellis,
            &self.survivors[..stages * wps],
            stages,
            &mut self.inputs,
        );
        // Pack the recovered info bits via the packing module (sole owner of the MSB-first
        // convention, R16). Allocation stays fallible: reserve exactly the needed capacity, then let
        // `pack_into` clear+resize into it — since the capacity already matches, no reallocation occurs.
        let mut bytes: Vec<u8> = Vec::new();
        bytes
            .try_reserve_exact(nbits.div_ceil(8))
            .map_err(|_| DecodeError::AllocationFailed)?;
        pack_into(&self.inputs[..nbits], nbits, &mut bytes);
        Ok(DecodedBlock { bytes, nbits })
    }
}

impl<const S: usize> ViterbiDecoder<S, HardHamming> {
    /// Hard-decision block decode — reads (by reference) the `CodedBlock` (symmetric with `encode`).
    ///
    /// **Validates `coded.nbits` fully before unpacking** (R16, no OOB): must be a multiple of the
    /// code's rate `n` (2 or 3), within the byte buffer, and long enough for the `m = k-1` tail. The
    /// rate is the code's own `n`, not the CCSDS-default const `N`, so a rate-1/3 block decodes here too.
    ///
    /// # Errors
    /// - [`DecodeError::LengthMismatch`] — `nbits` not a multiple of the code's `n`, outside the byte
    ///   buffer, or too short for the `m = k-1` tail.
    /// - [`DecodeError::InputTooLong`] — the implied payload exceeds `max_info_bits`.
    /// - [`DecodeError::AllocationFailed`] — the recovered-bytes output allocation failed.
    pub fn decode_block(&mut self, coded: &CodedBlock) -> Result<DecodedBlock, DecodeError> {
        let n = self.trellis.n; // runtime rate 1/n (2 or 3) — never the CCSDS-default const N
        if !coded.nbits.is_multiple_of(n) {
            return Err(DecodeError::LengthMismatch);
        }
        // `div_ceil` (no multiplication) is overflow-proof and provably correct: unpacking `coded.nbits`
        // bits needs `⌈nbits/8⌉` bytes, so reject when that exceeds the buffer. Avoids any saturating-mul
        // edge case on an attacker-controlled `bytes.len()`.
        if coded.nbits.div_ceil(8) > coded.bytes.len() {
            return Err(DecodeError::LengthMismatch);
        }
        let stages = coded.nbits / n;
        let m = self.trellis.k as usize - 1; // memory = k-1 (runtime tail length; 6 for K=7, 8 for K=9)
        if stages < m {
            return Err(DecodeError::LengthMismatch);
        }
        let nbits = stages - m;
        if nbits > self.max_info_bits {
            return Err(DecodeError::InputTooLong {
                max_bits: self.max_info_bits,
                got_bits: nbits,
            });
        }
        // Unpack into preallocated scratch via the packing module (sole owner, R16). The buffer is
        // moved out (borrow checker) and restored below; safe because `decode` never panics, so the
        // restore always runs. Even in a hypothetical unwind the empty scratch **self-heals** on the
        // next call (`unpack_into` refills it) — no correctness impact. No allocation (preallocated).
        let mut samples = core::mem::take(&mut self.samples);
        samples.clear();
        // The buffer is preallocated in `new` to `(max_info_bits + m)·n ≥ coded.nbits` (guaranteed by the
        // checks above), so this reserve is a no-op in practice; making it fallible removes any OOM-panic
        // path (R11) — an allocation failure returns `Err`, never aborts.
        if samples.try_reserve(coded.nbits).is_err() {
            self.samples = samples; // restore before bailing
            return Err(DecodeError::AllocationFailed);
        }
        unpack_into(&coded.bytes, coded.nbits, &mut samples);
        let result = self.decode(&samples, nbits);
        self.samples = samples; // restore the scratch buffer
        result
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::encoder::ViterbiEncoder;
    use crate::params::CodeParams;
    #[test]
    fn clean_round_trip_recovers_payload() {
        let enc = ViterbiEncoder::new(CodeParams::ccsds_r1_2()).unwrap();
        let mut dec = CcsdsViterbiDecoder::new(CodeParams::ccsds_r1_2(), 4096).unwrap();
        let raw = b"Viterbi works!";
        let coded = enc.encode(raw).unwrap();
        let out = dec.decode_block(&coded).unwrap();
        assert_eq!(&out.bytes, raw);
        assert_eq!(out.nbits, raw.len() * 8);
    }
    #[test]
    fn new_rejects_zero_and_oversized_max_bits() {
        assert_eq!(
            CcsdsViterbiDecoder::new(CodeParams::ccsds_r1_2(), 0)
                .err()
                .unwrap(),
            ConfigError::MaxBlockZero
        );
        assert!(matches!(
            CcsdsViterbiDecoder::new(CodeParams::ccsds_r1_2(), 2_000_000),
            Err(ConfigError::MaxBlockTooLarge { .. })
        ));
    }
    #[test]
    fn decode_rejects_input_longer_than_max() {
        let enc = ViterbiEncoder::new(CodeParams::ccsds_r1_2()).unwrap();
        let mut dec = CcsdsViterbiDecoder::new(CodeParams::ccsds_r1_2(), 8).unwrap();
        let coded = enc.encode(&[0u8; 4]).unwrap(); // 32 info bits > 8 cap
        assert!(matches!(
            dec.decode_block(&coded),
            Err(DecodeError::InputTooLong { .. })
        ));
    }
    #[test]
    fn decode_block_rejects_malformed_lengths() {
        use crate::packing::CodedBlock;
        let mut dec = CcsdsViterbiDecoder::new(CodeParams::ccsds_r1_2(), 4096).unwrap();
        // nbits not a multiple of N=2
        assert!(matches!(
            dec.decode_block(&CodedBlock {
                bytes: vec![0; 4],
                nbits: 15
            }),
            Err(DecodeError::LengthMismatch)
        ));
        // nbits exceeds the byte buffer (100 > 2*8)
        assert!(matches!(
            dec.decode_block(&CodedBlock {
                bytes: vec![0; 2],
                nbits: 100
            }),
            Err(DecodeError::LengthMismatch)
        ));
        // too short for the M=6 tail (stages = 8/2 = 4 < 6)
        assert!(matches!(
            dec.decode_block(&CodedBlock {
                bytes: vec![0; 1],
                nbits: 8
            }),
            Err(DecodeError::LengthMismatch)
        ));
    }
    #[test]
    fn new_accepts_rate_1_3_and_rejects_unsupported_rate_and_k() {
        // Rate 1/3 (n=3) is now SUPPORTED at S=64, K=7 (Task 6 relaxed the guard to n ∈ {2, 3}).
        assert!(ViterbiDecoder::<64, HardHamming>::new(CodeParams::ccsds_r1_3(), 64).is_ok());
        // n=1 (single monomial generator 0o100 = x^6, non-catastrophic) is outside the supported
        // rate set {1/2, 1/3} → RateUnsupported (checked before the S/trellis validation).
        let n1 = CodeParams::new(7, vec![0o100], vec![false]).unwrap();
        assert!(matches!(
            ViterbiDecoder::<64, HardHamming>::new(n1, 64),
            Err(ConfigError::Param(ParamError::RateUnsupported { n: 1 }))
        ));
        // K=3 with the K=7 const S=64 (should be 2^2=4): the hard `k != K` reject is GONE, so the
        // mismatch now surfaces as StateCountMismatch from `Trellis::from_params` — no panic.
        let k3 = CodeParams::new(3, vec![0b101, 0b111], vec![false, false]).unwrap();
        assert!(matches!(
            ViterbiDecoder::<64, HardHamming>::new(k3, 64),
            Err(ConfigError::Param(ParamError::StateCountMismatch {
                states: 64,
                k: 3
            }))
        ));
    }
    #[test]
    fn new_rejects_catastrophic_params_defensively() {
        // Defense-in-depth (release-active): even an in-crate struct literal bypassing `CodeParams::new`
        // (fields are pub(crate)) with catastrophic generators [0o12, 0o12] (gcd = 0o12, not a monomial)
        // is rejected by `new`'s `validate()` re-check — the engine never trusts params blindly.
        let bad = CodeParams {
            k: 7,
            generators: vec![0o12, 0o12],
            invert_outputs: vec![false, false],
        };
        assert!(matches!(
            CcsdsViterbiDecoder::new(bad, 64),
            Err(ConfigError::Param(ParamError::Catastrophic))
        ));
    }
}