geographdb-core 0.5.4

//! Cartan Moving Frames KV Cache Codec.
//!
//! Instead of storing absolute key/value vectors `K_t`, `V_t`, this codec
//! stores incremental *frame transitions*:
//!
//!   `F_t = K_t · pinv(K_{t-1})`  (Moore-Penrose pseudoinverse)
//!
//! Reconstruction: `K_t = F_t · F_{t-1} · … · F_1 · K_0`
//!
//! **Why this compresses:** For coherent text, consecutive key vectors lie on
//! a smooth manifold.  The frame transition `F_t - I` has small Frobenius norm
//! and low rank — far more compressible than absolute vectors.
//!
//! This is the differential-geometry analog of storing velocity instead of
//! position.  Élie Cartan (1900s) developed the moving-frame method for
//! exactly this purpose: encoding manifold geometry through infinitesimal
//! frame changes rather than absolute coordinates.
//!
//! # Usage
//! ```ignore
//! let mut codec = FrameCodec::new(head_dim);
//! for (k, v) in kv_pairs { codec.encode_step(&k, &v); }
//! let output = codec.attend(&query);
//! ```

// ── Types ─────────────────────────────────────────────────────────────────────

/// Cartan moving-frame KV cache codec.
///
/// Stores the anchor `(K_0, V_0)` plus a sequence of `d×d` frame transition
/// matrices `F_t` for `t = 1..T`.  Decoding reconstructs `K_t` and `V_t`
/// by cumulative matrix product from the anchor.
pub struct FrameCodec {
    /// Physical dimension of each key/value head vector.
    head_dim: usize,
    /// Anchor key K_0 (length `head_dim`).
    k0: Vec<f32>,
    /// Anchor value V_0 (length `head_dim`).
    v0: Vec<f32>,
    /// Frame transitions for keys: F_1, F_2, …, F_T.
    /// Each entry is a flat `head_dim × head_dim` row-major matrix.
    frames_k: Vec<Vec<f32>>,
    /// Frame transitions for values.
    frames_v: Vec<Vec<f32>>,
}

impl FrameCodec {
    /// Create a new empty codec for vectors of length `head_dim`.
    pub fn new(head_dim: usize) -> Self {
        Self {
            head_dim,
            k0: Vec::new(),
            v0: Vec::new(),
            frames_k: Vec::new(),
            frames_v: Vec::new(),
        }
    }

    /// Encode one `(k, v)` pair. First call sets the anchor; subsequent calls
    /// store the frame transition relative to the previous step.
    ///
    /// `k` and `v` must each have length `head_dim`.
    pub fn encode_step(&mut self, k: &[f32], v: &[f32]) {
        assert_eq!(k.len(), self.head_dim, "k length must equal head_dim");
        assert_eq!(v.len(), self.head_dim, "v length must equal head_dim");

        if self.k0.is_empty() {
            // First token: store as anchor
            self.k0 = k.to_vec();
            self.v0 = v.to_vec();
            return;
        }

        // Decode the previous key/value to compute frame transition
        let k_prev = self.decode_key(self.token_count() - 1);
        let v_prev = self.decode_val(self.token_count() - 1);

        // F_t = K_t ⊗ pinv(K_{t-1})
        // For vectors (d×1 outer product interpretation), the "frame" that maps
        // k_prev → k is the rank-1 outer product update: F = k * k_prev^T / ‖k_prev‖²
        // This is the closest rank-1 frame; for full d×d we use the identity
        // + a rank-1 correction: F = I + (k - k_prev) ⊗ k_prev / ‖k_prev‖²
        let fk = rank1_frame_transition(k, &k_prev, self.head_dim);
        let fv = rank1_frame_transition(v, &v_prev, self.head_dim);

        self.frames_k.push(fk);
        self.frames_v.push(fv);
    }

    /// Number of tokens stored (including anchor).
    pub fn token_count(&self) -> usize {
        if self.k0.is_empty() {
            0
        } else {
            1 + self.frames_k.len()
        }
    }

    /// Decode the key vector at position `t` by applying cumulative frame
    /// transitions from the anchor.
    pub fn decode_key(&self, t: usize) -> Vec<f32> {
        self.decode_vector(&self.k0, &self.frames_k, t)
    }

    /// Decode the value vector at position `t`.
    pub fn decode_val(&self, t: usize) -> Vec<f32> {
        self.decode_vector(&self.v0, &self.frames_v, t)
    }

