congee/
node_16.rs

1use crate::{
2    base_node::{BaseNode, Node, NodeIter, NodeType},
3    node_ptr::NodePtr,
4};
5
6#[repr(C)]
7#[repr(align(8))] // Node 16 doesn't need to align to 64 bc it occupies 3 cache lines anyway
8pub(crate) struct Node16 {
9    base: BaseNode,
10    children: [NodePtr; 16],
11    keys: [u8; 16],
12}
13
14#[cfg(not(feature = "shuttle"))]
15const _: () = assert!(std::mem::size_of::<Node16>() == 168);
16
17#[cfg(not(feature = "shuttle"))]
18const _: () = assert!(std::mem::align_of::<Node16>() == 8);
19
20impl Node16 {
21    fn get_insert_pos(&self, key: u8) -> usize {
22        let mut pos = 0;
23        while pos < self.base.meta.count {
24            if self.keys[pos as usize] >= key {
25                return pos as usize;
26            }
27            pos += 1;
28        }
29        pos as usize
30    }
31
32    fn get_child_pos(&self, key: u8) -> Option<usize> {
33        // TODO: xiangpeng check this code is being auto-vectorized
34
35        self.keys
36            .iter()
37            .take(self.base.meta.count as usize)
38            .position(|k| *k == key)
39    }
40}
41
42pub(crate) struct Node16Iter<'a> {
43    node: &'a Node16,
44    start_pos: usize,
45    end_pos: usize,
46}
47
48impl Iterator for Node16Iter<'_> {
49    type Item = (u8, NodePtr);
50
51    fn next(&mut self) -> Option<Self::Item> {
52        if self.start_pos > self.end_pos {
53            return None;
54        }
55        let key = self.node.keys[self.start_pos];
56        let child = self.node.children[self.start_pos];
57        self.start_pos += 1;
58        Some((key, child))
59    }
60}
61
62impl Node for Node16 {
63    fn get_type() -> NodeType {
64        NodeType::N16
65    }
66
67    fn get_children(&self, start: u8, end: u8) -> NodeIter {
68        if self.base.meta.count == 0 {
69            // FIXME: the node may be empty due to deletion, this is not intended, we should fix the delete logic
70            return NodeIter::N16(Node16Iter {
71                node: self,
72                start_pos: 1,
73                end_pos: 0,
74            });
75        }
76        let start_pos = self.get_child_pos(start).unwrap_or(0);
77        let end_pos = self
78            .get_child_pos(end)
79            .unwrap_or(self.base.meta.count as usize - 1);
80
81        debug_assert!(end_pos < 16);
82
83        NodeIter::N16(Node16Iter {
84            node: self,
85            start_pos,
86            end_pos,
87        })
88    }
89
90    fn remove(&mut self, k: u8) {
91        let pos = self
92            .get_child_pos(k)
93            .expect("trying to delete a non-existing key");
94
95        self.keys
96            .copy_within(pos + 1..self.base.meta.count as usize, pos);
97        self.children
98            .copy_within(pos + 1..self.base.meta.count as usize, pos);
99
100        self.base.meta.count -= 1;
101        debug_assert!(self.get_child(k).is_none());
102    }
103
104    fn copy_to<N: Node>(&self, dst: &mut N) {
105        for i in 0..self.base.meta.count {
106            dst.insert(self.keys[i as usize], self.children[i as usize]);
107        }
108    }
109
110    fn base(&self) -> &BaseNode {
111        &self.base
112    }
113
114    fn is_full(&self) -> bool {
115        self.base.meta.count == 16
116    }
117
118    // Insert must keep keys sorted, is this necessary?
119    fn insert(&mut self, key: u8, node: NodePtr) {
120        let pos = self.get_insert_pos(key);
121
122        if pos < self.base.meta.count as usize {
123            self.keys
124                .copy_within(pos..self.base.meta.count as usize, pos + 1);
125            self.children
126                .copy_within(pos..self.base.meta.count as usize, pos + 1);
127        }
128
129        self.keys[pos] = key;
130        self.children[pos] = node;
131        self.base.meta.count += 1;
132
133        assert!(self.base.meta.count <= 16);
134    }
135
136    fn change(&mut self, key: u8, val: NodePtr) -> NodePtr {
137        let pos = self.get_child_pos(key).unwrap();
138        let old = self.children[pos];
139        self.children[pos] = val;
140        old
141    }
142
143    fn get_child(&self, key: u8) -> Option<NodePtr> {
144        let pos = self.get_child_pos(key)?;
145        let child = unsafe { self.children.get_unchecked(pos) };
146        Some(*child)
147    }
148}