test: improve test coverage of BTreeMap::Node

src/btreemap/node.rs

@@ -24,6 +24,7 @@ const INTERNAL_NODE_TYPE: u8 = 1;
 const U32_SIZE: Bytes = Bytes::new(4);
 
 #[derive(Debug, PartialEq, Copy, Clone, Eq)]
+#[cfg_attr(test, derive(test_strategy::Arbitrary))]
 pub enum NodeType {
     Leaf,
     Internal,
@@ -32,18 +33,7 @@ pub enum NodeType {
 pub type Entry<K> = (K, Vec<u8>);
 
 /// A node of a B-Tree.
-///
-/// The node is stored in stable memory with the following layout:
-///
-///    |  NodeHeader  |  Entries (keys and values) |  Children  |
-///
-/// Each node contains up to `CAPACITY` entries, each entry contains:
-///     - size of key (4 bytes)
-///     - key (`max_key_size` bytes)
-///     - size of value (4 bytes)
-///     - value (`max_value_size` bytes)
-///
-/// Each node can contain up to `CAPACITY + 1` children, each child is 8 bytes.
+/// See `v1.rs` for more details on the memory layout.
 #[derive(Debug)]
 pub struct Node<K: Storable + Ord + Clone> {
     address: Address,
@@ -67,15 +57,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
         max_key_size: u32,
         max_value_size: u32,
     ) -> Node<K> {
-        Node {
-            address,
-            keys: vec![],
-            encoded_values: RefCell::default(),
-            children: vec![],
-            node_type,
-            max_key_size,
-            max_value_size,
-        }
+        Node::new_v1(address, node_type, max_key_size, max_value_size)
     }
 
     /// Loads a node from memory at the given address.
@@ -86,7 +68,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
         max_value_size: u32,
     ) -> Self {
         // NOTE: new versions of `Node` will be introduced.
-        Self::load_v1(address, memory, max_key_size, max_value_size)
+        Self::load_v1(address, max_key_size, max_value_size, memory)
     }
 
     /// Saves the node to memory.
@@ -371,16 +353,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
     ///
     /// See the documentation of [`Node`] for the memory layout.
     pub fn size(max_key_size: u32, max_value_size: u32) -> Bytes {
-        let max_key_size = Bytes::from(max_key_size);
-        let max_value_size = Bytes::from(max_value_size);
-
-        let node_header_size = NodeHeader::size();
-        let entry_size = U32_SIZE + max_key_size + max_value_size + U32_SIZE;
-        let child_size = Address::size();
-
-        node_header_size
-            + Bytes::from(CAPACITY as u64) * entry_size
-            + Bytes::from((CAPACITY + 1) as u64) * child_size
+        v1::size_v1(max_key_size, max_value_size)
     }
 
     /// Returns true if the node is at the minimum required size, false otherwise.

src/btreemap/node/tests.rs

@@ -1,39 +1,91 @@
 use super::*;
-use crate::types::NULL;
 use proptest::collection::btree_map as pmap;
 use proptest::collection::vec as pvec;
 use std::cell::RefCell;
 use std::collections::BTreeMap;
 use std::rc::Rc;
-use test_strategy::proptest;
+use test_strategy::{proptest, Arbitrary};
 
 fn make_memory() -> Rc<RefCell<Vec<u8>>> {
     Rc::new(RefCell::new(Vec::new()))
 }
 
-#[proptest]
-fn loading_and_reloading_entries(
-    #[strategy(1..100u32)] max_key_size: u32,
-    #[strategy(1..1_000u32)] max_value_size: u32,
+/// Generates arbitrary v1 nodes.
+#[derive(Arbitrary, Debug)]
+struct NodeV1Data {
+    #[strategy(0..1_000u32)]
+    max_key_size: u32,
+    #[strategy(0..1_000u32)]
+    max_value_size: u32,
+
+    // NOTE: A BTreeMap is used for creating the entries so that they're in sorted order.
     #[strategy(
         pmap(
-            pvec(0..u8::MAX, 0..#max_key_size as usize),
-            pvec(0..u8::MAX, 0..#max_value_size as usize),
+            pvec(0..u8::MAX, 0..=#max_key_size as usize),
+            pvec(0..u8::MAX, 0..=#max_value_size as usize),
             1..CAPACITY
         )
     )]
     entries: BTreeMap<Vec<u8>, Vec<u8>>,
-) {
+    node_type: NodeType,
+}
+
+impl NodeV1Data {
+    /// Returns a v1 node with the data generated by this struct.
+    fn get(&self, address: Address) -> Node<Vec<u8>> {
+        let mut node = Node::new_v1(
+            address,
+            self.node_type,
+            self.max_key_size,
+            self.max_value_size,
+        );
+
+        // Push the entries
+        for entry in self.entries.clone().into_iter() {
+            node.push_entry(entry);
+        }
+
+        // Push the children
+        for child in self.children() {
+            node.push_child(child);
+        }
+
+        node
+    }
+
+    fn children(&self) -> Vec<Address> {
+        match self.node_type {
+            // A leaf node doesn't have any children.
+            NodeType::Leaf => vec![],
+            // An internal node has entries.len() + 1 children.
+            // Here we generate a list of addresses.
+            NodeType::Internal => (0..=self.entries.len())
+                .map(|i| Address::from(i as u64))
+                .collect(),
+        }
+    }
+}
+
+#[proptest]
+fn saving_and_loading_v1_preserves_data(node_data: NodeV1Data) {
     let mem = make_memory();
 
     // Create a new node and save it into memory.
-    let mut node = Node::new(NULL, NodeType::Leaf, max_key_size, max_value_size);
-    for entry in entries.clone().into_iter() {
-        node.push_entry(entry);
-    }
-    node.save(&mem);
+    let node_addr = Address::from(0);
+    let node = node_data.get(node_addr);
+    node.save_v1(&mem);
+
+    // Load the node and double check all the entries and children are correct.
+    let node = Node::load_v1(
+        node_addr,
+        node_data.max_key_size,
+        node_data.max_value_size,
+        &mem,
+    );
 
-    // Reload the node and double check all the entries are correct.
-    let node = Node::load(NULL, &mem, max_key_size, max_value_size);
-    assert_eq!(node.entries(&mem), entries.into_iter().collect::<Vec<_>>());
+    assert_eq!(node.children, node_data.children());
+    assert_eq!(
+        node.entries(&mem),
+        node_data.entries.into_iter().collect::<Vec<_>>()
+    );
 }

src/btreemap/node/v1.rs

@@ -1,12 +1,68 @@
+//! Node V1
+//!
+//! A v1 node is the first node layout, with fixed node sizes and support for
+//! bounded types only.
+//!
+//! # Memory Layout
+//!
+//! ```text
+//! ---------------------------------------- <-- Header
+//! Magic "BTN"             ↕ 3 bytes
+//! ----------------------------------------
+//! Layout version (2)      ↕ 1 byte
+//! ----------------------------------------
+//! Node type               ↕ 1 byte
+//! ----------------------------------------
+//! # Entries (k)           ↕ 2 bytes
+//! ---------------------------------------- <-- Entries (upto `CAPACITY` entries)
+//! Key(0)
+//! ----------------------------------------
+//! Value(0)
+//! ----------------------------------------
+//! Key(1) size             ↕ 4 bytes
+//! ----------------------------------------
+//! Key(1)                  ↕ `max_key_size` bytes
+//! ----------------------------------------
+//! Value(1) size           ↕ 4 bytes
+//! ----------------------------------------
+//! Value(1)                ↕ `max_value_size` bytes
+//! ----------------------------------------
+//! ...
+//! ---------------------------------------- <-- Children (upto `CAPACITY + 1` children)
+//! Child(0) address        ↕ 8 bytes
+//! ----------------------------------------
+//! ...
+//! ----------------------------------------
+//! Child(k + 1) address    ↕ 8 bytes
+//! ----------------------------------------
+//! ```
 use super::*;
 
 impl<K: Storable + Ord + Clone> Node<K> {
-    /// Loads a node from memory at the given address.
+    /// Creates a new v1 node at the given address.
+    pub(super) fn new_v1(
+        address: Address,
+        node_type: NodeType,
+        max_key_size: u32,
+        max_value_size: u32,
+    ) -> Node<K> {
+        Node {
+            address,
+            node_type,
+            max_key_size,
+            max_value_size,
+            keys: vec![],
+            encoded_values: RefCell::default(),
+            children: vec![],
+        }
+    }
+
+    /// Loads a v1 node from memory at the given address.
     pub(super) fn load_v1<M: Memory>(
         address: Address,
-        memory: &M,
         max_key_size: u32,
         max_value_size: u32,
+        memory: &M,
     ) -> Self {
         // Load the header.
         let header: NodeHeader = read_struct(address, memory);
@@ -131,3 +187,17 @@ impl<K: Storable + Ord + Clone> Node<K> {
         }
     }
 }
+
+/// Returns the size of a v1 node in bytes.
+pub(super) fn size_v1(max_key_size: u32, max_value_size: u32) -> Bytes {
+    let max_key_size = Bytes::from(max_key_size);
+    let max_value_size = Bytes::from(max_value_size);
+
+    let node_header_size = NodeHeader::size();
+    let entry_size = U32_SIZE + max_key_size + max_value_size + U32_SIZE;
+    let child_size = Address::size();
+
+    node_header_size
+        + Bytes::from(CAPACITY as u64) * entry_size
+        + Bytes::from((CAPACITY + 1) as u64) * child_size
+}