Skip to main content

reddb_server/storage/engine/turboquant/
assigner.rs

1//! Block placement for the TurboQuant blocked-by-32 layout.
2//!
3//! Trivial deep module: callers ask "where does the next vector go?"
4//! and the assigner answers without exposing block bookkeeping. Split
5//! out from [`super::storage::BlockedCodeStorage`] so the placement
6//! policy can be reasoned about and unit-tested independently of the
7//! aligned-buffer machinery.
8
9use super::storage::BLOCK_LANES;
10
11#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
12pub struct BlockPlacement {
13    pub block_idx: u32,
14    pub lane: u8,
15}
16
17#[derive(Debug, Default, Clone, Copy)]
18pub struct BlockAssigner;
19
20impl BlockAssigner {
21    pub fn new() -> Self {
22        Self
23    }
24
25    /// Decide where the next vector should be written.
26    ///
27    /// - `n_blocks == 0` ⇒ open block 0 at lane 0.
28    /// - trailing block has `< BLOCK_LANES` lanes filled ⇒ append at the
29    ///   next free lane in that block.
30    /// - trailing block is full ⇒ open a new block at lane 0.
31    pub fn next_placement(&self, n_blocks: usize, trailing_lanes_filled: usize) -> BlockPlacement {
32        if n_blocks == 0 {
33            return BlockPlacement {
34                block_idx: 0,
35                lane: 0,
36            };
37        }
38        if trailing_lanes_filled < BLOCK_LANES {
39            return BlockPlacement {
40                block_idx: (n_blocks - 1) as u32,
41                lane: trailing_lanes_filled as u8,
42            };
43        }
44        BlockPlacement {
45            block_idx: n_blocks as u32,
46            lane: 0,
47        }
48    }
49}
50
51#[cfg(test)]
52mod tests {
53    use super::*;
54
55    #[test]
56    fn empty_assigns_block_zero_lane_zero() {
57        let placement = BlockAssigner::new().next_placement(0, 0);
58        assert_eq!(
59            placement,
60            BlockPlacement {
61                block_idx: 0,
62                lane: 0
63            }
64        );
65    }
66
67    #[test]
68    fn partial_block_assigns_next_lane() {
69        let assigner = BlockAssigner::new();
70        for filled in 1..BLOCK_LANES {
71            let placement = assigner.next_placement(1, filled);
72            assert_eq!(
73                placement,
74                BlockPlacement {
75                    block_idx: 0,
76                    lane: filled as u8,
77                },
78                "partial trailing block, {filled} lanes filled",
79            );
80        }
81
82        let placement = assigner.next_placement(3, 7);
83        assert_eq!(
84            placement,
85            BlockPlacement {
86                block_idx: 2,
87                lane: 7
88            }
89        );
90    }
91
92    #[test]
93    fn full_block_opens_new_block_at_lane_zero() {
94        let placement = BlockAssigner::new().next_placement(1, BLOCK_LANES);
95        assert_eq!(
96            placement,
97            BlockPlacement {
98                block_idx: 1,
99                lane: 0
100            }
101        );
102
103        let placement = BlockAssigner::new().next_placement(4, BLOCK_LANES);
104        assert_eq!(
105            placement,
106            BlockPlacement {
107                block_idx: 4,
108                lane: 0
109            }
110        );
111    }
112}