feat: v2 of BTreeMap::Node to support unbounded types.

Introduces V2 of `BTreeMap::Node`, which has a smaller memory footprint
than V1 and includes support unbounded types.

The implementation is currently not very efficient, as it deserializes
all entries on load and serializes all entries on save. Performance
enhancements will be included in subsequent PRs.

Note that V2 here is only used in tests and production continues to use
V1.

src/btreemap.rs

@@ -1046,8 +1046,10 @@ where
         Node::load(
             address,
             self.memory(),
-            self.max_key_size,
-            self.max_value_size,
+            node::Version::V1 {
+                max_key_size: self.max_key_size,
+                max_value_size: self.max_value_size,
+            },
         )
     }
 

src/btreemap/node.rs

@@ -11,6 +11,10 @@ use std::cell::{Ref, RefCell};
 mod tests;
 mod v1;
 
+// V2 nodes are currently only used in tests.
+#[allow(dead_code)]
+mod v2;
+
 // The minimum degree to use in the btree.
 // This constant is taken from Rust's std implementation of BTreeMap.
 const B: usize = 6;
@@ -33,7 +37,13 @@ pub enum NodeType {
 pub type Entry<K> = (K, Vec<u8>);
 
 /// A node of a B-Tree.
-/// See `v1.rs` for more details on the memory layout.
+///
+/// There are two versions of a `Node`:
+///
+/// 1. `V1`, which supports only bounded types.
+/// 2. `V2`, which supports both bounded and unbounded types.
+///
+/// See `v1.rs` and `v2.rs` for more details.
 #[derive(Debug)]
 pub struct Node<K: Storable + Ord + Clone> {
     address: Address,
@@ -45,8 +55,9 @@ pub struct Node<K: Storable + Ord + Clone> {
     // child of this key and children[I + 1] points to the right child.
     children: Vec<Address>,
     node_type: NodeType,
-    max_key_size: u32,
-    max_value_size: u32,
+    version: Version,
+    #[allow(dead_code)]
+    overflow: Option<Address>,
 }
 
 impl<K: Storable + Ord + Clone> Node<K> {
@@ -61,14 +72,16 @@ impl<K: Storable + Ord + Clone> Node<K> {
     }
 
     /// Loads a node from memory at the given address.
-    pub fn load<M: Memory>(
-        address: Address,
-        memory: &M,
-        max_key_size: u32,
-        max_value_size: u32,
-    ) -> Self {
-        // NOTE: new versions of `Node` will be introduced.
-        Self::load_v1(address, max_key_size, max_value_size, memory)
+    pub fn load<M: Memory>(address: Address, memory: &M, version: Version) -> Self {
+        if let Version::V1 {
+            max_key_size,
+            max_value_size,
+        } = version
+        {
+            Self::load_v1(address, max_key_size, max_value_size, memory)
+        } else {
+            unreachable!("Only v1 is currently supported.");
+        }
     }
 
     /// Saves the node to memory.
@@ -103,8 +116,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
                         .last()
                         .expect("An internal node must have children."),
                     memory,
-                    self.max_key_size,
-                    self.max_value_size,
+                    self.version,
                 );
                 last_child.get_max(memory)
             }
@@ -123,8 +135,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
                     // NOTE: an internal node must have children, so this access is safe.
                     self.children[0],
                     memory,
-                    self.max_key_size,
-                    self.max_value_size,
+                    self.version,
                 );
                 first_child.get_min(memory)
             }
@@ -413,3 +424,61 @@ enum Value {
     // The value's offset in the node.
     ByRef(Bytes),
 }
+
+/// Stores version-specific data.
+#[derive(Debug, PartialEq, Copy, Clone, Eq)]
+pub enum Version {
+    /// V1 nodes must have bounds for the keys and values.
+    V1 {
+        max_key_size: u32,
+        max_value_size: u32,
+    },
+    /// V2 nodes have a fixed page size.
+    V2(PageSize),
+}
+
+impl Version {
+    fn page_size(&self) -> u32 {
+        match self {
+            Self::V2(page_size) => page_size.get(),
+            Self::V1 {
+                max_key_size,
+                max_value_size,
+            } => {
+                // Page size can be computed from the max key/value sizes.
+                v1::size_v1(*max_key_size, *max_value_size).get() as u32
+            }
+        }
+    }
+}
+
+/// The size of an individual page in the memory where nodes are stored.
+/// A node, if it's bigger than a single page, overflows into multiple pages.
+#[allow(dead_code)]
+#[derive(Debug, PartialEq, Copy, Clone, Eq)]
+pub enum PageSize {
+    /// The page size is derived from the maximum sizes of the keys and values.
+    ///
+    /// Derived page sizes are used when migrating nodes from v1 to v2.
+    /// A migration from v1 nodes to v2 is done incrementally. Children of a v2 node
+    /// may be a v1 node, and keeping the maximum sizes around is necessary to store
+    /// to be able to load v1 nodes.
+    Derived {
+        max_key_size: u32,
+        max_value_size: u32,
+    },
+
+    Value(u32),
+}
+
+impl PageSize {
+    fn get(&self) -> u32 {
+        match self {
+            Self::Value(page_size) => *page_size,
+            Self::Derived {
+                max_key_size,
+                max_value_size,
+            } => v1::size_v1(*max_key_size, *max_value_size).get() as u32,
+        }
+    }
+}

src/btreemap/node/tests.rs