    /// Compute standard softmax attention over all stored tokens.
    ///
    /// `query` has length `head_dim`. Returns the weighted sum of values.
    pub fn attend(&self, query: &[f32]) -> Vec<f32> {
        let n = self.token_count();
        if n == 0 {
            return vec![0.0; self.head_dim];
        }

        let scale = 1.0 / (self.head_dim as f32).sqrt();
        let mut scores = Vec::with_capacity(n);

        // Decode all keys to compute scores
        let mut keys = Vec::with_capacity(n);
        let mut vals = Vec::with_capacity(n);
        keys.push(self.k0.clone());
        vals.push(self.v0.clone());
        for t in 1..n {
            keys.push(self.decode_key(t));
            vals.push(self.decode_val(t));
        }

        for k in &keys {
            let dot: f32 = query.iter().zip(k.iter()).map(|(q, ki)| q * ki).sum();
            scores.push(dot * scale);
        }

        // Softmax
        let max_s = scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let mut exp_scores: Vec<f32> = scores.iter().map(|s| (s - max_s).exp()).collect();
        let sum_exp: f32 = exp_scores.iter().sum();
        for e in &mut exp_scores {
            *e /= sum_exp;
        }

        // Weighted sum of values
        let mut out = vec![0.0f32; self.head_dim];
        for (weight, v) in exp_scores.iter().zip(vals.iter()) {
            for (o, vi) in out.iter_mut().zip(v.iter()) {
                *o += weight * vi;
            }
        }
        out
    }

    /// Number of bytes used by the codec (f32 = 4 bytes each).
    pub fn bytes_used(&self) -> usize {
        let anchor = 2 * self.head_dim * 4;
        let frames =
            (self.frames_k.len() + self.frames_v.len()) * self.head_dim * self.head_dim * 4;
        anchor + frames
    }

    /// Number of bytes a naive absolute-vector store would use (f32 K+V per token).
    pub fn bytes_absolute(&self) -> usize {
        self.token_count() * 2 * self.head_dim * 4
    }

    // ── Private ───────────────────────────────────────────────────────────────

    fn decode_vector(&self, anchor: &[f32], frames: &[Vec<f32>], t: usize) -> Vec<f32> {
        assert!(
            t < self.token_count(),
            "t={t} out of range (n={})",
            self.token_count()
        );
        if t == 0 {
            return anchor.to_vec();
        }
        let d = self.head_dim;
        // Apply frames[0..t] sequentially: v = F_{t-1} · F_{t-2} · … · F_0 · anchor
        let mut current = anchor.to_vec();
        for frame in &frames[..t] {
            current = matvec(frame, &current, d);
        }
        current
    }
}

// ── Math helpers ──────────────────────────────────────────────────────────────

/// Rank-1 frame transition that maps `prev` → `curr`:
///   F = I + (curr - prev) ⊗ prev / ‖prev‖²
///
/// This is the minimum-norm matrix correction satisfying `F · prev = curr`.
/// When `prev ≈ curr` (coherent text), `F ≈ I` → low Frobenius norm → high
/// compressibility.
fn rank1_frame_transition(curr: &[f32], prev: &[f32], d: usize) -> Vec<f32> {
    let norm_sq: f32 = prev.iter().map(|x| x * x).sum();
    let mut frame = identity_matrix(d);

    if norm_sq < 1e-12 {
        // prev is zero vector — store curr as absolute offset from I
        for i in 0..d {
            frame[i * d + i] += curr[i]; // diagonal: I + curr (approx)
        }
        return frame;
    }

    // delta = (curr - prev)
    let delta: Vec<f32> = curr.iter().zip(prev.iter()).map(|(c, p)| c - p).collect();

    // F[i,j] += delta[i] * prev[j] / norm_sq
    for i in 0..d {
        for j in 0..d {
            frame[i * d + j] += delta[i] * prev[j] / norm_sq;
        }
    }
    frame
}

/// `d×d` identity matrix (flat, row-major).
fn identity_matrix(d: usize) -> Vec<f32> {
    let mut m = vec![0.0f32; d * d];
    for i in 0..d {
        m[i * d + i] = 1.0;
    }
    m
}

/// Matrix-vector product: `y = M · x` where `M` is `d×d` row-major.
fn matvec(m: &[f32], x: &[f32], d: usize) -> Vec<f32> {
    let mut y = vec![0.0f32; d];
    for i in 0..d {
        for j in 0..d {
            y[i] += m[i * d + j] * x[j];
        }
    }
    y
}

