Skip to main content

reddb_server/storage/engine/turboquant/
rotation.rs

1//! Deterministic orthogonal rotation for TurboQuant.
2//!
3//! MIT notice: clean-room RedDB implementation for the turbovec-compatible
4//! TurboQuant surface; no upstream turbovec source is copied.
5
6#[derive(Debug, Clone, PartialEq)]
7pub struct RotationMatrix {
8    dim: usize,
9    seed: u64,
10    permutation: Vec<usize>,
11    signs: Vec<f32>,
12}
13
14impl RotationMatrix {
15    pub fn new(dim: usize, seed: u64) -> Self {
16        let mut permutation: Vec<usize> = (0..dim).collect();
17        let mut state = seed ^ ((dim as u64) << 32) ^ 0x9E37_79B9_7F4A_7C15;
18        for i in (1..dim).rev() {
19            let j = (splitmix64(&mut state) as usize) % (i + 1);
20            permutation.swap(i, j);
21        }
22        let signs = (0..dim)
23            .map(|_| {
24                if splitmix64(&mut state) & 1 == 0 {
25                    1.0
26                } else {
27                    -1.0
28                }
29            })
30            .collect();
31        Self {
32            dim,
33            seed,
34            permutation,
35            signs,
36        }
37    }
38
39    pub fn dim(&self) -> usize {
40        self.dim
41    }
42
43    pub fn seed(&self) -> u64 {
44        self.seed
45    }
46
47    pub fn rotate(&self, input: &[f32]) -> Vec<f32> {
48        assert_eq!(input.len(), self.dim, "rotation dimension mismatch");
49        self.permutation
50            .iter()
51            .zip(&self.signs)
52            .map(|(&source, &sign)| input[source] * sign)
53            .collect()
54    }
55
56    pub fn row_entries(&self, row: usize) -> Option<(usize, f32)> {
57        (row < self.dim).then(|| (self.permutation[row], self.signs[row]))
58    }
59}
60
61fn splitmix64(state: &mut u64) -> u64 {
62    *state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
63    let mut z = *state;
64    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
65    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
66    z ^ (z >> 31)
67}
68
69#[cfg(test)]
70mod tests {
71    use super::*;
72
73    #[test]
74    fn rotation_is_bit_identical_for_same_dim_and_seed() {
75        let a = RotationMatrix::new(1536, 42);
76        let b = RotationMatrix::new(1536, 42);
77        assert_eq!(a, b);
78        assert_eq!(a.rotate(&vec![1.0; 1536]), b.rotate(&vec![1.0; 1536]));
79    }
80}