@@ -1,4 +1,5 @@
 use super::*;
+use crate::btreemap::Allocator;
 use proptest::collection::btree_map as pmap;
 use proptest::collection::vec as pvec;
 use std::cell::RefCell;
@@ -66,6 +67,46 @@ impl NodeV1Data {
     }
 }
 
+#[derive(Arbitrary, Debug)]
+struct NodeV2Data {
+    #[strategy(128..10_000_u32)]
+    page_size: u32,
+    #[strategy(
+        pmap(
+            pvec(0..u8::MAX, 0..1000),
+            pvec(0..u8::MAX, 0..1000),
+            1..CAPACITY
+        )
+    )]
+    entries: BTreeMap<Vec<u8>, Vec<u8>>,
+    node_type: NodeType,
+}
+
+impl NodeV2Data {
+    fn get(&self, address: Address) -> Node<Vec<u8>> {
+        let mut node = Node::new_v2(address, self.node_type, PageSize::Value(self.page_size));
+        for entry in self.entries.clone().into_iter() {
+            node.push_entry(entry);
+        }
+        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 + 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();
@@ -89,3 +130,66 @@ fn saving_and_loading_v1_preserves_data(node_data: NodeV1Data) {
         node_data.entries.into_iter().collect::<Vec<_>>()
     );
 }
+
+#[proptest]
+fn saving_and_loading_v2(node_data: NodeV2Data) {
+    let mem = make_memory();
+    let allocator_addr = Address::from(0);
+    let mut allocator = Allocator::new(
+        mem.clone(),
+        allocator_addr,
+        Bytes::from(node_data.page_size as u64),
+    );
+
+    // Create a new node and save it into memory.
+    let node_addr = allocator.allocate();
+    let node = node_data.get(node_addr);
+    node.save_v2(&mut allocator);
+
+    // Reload the node and double check all the entries and children are correct.
+    let node = Node::load_v2(node_addr, PageSize::Value(node_data.page_size), &mem);
+
+    assert_eq!(node.children, node_data.children());
+    assert_eq!(
+        node.entries(&mem),
+        node_data.entries.into_iter().collect::<Vec<_>>()
+    );
+}
+
+#[proptest]
+fn migrating_v1_nodes_to_v2(node_data: NodeV1Data) {
+    let v1_size = v1::size_v1(node_data.max_key_size, node_data.max_value_size);
+    let mem = make_memory();
+    let allocator_addr = Address::from(0);
+    let mut allocator = Allocator::new(mem.clone(), allocator_addr, v1_size);
+
+    // Create a v1 node and save it into memory as v1.
+    let node_addr = allocator.allocate();
+    let node = node_data.get(node_addr);
+    node.save_v1(allocator.memory());
+
+    // Reload the v1 node and save it as v2.
+    let node = Node::<Vec<u8>>::load_v1(
+        node_addr,
+        node_data.max_key_size,
+        node_data.max_value_size,
+        allocator.memory(),
+    );
+    node.save_v2(&mut allocator);
+
+    // Reload the now v2 node and double check all the entries and children are correct.
+    let node = Node::load_v2(
+        node_addr,
+        PageSize::Derived {
+            max_key_size: node_data.max_key_size,
+            max_value_size: node_data.max_value_size,
+        },
+        &mem,
+    );
+
+    assert_eq!(node.children, node_data.children());
+    assert_eq!(
+        node.entries(&mem),
+        node_data.entries.into_iter().collect::<Vec<_>>()
+    );
+}

src/btreemap/node/v1.rs

@@ -49,11 +49,14 @@ impl<K: Storable + Ord + Clone> Node<K> {
         Node {
             address,
             node_type,
-            max_key_size,
-            max_value_size,
             keys: vec![],
             encoded_values: RefCell::default(),
             children: vec![],
+            version: Version::V1 {
+                max_key_size,
+                max_value_size,
+            },
+            overflow: None,
         }
     }
 
@@ -114,8 +117,11 @@ impl<K: Storable + Ord + Clone> Node<K> {
                 INTERNAL_NODE_TYPE => NodeType::Internal,
                 other => unreachable!("Unknown node type {}", other),
             },
-            max_key_size,
-            max_value_size,
+            version: Version::V1 {
+                max_key_size,
+                max_value_size,
+            },
+            overflow: None,
         }
     }
 
@@ -135,6 +141,14 @@ impl<K: Storable + Ord + Clone> Node<K> {
         // Assert entries are sorted in strictly increasing order.
         assert!(self.keys.windows(2).all(|e| e[0] < e[1]));
 
+        let (max_key_size, max_value_size) = match self.version {
+            Version::V1 {
+                max_key_size,
+                max_value_size,
+            } => (max_key_size, max_value_size),
+            Version::V2 { .. } => unreachable!("cannot save v2 node as v1."),
+        };
+
         let header = NodeHeader {
             magic: *MAGIC,
             version: LAYOUT_VERSION,
@@ -164,7 +178,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
 
             // Write the key.
             write(memory, (self.address + offset).get(), key_bytes.borrow());
-            offset += Bytes::from(self.max_key_size);
+            offset += Bytes::from(max_key_size);
 
             // Write the size of the value.
             let value = self.value(idx, memory);
@@ -173,7 +187,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
 
             // Write the value.
             write(memory, (self.address + offset).get(), &value);
-            offset += Bytes::from(self.max_value_size);
+            offset += Bytes::from(max_value_size);
         }
 
         // Write the children