// ── Tests ─────────────────────────────────────────────────────────────────────

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

    fn approx_eq_vec(a: &[f32], b: &[f32], tol: f32) -> bool {
        a.len() == b.len() && a.iter().zip(b).all(|(x, y)| (x - y).abs() < tol)
    }

    // ── Encode / decode roundtrip ─────────────────────────────────────────────

    #[test]
    fn test_encode_decode_roundtrip_anchor() {
        let mut codec = FrameCodec::new(4);
        let k0 = vec![1.0, 0.0, 0.0, 0.0];
        let v0 = vec![0.0, 1.0, 0.0, 0.0];
        codec.encode_step(&k0, &v0);
        assert!(approx_eq_vec(&codec.decode_key(0), &k0, 1e-6));
        assert!(approx_eq_vec(&codec.decode_val(0), &v0, 1e-6));
    }

    #[test]
    fn test_encode_decode_roundtrip_second_token() {
        let mut codec = FrameCodec::new(4);
        let k0 = vec![1.0, 0.0, 0.0, 0.0];
        let v0 = vec![0.0, 1.0, 0.0, 0.0];
        let k1 = vec![0.5, 0.5, 0.0, 0.0];
        let v1 = vec![0.0, 0.5, 0.5, 0.0];
        codec.encode_step(&k0, &v0);
        codec.encode_step(&k1, &v1);
        assert!(
            approx_eq_vec(&codec.decode_key(1), &k1, 1e-5),
            "decode_key(1) = {:?}, expected {:?}",
            codec.decode_key(1),
            k1
        );
        assert!(approx_eq_vec(&codec.decode_val(1), &v1, 1e-5));
    }

    #[test]
    fn test_encode_decode_multiple_tokens() {
        let d = 4;
        let mut codec = FrameCodec::new(d);
        let kvs: Vec<(Vec<f32>, Vec<f32>)> = (0..5)
            .map(|i| {
                let k = (0..d).map(|j| ((i + j) as f32) * 0.1).collect();
                let v = (0..d).map(|j| ((i * 2 + j) as f32) * 0.1).collect();
                (k, v)
            })
            .collect();
        for (k, v) in &kvs {
            codec.encode_step(k, v);
        }
        for (t, (k, v)) in kvs.iter().enumerate() {
            assert!(
                approx_eq_vec(&codec.decode_key(t), k, 1e-4),
                "key mismatch at t={t}"
            );
            assert!(
                approx_eq_vec(&codec.decode_val(t), v, 1e-4),
                "val mismatch at t={t}"
            );
        }
    }

    #[test]
    fn test_identity_frame_for_constant_sequence() {
        let d = 3;
        let mut codec = FrameCodec::new(d);
        let k = vec![1.0, 0.0, 0.0];
        let v = vec![0.0, 1.0, 0.0];
        // Same vector repeated → F_t = I
        for _ in 0..4 {
            codec.encode_step(&k, &v);
        }
        // All decoded keys should equal k
        for t in 0..4 {
            assert!(approx_eq_vec(&codec.decode_key(t), &k, 1e-5), "t={t}");
        }
    }

    // ── Attention ─────────────────────────────────────────────────────────────

    #[test]
    fn test_attend_matches_direct_single_token() {
        let d = 4;
        let k0 = vec![1.0, 0.0, 0.0, 0.0];
        let v0 = vec![0.0, 0.0, 1.0, 0.0];
        let q = vec![1.0, 0.0, 0.0, 0.0];

        let mut codec = FrameCodec::new(d);
        codec.encode_step(&k0, &v0);

        // With 1 token, softmax(score)=1.0, output = v0
        let out = codec.attend(&q);
        assert!(
            approx_eq_vec(&out, &v0, 1e-5),
            "single-token attend = {out:?}"
        );
    }

    #[test]
    fn test_attend_matches_direct_multi_token() {
        let d = 4;
        let kvs = vec![
            (vec![1.0f32, 0.0, 0.0, 0.0], vec![1.0f32, 0.0, 0.0, 0.0]),
            (vec![0.0f32, 1.0, 0.0, 0.0], vec![0.0f32, 1.0, 0.0, 0.0]),
            (vec![0.0f32, 0.0, 1.0, 0.0], vec![0.0f32, 0.0, 1.0, 0.0]),
        ];
        let q = vec![1.0f32, 0.0, 0.0, 0.0]; // aligns with k0

        let mut codec = FrameCodec::new(d);
        for (k, v) in &kvs {
            codec.encode_step(k, v);
        }

        // Direct reference
        let scale = 1.0 / (d as f32).sqrt();
        let scores_raw: Vec<f32> = kvs
            .iter()
            .map(|(k, _)| k.iter().zip(&q).map(|(ki, qi)| ki * qi).sum::<f32>() * scale)
            .collect();
        let max_s = scores_raw.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let exp: Vec<f32> = scores_raw.iter().map(|s| (s - max_s).exp()).collect();
        let sum_exp: f32 = exp.iter().sum();
        let weights: Vec<f32> = exp.iter().map(|e| e / sum_exp).collect();
        let mut expected = vec![0.0f32; d];
        for (w, (_, v)) in weights.iter().zip(kvs.iter()) {
            for (o, vi) in expected.iter_mut().zip(v.iter()) {
                *o += w * vi;
            }
        }

        let out = codec.attend(&q);
        assert!(
            approx_eq_vec(&out, &expected, 1e-4),
            "attend mismatch: got {out:?}, expected {expected:?}"
        );
    }

    #[test]
    fn test_attend_empty_returns_zero() {
        let codec = FrameCodec::new(4);
        let out = codec.attend(&[1.0, 0.0, 0.0, 0.0]);
        assert_eq!(out, vec![0.0; 4]);
    }

    // ── Bytes / compression ───────────────────────────────────────────────────

    #[test]
    fn test_token_count() {
        let mut codec = FrameCodec::new(4);
        assert_eq!(codec.token_count(), 0);
        codec.encode_step(&[1.0, 0.0, 0.0, 0.0], &[0.0, 1.0, 0.0, 0.0]);
        assert_eq!(codec.token_count(), 1);
        codec.encode_step(&[0.5, 0.5, 0.0, 0.0], &[0.0, 0.5, 0.5, 0.0]);
        assert_eq!(codec.token_count(), 2);
    }
}