Skip to main content

btree/
node.rs

1use crate::error::Error;
2use crate::node_type::{Key, KeyValuePair, NodeType, Offset};
3use crate::page::Page;
4use crate::page_layout::{
5    FromByte, INTERNAL_NODE_HEADER_SIZE, INTERNAL_NODE_NUM_CHILDREN_OFFSET, IS_ROOT_OFFSET,
6    KEY_SIZE, LEAF_NODE_HEADER_SIZE, LEAF_NODE_NUM_PAIRS_OFFSET, NODE_TYPE_OFFSET,
7    PARENT_POINTER_OFFSET, PTR_SIZE, VALUE_SIZE,
8};
9use std::convert::TryFrom;
10use std::str;
11
12/// Node represents a node in the BTree occupied by a single page in memory.
13#[derive(Clone, Debug)]
14pub struct Node {
15    pub node_type: NodeType,
16    pub is_root: bool,
17    pub parent_offset: Option<Offset>,
18}
19
20// Node represents a node in the B-Tree.
21impl Node {
22    pub fn new(node_type: NodeType, is_root: bool, parent_offset: Option<Offset>) -> Node {
23        Node {
24            node_type,
25            is_root,
26            parent_offset,
27        }
28    }
29
30    /// split creates a sibling node from a given node by splitting the node in two around a median.
31    /// split will split the child at b leaving the [0, b-1] keys
32    /// while moving the set of [b, 2b-1] keys to the sibling.
33    pub fn split(&mut self, b: usize) -> Result<(Key, Node), Error> {
34        match self.node_type {
35            NodeType::Internal(ref mut children, ref mut keys) => {
36                // Populate siblings keys.
37                let mut sibling_keys = keys.split_off(b - 1);
38                // Pop median key - to be added to the parent..
39                let median_key = sibling_keys.remove(0);
40                // Populate siblings children.
41                let sibling_children = children.split_off(b);
42                Ok((
43                    median_key,
44                    Node::new(
45                        NodeType::Internal(sibling_children, sibling_keys),
46                        false,
47                        self.parent_offset.clone(),
48                    ),
49                ))
50            }
51            NodeType::Leaf(ref mut pairs) => {
52                // Populate siblings pairs.
53                let sibling_pairs = pairs.split_off(b);
54                // Pop median key.
55                let median_pair = pairs.get(b - 1).ok_or(Error::UnexpectedError)?.clone();
56
57                Ok((
58                    Key(median_pair.key),
59                    Node::new(
60                        NodeType::Leaf(sibling_pairs),
61                        false,
62                        self.parent_offset.clone(),
63                    ),
64                ))
65            }
66            NodeType::Unexpected => Err(Error::UnexpectedError),
67        }
68    }
69}
70
71/// Implement TryFrom<Page> for Node allowing for easier
72/// deserialization of data from a Page.
73impl TryFrom<Page> for Node {
74    type Error = Error;
75    fn try_from(page: Page) -> Result<Node, Error> {
76        let raw = page.get_data();
77        let node_type = NodeType::from(raw[NODE_TYPE_OFFSET]);
78        let is_root = raw[IS_ROOT_OFFSET].from_byte();
79        let parent_offset: Option<Offset>;
80        if is_root {
81            parent_offset = None;
82        } else {
83            parent_offset = Some(Offset(page.get_value_from_offset(PARENT_POINTER_OFFSET)?));
84        }
85
86        match node_type {
87            NodeType::Internal(mut children, mut keys) => {
88                let num_children = page.get_value_from_offset(INTERNAL_NODE_NUM_CHILDREN_OFFSET)?;
89                let mut offset = INTERNAL_NODE_HEADER_SIZE;
90                for _i in 1..=num_children {
91                    let child_offset = page.get_value_from_offset(offset)?;
92                    children.push(Offset(child_offset));
93                    offset += PTR_SIZE;
94                }
95
96                // Number of keys is always one less than the number of children (i.e. branching factor)
97                for _i in 1..num_children {
98                    let key_raw = page.get_ptr_from_offset(offset, KEY_SIZE);
99                    let key = match str::from_utf8(key_raw) {
100                        Ok(key) => key,
101                        Err(_) => return Err(Error::UTF8Error),
102                    };
103                    offset += KEY_SIZE;
104                    // Trim leading or trailing zeros.
105                    keys.push(Key(key.trim_matches(char::from(0)).to_string()));
106                }
107                Ok(Node::new(
108                    NodeType::Internal(children, keys),
109                    is_root,
110                    parent_offset,
111                ))
112            }
113
114            NodeType::Leaf(mut pairs) => {
115                let mut offset = LEAF_NODE_NUM_PAIRS_OFFSET;
116                let num_keys_val_pairs = page.get_value_from_offset(offset)?;
117                offset = LEAF_NODE_HEADER_SIZE;
118
119                for _i in 0..num_keys_val_pairs {
120                    let key_raw = page.get_ptr_from_offset(offset, KEY_SIZE);
121                    let key = match str::from_utf8(key_raw) {
122                        Ok(key) => key,
123                        Err(_) => return Err(Error::UTF8Error),
124                    };
125                    offset += KEY_SIZE;
126
127                    let value_raw = page.get_ptr_from_offset(offset, VALUE_SIZE);
128                    let value = match str::from_utf8(value_raw) {
129                        Ok(val) => val,
130                        Err(_) => return Err(Error::UTF8Error),
131                    };
132                    offset += VALUE_SIZE;
133
134                    // Trim leading or trailing zeros.
135                    pairs.push(KeyValuePair::new(
136                        key.trim_matches(char::from(0)).to_string(),
137                        value.trim_matches(char::from(0)).to_string(),
138                    ))
139                }
140                Ok(Node::new(NodeType::Leaf(pairs), is_root, parent_offset))
141            }
142
143            NodeType::Unexpected => Err(Error::UnexpectedError),
144        }
145    }
146}
147
148////////////////////
149///              ///
150///  Unit Tests. ///
151///              ///
152////////////////////
153
154#[cfg(test)]
155mod tests {
156    use crate::error::Error;
157    use crate::node::{
158        Node, Page, INTERNAL_NODE_HEADER_SIZE, KEY_SIZE, LEAF_NODE_HEADER_SIZE, PTR_SIZE,
159        VALUE_SIZE,
160    };
161    use crate::node_type::{Key, NodeType};
162    use crate::page_layout::PAGE_SIZE;
163    use std::convert::TryFrom;
164
165    #[test]
166    fn page_to_node_works_for_leaf_node() -> Result<(), Error> {
167        const DATA_LEN: usize = LEAF_NODE_HEADER_SIZE + KEY_SIZE + VALUE_SIZE;
168        let page_data: [u8; DATA_LEN] = [
169            0x01, // Is-Root byte.
170            0x02, // Leaf Node type byte.
171            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // Parent offset.
172            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, // Number of Key-Value pairs.
173            0x68, 0x65, 0x6c, 0x6c, 0x6f, 0x00, 0x00, 0x00, 0x00, 0x00, // "hello"
174            0x77, 0x6f, 0x72, 0x6c, 0x64, 0x00, 0x00, 0x00, 0x00, 0x00, // "world"
175        ];
176        let junk: [u8; PAGE_SIZE - DATA_LEN] = [0x00; PAGE_SIZE - DATA_LEN];
177        let mut page = [0x00; PAGE_SIZE];
178        for (to, from) in page.iter_mut().zip(page_data.iter().chain(junk.iter())) {
179            *to = *from
180        }
181
182        let node = Node::try_from(Page::new(page))?;
183
184        assert_eq!(node.is_root, true);
185        Ok(())
186    }
187
188    #[test]
189    fn page_to_node_works_for_internal_node() -> Result<(), Error> {
190        use crate::node_type::Key;
191        const DATA_LEN: usize = INTERNAL_NODE_HEADER_SIZE + 3 * PTR_SIZE + 2 * KEY_SIZE;
192        let page_data: [u8; DATA_LEN] = [
193            0x01, // Is-Root byte.
194            0x01, // Internal Node type byte.
195            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // Parent offset.
196            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, // Number of children.
197            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x00, // 4096  (2nd Page)
198            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x20, 0x00, // 8192  (3rd Page)
199            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x30, 0x00, // 12288 (4th Page)
200            0x68, 0x65, 0x6c, 0x6c, 0x6f, 0x00, 0x00, 0x00, 0x00, 0x00, // "hello"
201            0x77, 0x6f, 0x72, 0x6c, 0x64, 0x00, 0x00, 0x00, 0x00, 0x00, // "world"
202        ];
203        let junk: [u8; PAGE_SIZE - DATA_LEN] = [0x00; PAGE_SIZE - DATA_LEN];
204
205        // Concatenate the two arrays; page_data and junk.
206        let mut page = [0x00; PAGE_SIZE];
207        for (to, from) in page.iter_mut().zip(page_data.iter().chain(junk.iter())) {
208            *to = *from
209        }
210
211        let node = Node::try_from(Page::new(page))?;
212
213        if let NodeType::Internal(_, keys) = node.node_type {
214            assert_eq!(keys.len(), 2);
215
216            let Key(first_key) = match keys.get(0) {
217                Some(key) => key,
218                None => return Err(Error::UnexpectedError),
219            };
220            assert_eq!(first_key, "hello");
221
222            let Key(second_key) = match keys.get(1) {
223                Some(key) => key,
224                None => return Err(Error::UnexpectedError),
225            };
226            assert_eq!(second_key, "world");
227            return Ok(());
228        }
229
230        Err(Error::UnexpectedError)
231    }
232
233    #[test]
234    fn split_leaf_works() -> Result<(), Error> {
235        use crate::node::Node;
236        use crate::node_type::KeyValuePair;
237        let mut node = Node::new(
238            NodeType::Leaf(vec![
239                KeyValuePair::new("foo".to_string(), "bar".to_string()),
240                KeyValuePair::new("lebron".to_string(), "james".to_string()),
241                KeyValuePair::new("ariana".to_string(), "grande".to_string()),
242            ]),
243            true,
244            None,
245        );
246
247        let (median, sibling) = node.split(2)?;
248        assert_eq!(median, Key("lebron".to_string()));
249        assert_eq!(
250            node.node_type,
251            NodeType::Leaf(vec![
252                KeyValuePair {
253                    key: "foo".to_string(),
254                    value: "bar".to_string()
255                },
256                KeyValuePair {
257                    key: "lebron".to_string(),
258                    value: "james".to_string()
259                }
260            ])
261        );
262        assert_eq!(
263            sibling.node_type,
264            NodeType::Leaf(vec![KeyValuePair::new(
265                "ariana".to_string(),
266                "grande".to_string()
267            )])
268        );
269        Ok(())
270    }
271
272    #[test]
273    fn split_internal_works() -> Result<(), Error> {
274        use crate::node::Node;
275        use crate::node_type::NodeType;
276        use crate::node_type::{Key, Offset};
277        use crate::page_layout::PAGE_SIZE;
278        let mut node = Node::new(
279            NodeType::Internal(
280                vec![
281                    Offset(PAGE_SIZE),
282                    Offset(PAGE_SIZE * 2),
283                    Offset(PAGE_SIZE * 3),
284                    Offset(PAGE_SIZE * 4),
285                ],
286                vec![
287                    Key("foo bar".to_string()),
288                    Key("lebron".to_string()),
289                    Key("ariana".to_string()),
290                ],
291            ),
292            true,
293            None,
294        );
295
296        let (median, sibling) = node.split(2)?;
297        assert_eq!(median, Key("lebron".to_string()));
298        assert_eq!(
299            node.node_type,
300            NodeType::Internal(
301                vec![Offset(PAGE_SIZE), Offset(PAGE_SIZE * 2)],
302                vec![Key("foo bar".to_string())]
303            )
304        );
305        assert_eq!(
306            sibling.node_type,
307            NodeType::Internal(
308                vec![Offset(PAGE_SIZE * 3), Offset(PAGE_SIZE * 4)],
309                vec![Key("ariana".to_string())]
310            )
311        );
312        Ok(())
313    }
314}