helena 0.1.0

//! The token latent payload: frame-aligned discrete codes.

use serde::{Deserialize, Serialize};

use crate::error::{Error, Result};

use super::kind::Framed;

/// Discrete codec tokens laid out as `[codebooks][frames]`.
///
/// Invariant, enforced by [`new`](Self::new) and re-checked on
/// deserialization: at least one codebook, and every codebook holds the same
/// number of frames. RVQ levels are frame-aligned (a frame's coarse and fine
/// codes describe the same instant), so a ragged grid is not a representable
/// token latent. An empty grid (zero frames) is permitted: emptiness is a
/// downstream concern, frame-alignment is the structural one. Codebook count
/// is a *fact of the grid*, not a kind — the old structure-derived
/// discrete-vs-RVQ tag is gone.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(into = "Vec<Vec<u32>>", try_from = "Vec<Vec<u32>>")]
pub struct TokenGrid(Vec<Vec<u32>>);

impl TokenGrid {
    /// Build a grid from per-codebook token streams laid out
    /// `[codebooks][frames]`.
    ///
    /// Returns [`Error::Validation`] if there are no codebooks, and
    /// [`Error::Shape`] if the codebooks disagree on frame count.
    pub fn new(codes: impl Into<Vec<Vec<u32>>>) -> Result<Self> {
        let codes = codes.into();
        let Some(frames) = codes.first().map(Vec::len) else {
            return Err(Error::validation("token grid needs at least one codebook"));
        };
        if let Some(level) = codes.iter().position(|stream| stream.len() != frames) {
            return Err(Error::shape(
                vec![codes.len(), frames],
                vec![codes.len(), codes[level].len()],
            ));
        }
        Ok(Self(codes))
    }

    /// Number of codebooks (RVQ levels); always at least one.
    pub fn codebooks(&self) -> usize {
        self.0.len()
    }

    /// Borrow the per-codebook token streams, `[codebooks][frames]`.
    pub fn as_slice(&self) -> &[Vec<u32>] {
        &self.0
    }

    /// Consume into the raw `[codebooks][frames]` token streams.
    pub fn into_inner(self) -> Vec<Vec<u32>> {
        self.0
    }
}

impl Framed for TokenGrid {
    fn frames(&self) -> usize {
        self.0[0].len()
    }

    fn slice_frames(&self, start: usize, len: usize) -> Self {
        Self(
            self.0
                .iter()
                .map(|stream| stream[start..start + len].to_vec())
                .collect(),
        )
    }
}

impl TryFrom<Vec<Vec<u32>>> for TokenGrid {
    type Error = Error;

    fn try_from(codes: Vec<Vec<u32>>) -> Result<Self> {
        Self::new(codes)
    }
}

impl From<TokenGrid> for Vec<Vec<u32>> {
    fn from(grid: TokenGrid) -> Self {
        grid.0
    }
}

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

    #[test]
    fn alignment_is_the_structural_invariant() {
        assert!(TokenGrid::new(Vec::<Vec<u32>>::new()).is_err());
        assert!(matches!(
            TokenGrid::new(vec![vec![1, 2], vec![3]]),
            Err(Error::Shape { .. })
        ));
        let grid = TokenGrid::new(vec![vec![1, 2, 3], vec![4, 5, 6]]).unwrap();
        assert_eq!(grid.codebooks(), 2);
        assert_eq!(grid.frames(), 3);
        let empty = TokenGrid::new(vec![vec![]]).unwrap();
        assert_eq!(empty.frames(), 0);
    }

    #[test]
    fn slices_cut_every_codebook() {
        let grid = TokenGrid::new(vec![vec![1, 2, 3, 4], vec![5, 6, 7, 8]]).unwrap();
        let mid = grid.slice_frames(1, 2);
        assert_eq!(mid.as_slice(), &[vec![2, 3], vec![6, 7]]);
    }

    #[test]
    fn serde_fails_closed_on_ragged_grids() {
        let grid = TokenGrid::new(vec![vec![9, 8]]).unwrap();
        let json = serde_json::to_string(&grid).unwrap();
        assert_eq!(serde_json::from_str::<TokenGrid>(&json).unwrap(), grid);
        assert!(serde_json::from_str::<TokenGrid>("[[1,2],[3]]").is_err());
        assert!(serde_json::from_str::<TokenGrid>("[]").is_err());
    }
}