reddb_server/storage/engine/turboquant/
rotation.rs1#[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}