viterbi 0.0.1

// Author: BolivarTech
// Version: 1.0.0
// Date: 2026-07-02
//! Precomputed forward + predecessor trellis tables, const-generic over states (R12).
use crate::params::CodeParams;

/// One trellis branch. In the forward table `from` is the source and `next` the destination;
/// in the predecessor table the same fields describe the branch reaching a destination state.
/// `output` packs the N output bits (first generator / C1 in the high bit).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Branch {
    /// Source state of the branch.
    pub from: usize,
    /// Destination state after shifting `input` in.
    pub next: usize,
    /// Input bit driving the transition (`0` or `1`).
    pub input: u8,
    /// N-bit output symbol (first generator in the high bit), inversion applied.
    pub output: u32,
}

/// The trellis: `forward[state][input]` gives the branch taken, `preds[dest][0..2]` its predecessors.
pub struct Trellis<const S: usize> {
    /// Forward transitions indexed by `[state][input]`.
    ///
    /// Test-only scaffolding: the engine drives encoding through the [`output_symbol`] /
    /// [`next_state`] free functions and decoding through [`Self::preds`], so nothing in a
    /// production build reads this table. It is retained (and built) under `#[cfg(test)]` to let the
    /// unit tests cross-check the forward transitions against the predecessor inversion.
    #[cfg(test)]
    pub(crate) forward: [[Branch; 2]; S],
    /// Predecessor branches indexed by `[destination][slot]`. There are **exactly 2** predecessors per
    /// state because the code takes **one input bit per stage** (2 = 2¹ input symbols) — this holds for
    /// any rate 1/n, independent of the generator count. While inverting the forward table `from_params`
    /// writes predecessors through a bounds-checked `get_mut`, returning an error (never an out-of-bounds
    /// panic) if the count would exceed 2; the engine only builds trellises from validated, rate-1/2
    /// `CodeParams`, so that branch is unreachable and the invariant is upheld by construction.
    pub(crate) preds: [[Branch; 2]; S],
}

/// Next state after shifting `input` into the K-bit register: `(g >> 1) & (S - 1)`.
/// `pub(crate)` — an internal helper; callers must pass a valid `k` (1..=9), which the crate
/// guarantees because it is only ever invoked with a validated [`CodeParams`].
pub(crate) fn next_state(state: usize, input: u8, k: u8) -> usize {
    let m = (k - 1) as u32;
    let g = ((input as usize) << m) | state;
    (g >> 1) & ((1usize << m) - 1)
}

/// The N-bit output symbol for `(state, input)` with per-generator inversion applied
/// (first generator in the high bit). `pub(crate)` — an internal helper; the crate only ever
/// calls it with a validated [`CodeParams`], so `p.k` is always in range (1..=9).
pub(crate) fn output_symbol(state: usize, input: u8, p: &CodeParams) -> u32 {
    let m = (p.k - 1) as u32;
    let g = ((input as u32) << m) | state as u32;
    let mut sym = 0u32;
    for (i, (&gen, &inv)) in p.generators.iter().zip(&p.invert_outputs).enumerate() {
        let mut bit = (g & gen).count_ones() & 1;
        if inv {
            bit ^= 1;
        }
        sym |= bit << (p.generators.len() - 1 - i); // first generator → high bit
    }
    sym
}

impl<const S: usize> Trellis<S> {
    /// Build both forward and predecessor tables from `p`. Returns
    /// `Err(ParamError::StateCountMismatch)` if `S != 2^(k-1)` — **never panics** (R11).
    pub fn from_params(p: &CodeParams) -> Result<Self, crate::params::ParamError> {
        if S != 1usize << (p.k - 1) {
            return Err(crate::params::ParamError::StateCountMismatch { states: S, k: p.k });
        }
        let blank = Branch {
            from: 0,
            next: 0,
            input: 0,
            output: 0,
        };
        let mut preds = [[blank; 2]; S];
        let mut pred_count = [0usize; S];
        for state in 0..S {
            for input in 0u8..2 {
                let next = next_state(state, input, p.k);
                let output = output_symbol(state, input, p);
                let b = Branch {
                    from: state,
                    next,
                    input,
                    output,
                };
                let slot = pred_count[next];
                // A validated 1-input-bit (rate 1/n) code yields exactly 2 predecessors per state.
                // Guard the write with `get_mut` so a hypothetical inconsistency returns an error
                // instead of an out-of-bounds panic — release-safe defense-in-depth (R11). Unreachable
                // for the rate-1/2 params the engine actually builds trellises from.
                match preds[next].get_mut(slot) {
                    Some(cell) => *cell = b,
                    None => {
                        return Err(crate::params::ParamError::RateUnsupported {
                            n: p.generators.len(),
                        })
                    }
                }
                pred_count[next] += 1;
            }
        }
        // Test-only: build the forward table so unit tests can cross-check it against the
        // predecessor inversion. Iterator form avoids a needless range-index loop (clippy).
        #[cfg(test)]
        let forward = {
            let mut forward = [[blank; 2]; S];
            for (state, fwd) in forward.iter_mut().enumerate() {
                for (input, cell) in fwd.iter_mut().enumerate() {
                    let input = input as u8;
                    *cell = Branch {
                        from: state,
                        next: next_state(state, input, p.k),
                        input,
                        output: output_symbol(state, input, p),
                    };
                }
            }
            forward
        };
        Ok(Self {
            #[cfg(test)]
            forward,
            preds,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::params::CodeParams;
    #[test]
    fn state0_input1_matches_ccsds() {
        let p = CodeParams::ccsds_r1_2();
        // g=0b1000000: C1=parity(&0o171)=1, C2=parity(&0o133)=1 → inverted 0 ⇒ symbol 0b10=2; next=0b100000=32
        assert_eq!(super::output_symbol(0, 1, &p), 0b10);
        assert_eq!(super::next_state(0, 1, p.k), 32);
    }
    #[test]
    fn state0_input0_matches_ccsds() {
        let p = CodeParams::ccsds_r1_2();
        // g=0: C1=0, C2=0→inv 1 ⇒ symbol 0b01=1; next=0
        assert_eq!(super::output_symbol(0, 0, &p), 0b01);
        assert_eq!(super::next_state(0, 0, p.k), 0);
    }
    #[test]
    fn predecessors_invert_forward() {
        let t = Trellis::<64>::from_params(&CodeParams::ccsds_r1_2()).unwrap();
        // state 0 reached from state 0 (input 0) and state 1 (input 0): both shift a 0 out at bit0
        let preds: Vec<usize> = t.preds[0].iter().map(|b| b.from).collect();
        assert!(preds.contains(&0) && preds.contains(&1));
    }
    #[test]
    fn from_params_rejects_wrong_state_count() {
        // S=32 but K=7 ⇒ 2^6=64: must Err, not panic (R11).
        assert!(matches!(
            Trellis::<32>::from_params(&CodeParams::ccsds_r1_2()),
            Err(crate::params::ParamError::StateCountMismatch { states: 32, k: 7 })
        ));
    }
}