//! A key/value store based on a B-Tree.
mod allocator;
mod iter;
mod node;
use crate::{
read_struct,
types::{Address, Bytes, NULL},
write_struct, Memory, Storable,
};
use allocator::Allocator;
pub use iter::Iter;
use iter::{Cursor, Index};
use node::{Entry, Node, NodeType, B};
use std::marker::PhantomData;
const LAYOUT_VERSION: u8 = 1;
const MAGIC: &[u8; 3] = b"BTR";
/// A "stable" map based on a B-tree.
///
/// The implementation is based on the algorithm outlined in "Introduction to Algorithms"
/// by Cormen et al.
pub struct BTreeMap<M: Memory, K: Storable, V: Storable> {
// The address of the root node. If a root node doesn't exist, the address
// is set to NULL.
root_addr: Address,
// The maximum size a key can have.
max_key_size: u32,
// The maximum size a value can have.
max_value_size: u32,
// An allocator used for managing memory and allocating nodes.
allocator: Allocator<M>,
// The number of elements in the map.
length: u64,
memory: M,
// A marker to communicate to the Rust compiler that we own these types.
_phantom: PhantomData<(K, V)>,
}
#[repr(packed)]
struct BTreeHeader {
magic: [u8; 3],
version: u8,
max_key_size: u32,
max_value_size: u32,
root_addr: Address,
length: u64,
// Additional space reserved to add new fields without breaking backward-compatibility.
_buffer: [u8; 24],
}
impl BTreeHeader {
fn size() -> Bytes {
Bytes::from(core::mem::size_of::<Self>() as u64)
}
}
impl<M: Memory + Clone, K: Storable, V: Storable> BTreeMap<M, K, V> {
/// Initializes a `BTreeMap`.
///
/// If the memory provided already contains a `BTreeMap`, then that
/// map is loaded. Otherwise, a new `BTreeMap` instance is created.
pub fn init(memory: M, max_key_size: u32, max_value_size: u32) -> Self {
if memory.size() == 0 {
// Memory is empty. Create a new map.
return BTreeMap::new(memory, max_key_size, max_value_size);
}
// Check if the magic in the memory corresponds to a BTreeMap.
let mut dst = vec![0; 3];
memory.read(0, &mut dst);
if dst != MAGIC {
// No BTreeMap found. Create a new instance.
BTreeMap::new(memory, max_key_size, max_value_size)
} else {
// The memory already contains a BTreeMap. Load it.
BTreeMap::load(memory)
}
}
/// Creates a new instance a `BTreeMap`.
///
/// The given `memory` is assumed to be exclusively reserved for this data
/// structure and that it starts at address zero. Typically `memory` will
/// be an instance of `RestrictedMemory`.
///
/// When initialized, the data structure has the following memory layout:
///
/// | BTreeHeader | Allocator | ... free memory for nodes |
///
/// See `Allocator` for more details on its own memory layout.
pub fn new(memory: M, max_key_size: u32, max_value_size: u32) -> Self {
// Because we assume that we have exclusive access to the memory,
// we can store the `BTreeHeader` at address zero, and the allocator is
// stored directly after the `BTreeHeader`.
let allocator_addr = Address::from(0) + BTreeHeader::size();
let btree = Self {
memory: memory.clone(),
root_addr: NULL,
allocator: Allocator::new(
memory,
allocator_addr,
Node::size(max_key_size, max_value_size),
),
max_key_size,
max_value_size,
length: 0,
_phantom: PhantomData,
};
btree.save();
btree
}
/// Loads the map from memory.
pub fn load(memory: M) -> Self {
// Read the header from memory.
let header: BTreeHeader = read_struct(Address::from(0), &memory);
assert_eq!(&header.magic, MAGIC, "Bad magic.");
assert_eq!(header.version, LAYOUT_VERSION, "Unsupported version.");
let allocator_addr = Address::from(0) + BTreeHeader::size();
Self {
memory: memory.clone(),
root_addr: header.root_addr,
allocator: Allocator::load(memory, allocator_addr),
max_key_size: header.max_key_size,
max_value_size: header.max_value_size,
length: header.length,
_phantom: PhantomData,
}
}
/// Inserts a key-value pair into the map.
///
/// The previous value of the key, if present, is returned.
///
/// The size of the key/value must be <= the max key/value sizes configured
/// for the map. Otherwise, an `InsertError` is returned.
pub fn insert(&mut self, key: K, value: V) -> Result<Option<V>, InsertError> {
let key = key.to_bytes();
let value = value.to_bytes();
// Verify the size of the key.
if key.len() > self.max_key_size as usize {
return Err(InsertError::KeyTooLarge {
given: key.len(),
max: self.max_key_size as usize,
});
}
// Verify the size of the value.
if value.len() > self.max_value_size as usize {
return Err(InsertError::ValueTooLarge {
given: value.len(),
max: self.max_value_size as usize,
});
}
let key = key.to_vec();
let value = value.to_vec();
let root = if self.root_addr == NULL {
// No root present. Allocate one.
let node = self.allocate_node(NodeType::Leaf);
self.root_addr = node.address;
self.save();
node
} else {
// Load the root from memory.
let mut root = self.load_node(self.root_addr);
// Check if the key already exists in the root.
if let Ok(idx) = root.get_key_idx(&key) {
// The key exists. Overwrite it and return the previous value.
let (_, previous_value) = root.swap_entry(idx, (key, value));
root.save(&self.memory);
return Ok(Some(V::from_bytes(previous_value)));
}
// If the root is full, we need to introduce a new node as the root.
//
// NOTE: In the case where we are overwriting an existing key, then introducing
// a new root node isn't strictly necessary. However, that's a micro-optimization
// that adds more complexity than it's worth.
if root.is_full() {
// The root is full. Allocate a new node that will be used as the new root.
let mut new_root = self.allocate_node(NodeType::Internal);
// The new root has the old root as its only child.
new_root.children.push(self.root_addr);
// Update the root address.
self.root_addr = new_root.address;
self.save();
// Split the old (full) root.
self.split_child(&mut new_root, 0);
new_root
} else {
root
}
};
Ok(self.insert_nonfull(root, key, value).map(V::from_bytes))
}
// Inserts an entry into a node that is *not full*.
fn insert_nonfull(&mut self, mut node: Node, key: Vec<u8>, value: Vec<u8>) -> Option<Vec<u8>> {
// We're guaranteed by the caller that the provided node is not full.
assert!(!node.is_full());
// Look for the key in the node.
match node.entries.binary_search_by(|e| e.0.cmp(&key)) {
Ok(idx) => {
// The key is already in the node.
// Overwrite it and return the previous value.
let (_, previous_value) = node.swap_entry(idx, (key, value));
node.save(&self.memory);
Some(previous_value)
}
Err(idx) => {
// The key isn't in the node. `idx` is where that key should be inserted.
match node.node_type {
NodeType::Leaf => {
// The node is a non-full leaf.
// Insert the entry at the proper location.
node.entries.insert(idx, (key, value));
node.save(&self.memory);
// Update the length.
self.length += 1;
self.save();
// No previous value to return.
None
}
NodeType::Internal => {
// The node is an internal node.
// Load the child that we should add the entry to.
let mut child = self.load_node(node.children[idx]);
if child.is_full() {
// Check if the key already exists in the child.
if let Ok(idx) = child.get_key_idx(&key) {
// The key exists. Overwrite it and return the previous value.
let (_, previous_value) = child.swap_entry(idx, (key, value));
child.save(&self.memory);
return Some(previous_value);
}
// The child is full. Split the child.
self.split_child(&mut node, idx);
// The children have now changed. Search again for
// the child where we need to store the entry in.
let idx = node.get_key_idx(&key).unwrap_or_else(|idx| idx);
child = self.load_node(node.children[idx]);
}
// The child should now be not full.
assert!(!child.is_full());
self.insert_nonfull(child, key, value)
}
}
}
}
}
// Takes as input a nonfull internal `node` and index to its full child, then
// splits this child into two, adding an additional child to `node`.
//
// Example:
//
// [ ... M Y ... ]
// |
// [ N O P Q R S T U V W X ]
//
//
// After splitting becomes:
//
// [ ... M S Y ... ]
// / \
// [ N O P Q R ] [ T U V W X ]
//
fn split_child(&mut self, node: &mut Node, full_child_idx: usize) {
// The node must not be full.
assert!(!node.is_full());
// The node's child must be full.
let mut full_child = self.load_node(node.children[full_child_idx]);
assert!(full_child.is_full());
// Create a sibling to this full child (which has to be the same type).
let mut sibling = self.allocate_node(full_child.node_type);
assert_eq!(sibling.node_type, full_child.node_type);
// Move the values above the median into the new sibling.
sibling.entries = full_child.entries.split_off(B as usize);
if full_child.node_type == NodeType::Internal {
sibling.children = full_child.children.split_off(B as usize);
}
// Add sibling as a new child in the node.
node.children.insert(full_child_idx + 1, sibling.address);
// Move the median entry into the node.
let (median_key, median_value) = full_child
.entries
.pop()
.expect("A full child cannot be empty");
node.entries
.insert(full_child_idx, (median_key, median_value));
sibling.save(&self.memory);
full_child.save(&self.memory);
node.save(&self.memory);
}
/// Returns the value associated with the given key if it exists.
pub fn get(&self, key: &K) -> Option<V> {
if self.root_addr == NULL {
return None;
}
self.get_helper(self.root_addr, &key.to_bytes())
.map(V::from_bytes)
}
fn get_helper(&self, node_addr: Address, key: &[u8]) -> Option<Vec<u8>> {
let node = self.load_node(node_addr);
match node.entries.binary_search_by(|e| e.0.as_slice().cmp(key)) {
Ok(idx) => Some(node.entries[idx].1.clone()),
Err(idx) => {
match node.node_type {
NodeType::Leaf => None, // Key not found.
NodeType::Internal => {
// The key isn't in the node. Look for the key in the child.
self.get_helper(node.children[idx], key)
}
}
}
}
}
/// Returns `true` if the key exists in the map, `false` otherwise.
pub fn contains_key(&self, key: &K) -> bool {
self.get(key).is_some()
}
/// Returns `true` if the map contains no elements.
pub fn is_empty(&self) -> bool {
self.length == 0
}
/// Returns the number of elements in the map.
pub fn len(&self) -> u64 {
self.length
}
/// Returns a reference to the memory used by the map.
pub fn get_memory(&self) -> M {
self.memory.clone()
}
/// Removes a key from the map, returning the previous value at the key if it exists.
pub fn remove(&mut self, key: &K) -> Option<V> {
if self.root_addr == NULL {
return None;
}
self.remove_helper(self.root_addr, &key.to_bytes())
.map(V::from_bytes)
}
// A helper method for recursively removing a key from the B-tree.
fn remove_helper(&mut self, node_addr: Address, key: &[u8]) -> Option<Vec<u8>> {
let mut node = self.load_node(node_addr);
if node.address != self.root_addr {
// We're guaranteed that whenever this method is called the number
// of keys is >= `B`. Note that this is higher than the minimum required
// in a node, which is `B - 1`, and that's because this strengthened
// condition allows us to delete an entry in a single pass most of the
// time without having to back up.
assert!(node.entries.len() >= B as usize);
}
match node.node_type {
NodeType::Leaf => {
match node.entries.binary_search_by(|e| e.0.as_slice().cmp(key)) {
Ok(idx) => {
// Case 1: The node is a leaf node and the key exists in it.
// This is the simplest case. The key is removed from the leaf.
let value = node.entries.remove(idx).1;
self.length -= 1;
if node.entries.is_empty() {
assert_eq!(
node.address, self.root_addr,
"Removal can only result in an empty leaf node if that node is the root"
);
// Deallocate the empty node.
self.allocator.deallocate(node.address);
self.root_addr = NULL;
} else {
node.save(&self.memory);
}
self.save();
Some(value)
}
_ => None, // Key not found.
}
}
NodeType::Internal => {
match node.entries.binary_search_by(|e| e.0.as_slice().cmp(key)) {
Ok(idx) => {
// Case 2: The node is an internal node and the key exists in it.
// Check if the child that precedes `key` has at least `B` keys.
let left_child = self.load_node(node.children[idx]);
if left_child.entries.len() >= B as usize {
// Case 2.a: The node's left child has >= `B` keys.
//
// parent
// [..., key, ...]
// / \
// [left child] [...]
// / \
// [...] [..., key predecessor]
//
// In this case, we replace `key` with the key's predecessor from the
// left child's subtree, then we recursively delete the key's
// predecessor for the following end result:
//
// parent
// [..., key predecessor, ...]
// / \
// [left child] [...]
// / \
// [...] [...]
// Recursively delete the predecessor.
// TODO(EXC-1034): Do this in a single pass.
let predecessor = left_child.get_max(&self.memory);
self.remove_helper(node.children[idx], &predecessor.0)?;
// Replace the `key` with its predecessor.
let (_, old_value) = node.swap_entry(idx, predecessor);
// Save the parent node.
node.save(&self.memory);
return Some(old_value);
}
// Check if the child that succeeds `key` has at least `B` keys.
let right_child = self.load_node(node.children[idx + 1]);
if right_child.entries.len() >= B as usize {
// Case 2.b: The node's right child has >= `B` keys.
//
// parent
// [..., key, ...]
// / \
// [...] [right child]
// / \
// [key successor, ...] [...]
//
// In this case, we replace `key` with the key's successor from the
// right child's subtree, then we recursively delete the key's
// successor for the following end result:
//
// parent
// [..., key successor, ...]
// / \
// [...] [right child]
// / \
// [...] [...]
// Recursively delete the successor.
// TODO(EXC-1034): Do this in a single pass.
let successor = right_child.get_min(&self.memory);
self.remove_helper(node.children[idx + 1], &successor.0)?;
// Replace the `key` with its successor.
let (_, old_value) = node.swap_entry(idx, successor);
// Save the parent node.
node.save(&self.memory);
return Some(old_value);
}
// Case 2.c: Both the left child and right child have B - 1 keys.
//
// parent
// [..., key, ...]
// / \
// [left child] [right child]
//
// In this case, we merge (left child, key, right child) into a single
// node of size 2B - 1. The result will look like this:
//
// parent
// [... ...]
// |
// [left child, `key`, right child] <= new child
//
// We then recurse on this new child to delete `key`.
//
// If `parent` becomes empty (which can only happen if it's the root),
// then `parent` is deleted and `new_child` becomes the new root.
assert_eq!(left_child.entries.len(), B as usize - 1);
assert_eq!(right_child.entries.len(), B as usize - 1);
// Merge the right child into the left child.
let new_child =
self.merge(right_child, left_child, node.entries.remove(idx));
// Remove the right child from the parent node.
node.children.remove(idx + 1);
if node.entries.is_empty() {
// Can only happen if this node is root.
assert_eq!(node.address, self.root_addr);
assert_eq!(node.children, vec![new_child.address]);
self.root_addr = new_child.address;
// Deallocate the root node.
self.allocator.deallocate(node.address);
self.save();
}
node.save(&self.memory);
new_child.save(&self.memory);
// Recursively delete the key.
self.remove_helper(new_child.address, key)
}
Err(idx) => {
// Case 3: The node is an internal node and the key does NOT exist in it.
// If the key does exist in the tree, it will exist in the subtree at index
// `idx`.
let mut child = self.load_node(node.children[idx]);
if child.entries.len() >= B as usize {
// The child has enough nodes. Recurse to delete the `key` from the
// `child`.
return self.remove_helper(node.children[idx], key);
}
// The child has < `B` keys. Let's see if it has a sibling with >= `B` keys.
let mut left_sibling = if idx > 0 {
Some(self.load_node(node.children[idx - 1]))
} else {
None
};
let mut right_sibling = if idx + 1 < node.children.len() {
Some(self.load_node(node.children[idx + 1]))
} else {
None
};
if let Some(ref mut left_sibling) = left_sibling {
if left_sibling.entries.len() >= B as usize {
// Case 3.a (left): The child has a left sibling with >= `B` keys.
//
// [d] (parent)
// / \
// (left sibling) [a, b, c] [e, f] (child)
// \
// [c']
//
// In this case, we move a key down from the parent into the child
// and move a key from the left sibling up into the parent
// resulting in the following tree:
//
// [c] (parent)
// / \
// (left sibling) [a, b] [d, e, f] (child)
// /
// [c']
//
// We then recurse to delete the key from the child.
// Remove the last entry from the left sibling.
let (left_sibling_key, left_sibling_value) =
left_sibling.entries.pop().unwrap();
// Replace the parent's entry with the one from the left sibling.
let (parent_key, parent_value) = node
.swap_entry(idx - 1, (left_sibling_key, left_sibling_value));
// Move the entry from the parent into the child.
child.entries.insert(0, (parent_key, parent_value));
// Move the last child from left sibling into child.
if let Some(last_child) = left_sibling.children.pop() {
assert_eq!(left_sibling.node_type, NodeType::Internal);
assert_eq!(child.node_type, NodeType::Internal);
child.children.insert(0, last_child);
} else {
assert_eq!(left_sibling.node_type, NodeType::Leaf);
assert_eq!(child.node_type, NodeType::Leaf);
}
left_sibling.save(&self.memory);
child.save(&self.memory);
node.save(&self.memory);
return self.remove_helper(child.address, key);
}
}
if let Some(right_sibling) = &mut right_sibling {
if right_sibling.entries.len() >= B as usize {
// Case 3.a (right): The child has a right sibling with >= `B` keys.
//
// [c] (parent)
// / \
// (child) [a, b] [d, e, f] (left sibling)
// /
// [d']
//
// In this case, we move a key down from the parent into the child
// and move a key from the right sibling up into the parent
// resulting in the following tree:
//
// [d] (parent)
// / \
// (child) [a, b, c] [e, f] (right sibling)
// \
// [d']
//
// We then recurse to delete the key from the child.
// Remove the first entry from the right sibling.
let (right_sibling_key, right_sibling_value) =
right_sibling.entries.remove(0);
// Replace the parent's entry with the one from the right sibling.
let parent_entry =
node.swap_entry(idx, (right_sibling_key, right_sibling_value));
// Move the entry from the parent into the child.
child.entries.push(parent_entry);
// Move the first child of right_sibling into `child`.
match right_sibling.node_type {
NodeType::Internal => {
assert_eq!(child.node_type, NodeType::Internal);
child.children.push(right_sibling.children.remove(0));
}
NodeType::Leaf => {
assert_eq!(child.node_type, NodeType::Leaf);
}
}
right_sibling.save(&self.memory);
child.save(&self.memory);
node.save(&self.memory);
return self.remove_helper(child.address, key);
}
}
// Case 3.b: neither siblings of the child have >= `B` keys.
if let Some(left_sibling) = left_sibling {
// Merge child into left sibling if it exists.
let left_sibling_address = left_sibling.address;
self.merge(child, left_sibling, node.entries.remove(idx - 1));
// Removing child from parent.
node.children.remove(idx);
if node.entries.is_empty() {
self.allocator.deallocate(node.address);
if node.address == self.root_addr {
// Update the root.
self.root_addr = left_sibling_address;
self.save();
}
} else {
node.save(&self.memory);
}
return self.remove_helper(left_sibling_address, key);
}
if let Some(right_sibling) = right_sibling {
// Merge child into right sibling.
let right_sibling_address = right_sibling.address;
self.merge(child, right_sibling, node.entries.remove(idx));
// Removing child from parent.
node.children.remove(idx);
if node.entries.is_empty() {
self.allocator.deallocate(node.address);
if node.address == self.root_addr {
// Update the root.
self.root_addr = right_sibling_address;
self.save();
}
} else {
node.save(&self.memory);
}
return self.remove_helper(right_sibling_address, key);
}
unreachable!("At least one of the siblings must exist.");
}
}
}
}
}
/// Returns an iterator over the entries of the map, sorted by key.
pub fn iter(&self) -> Iter<M, K, V> {
Iter::new(self)
}
/// Returns an iterator over the entries in the map where keys begin with the given `prefix`.
/// If the optional `offset` is set, the iterator returned will start from the entry that
/// contains this `offset` (while still iterating over all remaining entries that begin
/// with the given `prefix`).
pub fn range(&self, prefix: Vec<u8>, offset: Option<Vec<u8>>) -> Iter<M, K, V> {
if self.root_addr == NULL {
// Map is empty.
return Iter::null(self);
}
let mut node = self.load_node(self.root_addr);
let mut cursors = vec![];
loop {
// Look for the prefix in the node.
let mut pivot = prefix.clone();
if let Some(offset) = &offset {
pivot.extend_from_slice(offset);
}
match node.entries.binary_search_by(|e| e.0.cmp(&pivot)) {
Ok(idx) | Err(idx) => {
// If `prefix` is a key in the node, then `idx` would return its
// location. Otherwise, `idx` is the location of the next key in
// lexicographical order.
// Load the next child of the node to visit if it exists.
// This is done first to avoid cloning the node.
let child = match node.node_type {
NodeType::Internal => {
// Note that loading a child node cannot fail since
// len(children) = len(entries) + 1
Some(self.load_node(node.children[idx]))
}
NodeType::Leaf => None,
};
// If the prefix is found in the node, then add a cursor starting from its index.
if idx < node.entries.len() && node.entries[idx].0.starts_with(&prefix) {
cursors.push(Cursor::Node {
node,
next: Index::Entry(idx),
});
}
match child {
None => {
// Leaf node. Return an iterator with the found cursors.
match offset {
Some(offset) => {
return Iter::new_with_prefix_and_offset(
self, prefix, offset, cursors,
);
}
None => {
return Iter::new_with_prefix(self, prefix, cursors);
}
}
}
Some(child) => {
// Iterate over the child node.
node = child;
}
}
}
}
}
}
// Merges one node (`source`) into another (`into`), along with a median entry.
//
// Example (values are not included for brevity):
//
// Input:
// Source: [1, 2, 3]
// Into: [5, 6, 7]
// Median: 4
//
// Output:
// [1, 2, 3, 4, 5, 6, 7] (stored in the `into` node)
// `source` is deallocated.
fn merge(&mut self, source: Node, into: Node, median: Entry) -> Node {
assert_eq!(source.node_type, into.node_type);
assert!(!source.entries.is_empty());
assert!(!into.entries.is_empty());
let into_address = into.address;
let source_address = source.address;
// Figure out which node contains lower values than the other.
let (mut lower, mut higher) = if source.entries[0].0 < into.entries[0].0 {
(source, into)
} else {
(into, source)
};
lower.entries.push(median);
lower.entries.append(&mut higher.entries);
lower.address = into_address;
// Move the children (if any exist).
lower.children.append(&mut higher.children);
lower.save(&self.memory);
self.allocator.deallocate(source_address);
lower
}
fn allocate_node(&mut self, node_type: NodeType) -> Node {
Node {
address: self.allocator.allocate(),
entries: vec![],
children: vec![],
node_type,
max_key_size: self.max_key_size,
max_value_size: self.max_value_size,
}
}
fn load_node(&self, address: Address) -> Node {
Node::load(
address,
&self.memory,
self.max_key_size,
self.max_value_size,
)
}
// Saves the map to memory.
fn save(&self) {
let header = BTreeHeader {
magic: *MAGIC,
version: LAYOUT_VERSION,
root_addr: self.root_addr,
max_key_size: self.max_key_size,
max_value_size: self.max_value_size,
length: self.length,
_buffer: [0; 24],
};
write_struct(&header, Address::from(0), &self.memory);
}
}
/// An error returned when inserting entries into the map.
#[derive(Debug, PartialEq, Eq)]
pub enum InsertError {
KeyTooLarge { given: usize, max: usize },
ValueTooLarge { given: usize, max: usize },
}
impl std::fmt::Display for InsertError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::KeyTooLarge { given, max } => {
write!(
f,
"InsertError::KeyTooLarge Expected key to be <= {} bytes but received key with {} bytes.",
max, given
)
}
Self::ValueTooLarge { given, max } => {
write!(
f,
"InsertError::ValueTooLarge Expected value to be <= {} bytes but received value with {} bytes.",
max, given
)
}
}
}
}
#[cfg(test)]
mod test {
use super::*;
use crate::btreemap::node::CAPACITY;
use std::cell::RefCell;
use std::rc::Rc;
fn make_memory() -> Rc<RefCell<Vec<u8>>> {
Rc::new(RefCell::new(Vec::new()))
}
// A helper method to succinctly create an entry.
fn e(x: u8) -> (Vec<u8>, Vec<u8>) {
(vec![x], vec![])
}
#[test]
fn init_preserves_data() {
let mem = make_memory();
let mut btree = BTreeMap::init(mem.clone(), 3, 4);
assert_eq!(btree.insert(vec![1, 2, 3], vec![4, 5, 6]), Ok(None));
assert_eq!(btree.get(&vec![1, 2, 3]), Some(vec![4, 5, 6]));
// Reload the btree
let btree = BTreeMap::init(mem, 3, 4);
// Data still exists.
assert_eq!(btree.get(&vec![1, 2, 3]), Some(vec![4, 5, 6]));
}
#[test]
fn insert_get() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 3, 4);
assert_eq!(btree.insert(vec![1, 2, 3], vec![4, 5, 6]), Ok(None));
assert_eq!(btree.get(&vec![1, 2, 3]), Some(vec![4, 5, 6]));
}
#[test]
fn insert_overwrites_previous_value() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
assert_eq!(btree.insert(vec![1, 2, 3], vec![4, 5, 6]), Ok(None));
assert_eq!(
btree.insert(vec![1, 2, 3], vec![7, 8, 9]),
Ok(Some(vec![4, 5, 6]))
);
assert_eq!(btree.get(&vec![1, 2, 3]), Some(vec![7, 8, 9]));
}
#[test]
fn insert_get_multiple() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
assert_eq!(btree.insert(vec![1, 2, 3], vec![4, 5, 6]), Ok(None));
assert_eq!(btree.insert(vec![4, 5], vec![7, 8, 9, 10]), Ok(None));
assert_eq!(btree.insert(vec![], vec![11]), Ok(None));
assert_eq!(btree.get(&vec![1, 2, 3]), Some(vec![4, 5, 6]));
assert_eq!(btree.get(&vec![4, 5]), Some(vec![7, 8, 9, 10]));
assert_eq!(btree.get(&vec![]), Some(vec![11]));
}
#[test]
fn insert_overwrite_median_key_in_full_child_node() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
for i in 1..=17 {
assert_eq!(btree.insert(vec![i], vec![]), Ok(None));
}
// The result should look like this:
// [6]
// / \
// [1, 2, 3, 4, 5] [7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17]
let root = btree.load_node(btree.root_addr);
assert_eq!(root.node_type, NodeType::Internal);
assert_eq!(root.entries, vec![(vec![6], vec![])]);
assert_eq!(root.children.len(), 2);
// The right child should now be full, with the median key being "12"
let right_child = btree.load_node(root.children[1]);
assert!(right_child.is_full());
let median_index = right_child.entries.len() / 2;
assert_eq!(right_child.entries[median_index].0, vec![12]);
// Overwrite the median key.
assert_eq!(btree.insert(vec![12], vec![1, 2, 3]), Ok(Some(vec![])));
// The key is overwritten successfully.
assert_eq!(btree.get(&vec![12]), Some(vec![1, 2, 3]));
// The child has not been split and is still full.
let right_child = btree.load_node(root.children[1]);
assert_eq!(right_child.node_type, NodeType::Leaf);
assert!(right_child.is_full());
}
#[test]
fn insert_overwrite_key_in_full_root_node() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
for i in 1..=11 {
assert_eq!(btree.insert(vec![i], vec![]), Ok(None));
}
// We now have a root that is full and looks like this:
//
// [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
let root = btree.load_node(btree.root_addr);
assert!(root.is_full());
// Overwrite an element in the root. It should NOT cause the node to be split.
assert_eq!(btree.insert(vec![6], vec![4, 5, 6]), Ok(Some(vec![])));
let root = btree.load_node(btree.root_addr);
assert_eq!(root.node_type, NodeType::Leaf);
assert_eq!(btree.get(&vec![6]), Some(vec![4, 5, 6]));
assert_eq!(root.entries.len(), 11);
}
#[test]
fn allocations() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
for i in 0..CAPACITY as u8 {
assert_eq!(btree.insert(vec![i], vec![]), Ok(None));
}
// Only need a single allocation to store up to `CAPACITY` elements.
assert_eq!(btree.allocator.num_allocated_chunks(), 1);
assert_eq!(btree.insert(vec![255], vec![]), Ok(None));
// The node had to be split into three nodes.
assert_eq!(btree.allocator.num_allocated_chunks(), 3);
}
#[test]
fn allocations_2() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
assert_eq!(btree.allocator.num_allocated_chunks(), 0);
assert_eq!(btree.insert(vec![], vec![]), Ok(None));
assert_eq!(btree.allocator.num_allocated_chunks(), 1);
assert_eq!(btree.remove(&vec![]), Some(vec![]));
assert_eq!(btree.allocator.num_allocated_chunks(), 0);
}
#[test]
fn insert_same_key_multiple() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
assert_eq!(btree.insert(vec![1], vec![2]), Ok(None));
for i in 2..10 {
assert_eq!(btree.insert(vec![1], vec![i + 1]), Ok(Some(vec![i])));
}
}
#[test]
fn insert_split_node() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
for i in 1..=11 {
assert_eq!(btree.insert(vec![i], vec![]), Ok(None));
}
// Should now split a node.
assert_eq!(btree.insert(vec![12], vec![]), Ok(None));
// The result should look like this:
// [6]
// / \
// [1, 2, 3, 4, 5] [7, 8, 9, 10, 11, 12]
for i in 1..=12 {
assert_eq!(btree.get(&vec![i]), Some(vec![]));
}
}
#[test]
fn overwrite_test() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
let num_elements: u8 = 255;
// Ensure that the number of elements we insert is significantly
// higher than `CAPACITY` so that we test interesting cases (e.g.
// overwriting the value in an internal node).
assert!(num_elements as u64 > 10 * CAPACITY);
for i in 0..num_elements {
assert_eq!(btree.insert(vec![i], vec![]), Ok(None));
}
// Overwrite the values.
for i in 0..num_elements {
// Assert we retrieved the old value correctly.
assert_eq!(btree.insert(vec![i], vec![1, 2, 3]), Ok(Some(vec![])));
// Assert we retrieved the new value correctly.
assert_eq!(btree.get(&vec![i]), Some(vec![1, 2, 3]));
}
}
#[test]
fn insert_split_multiple_nodes() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
for i in 1..=11 {
assert_eq!(btree.insert(vec![i], vec![]), Ok(None));
}
// Should now split a node.
assert_eq!(btree.insert(vec![12], vec![]), Ok(None));
// The result should look like this:
// [6]
// / \
// [1, 2, 3, 4, 5] [7, 8, 9, 10, 11, 12]
let root = btree.load_node(btree.root_addr);
assert_eq!(root.node_type, NodeType::Internal);
assert_eq!(root.entries, vec![(vec![6], vec![])]);
assert_eq!(root.children.len(), 2);
let child_0 = btree.load_node(root.children[0]);
assert_eq!(child_0.node_type, NodeType::Leaf);
assert_eq!(
child_0.entries,
vec![
(vec![1], vec![]),
(vec![2], vec![]),
(vec![3], vec![]),
(vec![4], vec![]),
(vec![5], vec![])
]
);
let child_1 = btree.load_node(root.children[1]);
assert_eq!(child_1.node_type, NodeType::Leaf);
assert_eq!(
child_1.entries,
vec![
(vec![7], vec![]),
(vec![8], vec![]),
(vec![9], vec![]),
(vec![10], vec![]),
(vec![11], vec![]),
(vec![12], vec![])
]
);
for i in 1..=12 {
assert_eq!(btree.get(&vec![i]), Some(vec![]));
}
// Insert more to cause more splitting.
assert_eq!(btree.insert(vec![13], vec![]), Ok(None));
assert_eq!(btree.insert(vec![14], vec![]), Ok(None));
assert_eq!(btree.insert(vec![15], vec![]), Ok(None));
assert_eq!(btree.insert(vec![16], vec![]), Ok(None));
assert_eq!(btree.insert(vec![17], vec![]), Ok(None));
// Should cause another split
assert_eq!(btree.insert(vec![18], vec![]), Ok(None));
for i in 1..=18 {
assert_eq!(btree.get(&vec![i]), Some(vec![]));
}
let root = btree.load_node(btree.root_addr);
assert_eq!(root.node_type, NodeType::Internal);
assert_eq!(root.entries, vec![(vec![6], vec![]), (vec![12], vec![])]);
assert_eq!(root.children.len(), 3);
let child_0 = btree.load_node(root.children[0]);
assert_eq!(child_0.node_type, NodeType::Leaf);
assert_eq!(
child_0.entries,
vec![
(vec![1], vec![]),
(vec![2], vec![]),
(vec![3], vec![]),
(vec![4], vec![]),
(vec![5], vec![])
]
);
let child_1 = btree.load_node(root.children[1]);
assert_eq!(child_1.node_type, NodeType::Leaf);
assert_eq!(
child_1.entries,
vec![
(vec![7], vec![]),
(vec![8], vec![]),
(vec![9], vec![]),
(vec![10], vec![]),
(vec![11], vec![]),
]
);
let child_2 = btree.load_node(root.children[2]);
assert_eq!(child_2.node_type, NodeType::Leaf);
assert_eq!(
child_2.entries,
vec![
(vec![13], vec![]),
(vec![14], vec![]),
(vec![15], vec![]),
(vec![16], vec![]),
(vec![17], vec![]),
(vec![18], vec![]),
]
);
}
#[test]
fn remove_simple() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
assert_eq!(btree.insert(vec![1, 2, 3], vec![4, 5, 6]), Ok(None));
assert_eq!(btree.get(&vec![1, 2, 3]), Some(vec![4, 5, 6]));
assert_eq!(btree.remove(&vec![1, 2, 3]), Some(vec![4, 5, 6]));
assert_eq!(btree.get(&vec![1, 2, 3]), None);
}
#[test]
fn remove_case_2a_and_2c() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem.clone(), 5, 5);
for i in 1..=11 {
assert_eq!(btree.insert(vec![i], vec![]), Ok(None));
}
// Should now split a node.
assert_eq!(btree.insert(vec![0], vec![]), Ok(None));
// The result should look like this:
// [6]
// / \
// [0, 1, 2, 3, 4, 5] [7, 8, 9, 10, 11]
for i in 0..=11 {
assert_eq!(btree.get(&vec![i]), Some(vec![]));
}
// Remove node 6. Triggers case 2.a
assert_eq!(btree.remove(&vec![6]), Some(vec![]));
// The result should look like this:
// [5]
// / \
// [0, 1, 2, 3, 4] [7, 8, 9, 10, 11]
let root = btree.load_node(btree.root_addr);
assert_eq!(root.node_type, NodeType::Internal);
assert_eq!(root.entries, vec![e(5)]);
assert_eq!(root.children.len(), 2);
let child_0 = btree.load_node(root.children[0]);
assert_eq!(child_0.node_type, NodeType::Leaf);
assert_eq!(child_0.entries, vec![e(0), e(1), e(2), e(3), e(4)]);
let child_1 = btree.load_node(root.children[1]);
assert_eq!(child_1.node_type, NodeType::Leaf);
assert_eq!(child_1.entries, vec![e(7), e(8), e(9), e(10), e(11)]);
// There are three allocated nodes.
assert_eq!(btree.allocator.num_allocated_chunks(), 3);
// Remove node 5. Triggers case 2c
assert_eq!(btree.remove(&vec![5]), Some(vec![]));
// Reload the btree to verify that we saved it correctly.
let btree = BTreeMap::<_, Vec<u8>, Vec<u8>>::load(mem);
// The result should look like this:
// [0, 1, 2, 3, 4, 7, 8, 9, 10, 11]
let root = btree.load_node(btree.root_addr);
assert_eq!(
root.entries,
vec![e(0), e(1), e(2), e(3), e(4), e(7), e(8), e(9), e(10), e(11)]
);
// There is only one node allocated.
assert_eq!(btree.allocator.num_allocated_chunks(), 1);
}
#[test]
fn remove_case_2b() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
for i in 1..=11 {
assert_eq!(btree.insert(vec![i], vec![]), Ok(None));
}
// Should now split a node.
assert_eq!(btree.insert(vec![12], vec![]), Ok(None));
// The result should look like this:
// [6]
// / \
// [1, 2, 3, 4, 5] [7, 8, 9, 10, 11, 12]
for i in 1..=12 {
assert_eq!(btree.get(&vec![i]), Some(vec![]));
}
// Remove node 6. Triggers case 2.b
assert_eq!(btree.remove(&vec![6]), Some(vec![]));
// The result should look like this:
// [7]
// / \
// [1, 2, 3, 4, 5] [8, 9, 10, 11, 12]
let root = btree.load_node(btree.root_addr);
assert_eq!(root.node_type, NodeType::Internal);
assert_eq!(root.entries, vec![e(7)]);
assert_eq!(root.children.len(), 2);
let child_0 = btree.load_node(root.children[0]);
assert_eq!(child_0.node_type, NodeType::Leaf);
assert_eq!(child_0.entries, vec![e(1), e(2), e(3), e(4), e(5)]);
let child_1 = btree.load_node(root.children[1]);
assert_eq!(child_1.node_type, NodeType::Leaf);
assert_eq!(child_1.entries, vec![e(8), e(9), e(10), e(11), e(12)]);
// Remove node 7. Triggers case 2.c
assert_eq!(btree.remove(&vec![7]), Some(vec![]));
// The result should look like this:
//
// [1, 2, 3, 4, 5, 8, 9, 10, 11, 12]
let root = btree.load_node(btree.root_addr);
assert_eq!(root.node_type, NodeType::Leaf);
assert_eq!(
root.entries,
vec![
e(1),
e(2),
e(3),
e(4),
e(5),
e(8),
e(9),
e(10),
e(11),
e(12)
]
);
}
#[test]
fn remove_case_3a_right() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
for i in 1..=11 {
assert_eq!(btree.insert(vec![i], vec![]), Ok(None));
}
// Should now split a node.
assert_eq!(btree.insert(vec![12], vec![]), Ok(None));
// The result should look like this:
// [6]
// / \
// [1, 2, 3, 4, 5] [7, 8, 9, 10, 11, 12]
// Remove node 3. Triggers case 3.a
assert_eq!(btree.remove(&vec![3]), Some(vec![]));
// The result should look like this:
// [7]
// / \
// [1, 2, 4, 5, 6] [8, 9, 10, 11, 12]
let root = btree.load_node(btree.root_addr);
assert_eq!(root.node_type, NodeType::Internal);
assert_eq!(root.entries, vec![(vec![7], vec![])]);
assert_eq!(root.children.len(), 2);
let child_0 = btree.load_node(root.children[0]);
assert_eq!(child_0.node_type, NodeType::Leaf);
assert_eq!(child_0.entries, vec![e(1), e(2), e(4), e(5), e(6)]);
let child_1 = btree.load_node(root.children[1]);
assert_eq!(child_1.node_type, NodeType::Leaf);
assert_eq!(child_1.entries, vec![e(8), e(9), e(10), e(11), e(12)]);
// There are three allocated nodes.
assert_eq!(btree.allocator.num_allocated_chunks(), 3);
}
#[test]
fn remove_case_3a_left() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
for i in 1..=11 {
assert_eq!(btree.insert(vec![i], vec![]), Ok(None));
}
// Should now split a node.
assert_eq!(btree.insert(vec![0], vec![]), Ok(None));
// The result should look like this:
// [6]
// / \
// [0, 1, 2, 3, 4, 5] [7, 8, 9, 10, 11]
// Remove node 8. Triggers case 3.a left
assert_eq!(btree.remove(&vec![8]), Some(vec![]));
// The result should look like this:
// [5]
// / \
// [0, 1, 2, 3, 4] [6, 7, 9, 10, 11]
let root = btree.load_node(btree.root_addr);
assert_eq!(root.node_type, NodeType::Internal);
assert_eq!(root.entries, vec![(vec![5], vec![])]);
assert_eq!(root.children.len(), 2);
let child_0 = btree.load_node(root.children[0]);
assert_eq!(child_0.node_type, NodeType::Leaf);
assert_eq!(child_0.entries, vec![e(0), e(1), e(2), e(3), e(4)]);
let child_1 = btree.load_node(root.children[1]);
assert_eq!(child_1.node_type, NodeType::Leaf);
assert_eq!(child_1.entries, vec![e(6), e(7), e(9), e(10), e(11)]);
// There are three allocated nodes.
assert_eq!(btree.allocator.num_allocated_chunks(), 3);
}
#[test]
fn remove_case_3b_merge_into_right() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem.clone(), 5, 5);
for i in 1..=11 {
assert_eq!(btree.insert(vec![i], vec![]), Ok(None));
}
// Should now split a node.
assert_eq!(btree.insert(vec![12], vec![]), Ok(None));
// The result should look like this:
// [6]
// / \
// [1, 2, 3, 4, 5] [7, 8, 9, 10, 11, 12]
for i in 1..=12 {
assert_eq!(btree.get(&vec![i]), Some(vec![]));
}
// Remove node 6. Triggers case 2.b
assert_eq!(btree.remove(&vec![6]), Some(vec![]));
// The result should look like this:
// [7]
// / \
// [1, 2, 3, 4, 5] [8, 9, 10, 11, 12]
let root = btree.load_node(btree.root_addr);
assert_eq!(root.node_type, NodeType::Internal);
assert_eq!(root.entries, vec![(vec![7], vec![])]);
assert_eq!(root.children.len(), 2);
let child_0 = btree.load_node(root.children[0]);
assert_eq!(child_0.node_type, NodeType::Leaf);
assert_eq!(child_0.entries, vec![e(1), e(2), e(3), e(4), e(5)]);
let child_1 = btree.load_node(root.children[1]);
assert_eq!(child_1.node_type, NodeType::Leaf);
assert_eq!(child_1.entries, vec![e(8), e(9), e(10), e(11), e(12)]);
// There are three allocated nodes.
assert_eq!(btree.allocator.num_allocated_chunks(), 3);
// Remove node 3. Triggers case 3.b
assert_eq!(btree.remove(&vec![3]), Some(vec![]));
// Reload the btree to verify that we saved it correctly.
let btree = BTreeMap::<_, Vec<u8>, Vec<u8>>::load(mem);
// The result should look like this:
//
// [1, 2, 4, 5, 7, 8, 9, 10, 11, 12]
let root = btree.load_node(btree.root_addr);
assert_eq!(root.node_type, NodeType::Leaf);
assert_eq!(
root.entries,
vec![
e(1),
e(2),
e(4),
e(5),
e(7),
e(8),
e(9),
e(10),
e(11),
e(12)
]
);
// There is only one allocated node remaining.
assert_eq!(btree.allocator.num_allocated_chunks(), 1);
}
#[test]
fn remove_case_3b_merge_into_left() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem.clone(), 5, 5);
for i in 1..=11 {
assert_eq!(btree.insert(vec![i], vec![]), Ok(None));
}
// Should now split a node.
assert_eq!(btree.insert(vec![12], vec![]), Ok(None));
// The result should look like this:
// [6]
// / \
// [1, 2, 3, 4, 5] [7, 8, 9, 10, 11, 12]
for i in 1..=12 {
assert_eq!(btree.get(&vec![i]), Some(vec![]));
}
// Remove node 6. Triggers case 2.b
assert_eq!(btree.remove(&vec![6]), Some(vec![]));
// The result should look like this:
// [7]
// / \
// [1, 2, 3, 4, 5] [8, 9, 10, 11, 12]
let root = btree.load_node(btree.root_addr);
assert_eq!(root.node_type, NodeType::Internal);
assert_eq!(root.entries, vec![(vec![7], vec![])]);
assert_eq!(root.children.len(), 2);
let child_0 = btree.load_node(root.children[0]);
assert_eq!(child_0.node_type, NodeType::Leaf);
assert_eq!(child_0.entries, vec![e(1), e(2), e(3), e(4), e(5)]);
let child_1 = btree.load_node(root.children[1]);
assert_eq!(child_1.node_type, NodeType::Leaf);
assert_eq!(child_1.entries, vec![e(8), e(9), e(10), e(11), e(12)]);
// There are three allocated nodes.
assert_eq!(btree.allocator.num_allocated_chunks(), 3);
// Remove node 10. Triggers case 3.b where we merge the right into the left.
assert_eq!(btree.remove(&vec![10]), Some(vec![]));
// Reload the btree to verify that we saved it correctly.
let btree = BTreeMap::<_, Vec<u8>, Vec<u8>>::load(mem);
// The result should look like this:
//
// [1, 2, 3, 4, 5, 7, 8, 9, 11, 12]
let root = btree.load_node(btree.root_addr);
assert_eq!(root.node_type, NodeType::Leaf);
assert_eq!(
root.entries,
vec![e(1), e(2), e(3), e(4), e(5), e(7), e(8), e(9), e(11), e(12)]
);
// There is only one allocated node remaining.
assert_eq!(btree.allocator.num_allocated_chunks(), 1);
}
#[test]
fn many_insertions() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem.clone(), 5, 5);
for j in 0..=10 {
for i in 0..=255 {
assert_eq!(btree.insert(vec![i, j], vec![i, j]), Ok(None));
}
}
for j in 0..=10 {
for i in 0..=255 {
assert_eq!(btree.get(&vec![i, j]), Some(vec![i, j]));
}
}
let mut btree = BTreeMap::load(mem);
for j in 0..=10 {
for i in 0..=255 {
assert_eq!(btree.remove(&vec![i, j]), Some(vec![i, j]));
}
}
for j in 0..=10 {
for i in 0..=255 {
assert_eq!(btree.get(&vec![i, j]), None);
}
}
// We've deallocated everything.
assert_eq!(btree.allocator.num_allocated_chunks(), 0);
}
#[test]
fn many_insertions_2() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem.clone(), 5, 5);
for j in (0..=10).rev() {
for i in (0..=255).rev() {
assert_eq!(btree.insert(vec![i, j], vec![i, j]), Ok(None));
}
}
for j in 0..=10 {
for i in 0..=255 {
assert_eq!(btree.get(&vec![i, j]), Some(vec![i, j]));
}
}
let mut btree = BTreeMap::load(mem);
for j in (0..=10).rev() {
for i in (0..=255).rev() {
assert_eq!(btree.remove(&vec![i, j]), Some(vec![i, j]));
}
}
for j in 0..=10 {
for i in 0..=255 {
assert_eq!(btree.get(&vec![i, j]), None);
}
}
// We've deallocated everything.
assert_eq!(btree.allocator.num_allocated_chunks(), 0);
}
#[test]
fn reloading() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem.clone(), 5, 5);
// The btree is initially empty.
assert_eq!(btree.len(), 0);
assert!(btree.is_empty());
// Add an entry into the btree.
assert_eq!(btree.insert(vec![1, 2, 3], vec![4, 5, 6]), Ok(None));
assert_eq!(btree.len(), 1);
assert!(!btree.is_empty());
// Reload the btree. The element should still be there, and `len()`
// should still be `1`.
let btree = BTreeMap::<_, Vec<u8>, Vec<u8>>::load(mem.clone());
assert_eq!(btree.get(&vec![1, 2, 3]), Some(vec![4, 5, 6]));
assert_eq!(btree.len(), 1);
assert!(!btree.is_empty());
// Remove an element. Length should be zero.
let mut btree = BTreeMap::load(mem.clone());
assert_eq!(btree.remove(&vec![1, 2, 3]), Some(vec![4, 5, 6]));
assert_eq!(btree.len(), 0);
assert!(btree.is_empty());
// Reload. Btree should still be empty.
let btree = BTreeMap::<_, Vec<u8>, Vec<u8>>::load(mem);
assert_eq!(btree.get(&vec![1, 2, 3]), None);
assert_eq!(btree.len(), 0);
assert!(btree.is_empty());
}
#[test]
fn len() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
for i in 0..1000u32 {
assert_eq!(btree.insert(i.to_le_bytes().to_vec(), vec![]), Ok(None));
}
assert_eq!(btree.len(), 1000);
assert!(!btree.is_empty());
for i in 0..1000u32 {
assert_eq!(btree.remove(&i.to_le_bytes().to_vec()), Some(vec![]));
}
assert_eq!(btree.len(), 0);
assert!(btree.is_empty());
}
#[test]
fn contains_key() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
// Insert even numbers from 0 to 1000.
for i in (0..1000u32).step_by(2) {
assert_eq!(btree.insert(i.to_le_bytes().to_vec(), vec![]), Ok(None));
}
// Contains key should return true on all the even numbers and false on all the odd
// numbers.
for i in 0..1000u32 {
assert_eq!(btree.contains_key(&i.to_le_bytes().to_vec()), i % 2 == 0);
}
}
#[test]
fn range_empty() {
let mem = make_memory();
let btree = BTreeMap::<_, Vec<u8>, Vec<u8>>::new(mem, 5, 5);
// Test prefixes that don't exist in the map.
assert_eq!(btree.range(vec![0], None).collect::<Vec<_>>(), vec![]);
assert_eq!(
btree.range(vec![1, 2, 3, 4], None).collect::<Vec<_>>(),
vec![]
);
}
// Tests the case where the prefix is larger than all the entries in a leaf node.
#[test]
fn range_leaf_prefix_greater_than_all_entries() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
btree.insert(vec![0], vec![]).unwrap();
// Test a prefix that's larger than the value in the leaf node. Should be empty.
assert_eq!(btree.range(vec![1], None).collect::<Vec<_>>(), vec![]);
}
// Tests the case where the prefix is larger than all the entries in an internal node.
#[test]
fn range_internal_prefix_greater_than_all_entries() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
for i in 1..=12 {
assert_eq!(btree.insert(vec![i], vec![]), Ok(None));
}
// The result should look like this:
// [6]
// / \
// [1, 2, 3, 4, 5] [7, 8, 9, 10, 11, 12]
// Test a prefix that's larger than the value in the internal node.
assert_eq!(
btree.range(vec![7], None).collect::<Vec<_>>(),
vec![(vec![7], vec![])]
);
}
#[test]
fn range_various_prefixes() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
btree.insert(vec![0, 1], vec![]).unwrap();
btree.insert(vec![0, 2], vec![]).unwrap();
btree.insert(vec![0, 3], vec![]).unwrap();
btree.insert(vec![0, 4], vec![]).unwrap();
btree.insert(vec![1, 1], vec![]).unwrap();
btree.insert(vec![1, 2], vec![]).unwrap();
btree.insert(vec![1, 3], vec![]).unwrap();
btree.insert(vec![1, 4], vec![]).unwrap();
btree.insert(vec![2, 1], vec![]).unwrap();
btree.insert(vec![2, 2], vec![]).unwrap();
btree.insert(vec![2, 3], vec![]).unwrap();
btree.insert(vec![2, 4], vec![]).unwrap();
// The result should look like this:
// [(1, 2)]
// / \
// [(0, 1), (0, 2), (0, 3), (0, 4), (1, 1)] [(1, 3), (1, 4), (2, 1), (2, 2), (2, 3), (2, 4)]
let root = btree.load_node(btree.root_addr);
assert_eq!(root.node_type, NodeType::Internal);
assert_eq!(root.entries, vec![(vec![1, 2], vec![])]);
assert_eq!(root.children.len(), 2);
// Tests a prefix that's smaller than the value in the internal node.
assert_eq!(
btree.range(vec![0], None).collect::<Vec<_>>(),
vec![
(vec![0, 1], vec![]),
(vec![0, 2], vec![]),
(vec![0, 3], vec![]),
(vec![0, 4], vec![]),
]
);
// Tests a prefix that crosses several nodes.
assert_eq!(
btree.range(vec![1], None).collect::<Vec<_>>(),
vec![
(vec![1, 1], vec![]),
(vec![1, 2], vec![]),
(vec![1, 3], vec![]),
(vec![1, 4], vec![]),
]
);
// Tests a prefix that's larger than the value in the internal node.
assert_eq!(
btree.range(vec![2], None).collect::<Vec<_>>(),
vec![
(vec![2, 1], vec![]),
(vec![2, 2], vec![]),
(vec![2, 3], vec![]),
(vec![2, 4], vec![]),
]
);
}
#[test]
fn range_various_prefixes_2() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
btree.insert(vec![0, 1], vec![]).unwrap();
btree.insert(vec![0, 2], vec![]).unwrap();
btree.insert(vec![0, 3], vec![]).unwrap();
btree.insert(vec![0, 4], vec![]).unwrap();
btree.insert(vec![1, 2], vec![]).unwrap();
btree.insert(vec![1, 4], vec![]).unwrap();
btree.insert(vec![1, 6], vec![]).unwrap();
btree.insert(vec![1, 8], vec![]).unwrap();
btree.insert(vec![1, 10], vec![]).unwrap();
btree.insert(vec![2, 1], vec![]).unwrap();
btree.insert(vec![2, 2], vec![]).unwrap();
btree.insert(vec![2, 3], vec![]).unwrap();
btree.insert(vec![2, 4], vec![]).unwrap();
btree.insert(vec![2, 5], vec![]).unwrap();
btree.insert(vec![2, 6], vec![]).unwrap();
btree.insert(vec![2, 7], vec![]).unwrap();
btree.insert(vec![2, 8], vec![]).unwrap();
btree.insert(vec![2, 9], vec![]).unwrap();
// The result should look like this:
// [(1, 4), (2, 3)]
// / | \
// [(0, 1), (0, 2), (0, 3), (0, 4), (1, 2)] | [(2, 4), (2, 5), (2, 6), (2, 7), (2, 8), (2, 9)]
// |
// [(1, 6), (1, 8), (1, 10), (2, 1), (2, 2)]
let root = btree.load_node(btree.root_addr);
assert_eq!(root.node_type, NodeType::Internal);
assert_eq!(
root.entries,
vec![(vec![1, 4], vec![]), (vec![2, 3], vec![])]
);
assert_eq!(root.children.len(), 3);
let child_0 = btree.load_node(root.children[0]);
assert_eq!(child_0.node_type, NodeType::Leaf);
assert_eq!(
child_0.entries,
vec![
(vec![0, 1], vec![]),
(vec![0, 2], vec![]),
(vec![0, 3], vec![]),
(vec![0, 4], vec![]),
(vec![1, 2], vec![]),
]
);
let child_1 = btree.load_node(root.children[1]);
assert_eq!(child_1.node_type, NodeType::Leaf);
assert_eq!(
child_1.entries,
vec![
(vec![1, 6], vec![]),
(vec![1, 8], vec![]),
(vec![1, 10], vec![]),
(vec![2, 1], vec![]),
(vec![2, 2], vec![]),
]
);
let child_2 = btree.load_node(root.children[2]);
assert_eq!(
child_2.entries,
vec![
(vec![2, 4], vec![]),
(vec![2, 5], vec![]),
(vec![2, 6], vec![]),
(vec![2, 7], vec![]),
(vec![2, 8], vec![]),
(vec![2, 9], vec![]),
]
);
// Tests a prefix that doesn't exist, but is in the middle of the root node.
assert_eq!(btree.range(vec![1, 5], None).collect::<Vec<_>>(), vec![]);
// Tests a prefix that crosses several nodes.
assert_eq!(
btree.range(vec![1], None).collect::<Vec<_>>(),
vec![
(vec![1, 2], vec![]),
(vec![1, 4], vec![]),
(vec![1, 6], vec![]),
(vec![1, 8], vec![]),
(vec![1, 10], vec![]),
]
);
// Tests a prefix that starts from a leaf node, then iterates through the root and right
// sibling.
assert_eq!(
btree.range(vec![2], None).collect::<Vec<_>>(),
vec![
(vec![2, 1], vec![]),
(vec![2, 2], vec![]),
(vec![2, 3], vec![]),
(vec![2, 4], vec![]),
(vec![2, 5], vec![]),
(vec![2, 6], vec![]),
(vec![2, 7], vec![]),
(vec![2, 8], vec![]),
(vec![2, 9], vec![]),
]
);
}
#[test]
fn range_large() {
let mem = make_memory();
let mut btree = BTreeMap::<_, Vec<u8>, Vec<u8>>::new(mem, 5, 5);
// Insert 1000 elements with prefix 0 and another 1000 elements with prefix 1.
for prefix in 0..=1 {
for i in 0..1000u32 {
assert_eq!(
btree.insert(
// The key is the prefix followed by the integer's encoding.
// The encoding is big-endian so that the byte representation of the
// integers are sorted.
vec![vec![prefix], i.to_be_bytes().to_vec()]
.into_iter()
.flatten()
.collect(),
vec![]
),
Ok(None)
);
}
}
// Getting the range with a prefix should return all 1000 elements with that prefix.
for prefix in 0..=1 {
let mut i: u32 = 0;
for (key, _) in btree.range(vec![prefix], None) {
assert_eq!(
key,
vec![vec![prefix], i.to_be_bytes().to_vec()]
.into_iter()
.flatten()
.collect::<Vec<_>>()
);
i += 1;
}
assert_eq!(i, 1000);
}
}
#[test]
fn range_various_prefixes_with_offset() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
btree.insert(vec![0, 1], vec![]).unwrap();
btree.insert(vec![0, 2], vec![]).unwrap();
btree.insert(vec![0, 3], vec![]).unwrap();
btree.insert(vec![0, 4], vec![]).unwrap();
btree.insert(vec![1, 1], vec![]).unwrap();
btree.insert(vec![1, 2], vec![]).unwrap();
btree.insert(vec![1, 3], vec![]).unwrap();
btree.insert(vec![1, 4], vec![]).unwrap();
btree.insert(vec![2, 1], vec![]).unwrap();
btree.insert(vec![2, 2], vec![]).unwrap();
btree.insert(vec![2, 3], vec![]).unwrap();
btree.insert(vec![2, 4], vec![]).unwrap();
// The result should look like this:
// [(1, 2)]
// / \
// [(0, 1), (0, 2), (0, 3), (0, 4), (1, 1)] [(1, 3), (1, 4), (2, 1), (2, 2), (2, 3), (2, 4)]
let root = btree.load_node(btree.root_addr);
assert_eq!(root.node_type, NodeType::Internal);
assert_eq!(root.entries, vec![(vec![1, 2], vec![])]);
assert_eq!(root.children.len(), 2);
// Tests a offset that's smaller than the value in the internal node.
assert_eq!(
btree.range(vec![0], Some(vec![0])).collect::<Vec<_>>(),
vec![
(vec![0, 1], vec![]),
(vec![0, 2], vec![]),
(vec![0, 3], vec![]),
(vec![0, 4], vec![]),
]
);
// Tests a offset that has a value somewhere in the range of values of an internal node.
assert_eq!(
btree.range(vec![1], Some(vec![3])).collect::<Vec<_>>(),
vec![(vec![1, 3], vec![]), (vec![1, 4], vec![]),]
);
// Tests a offset that's larger than the value in the internal node.
assert_eq!(
btree.range(vec![2], Some(vec![5])).collect::<Vec<_>>(),
vec![],
);
}
#[test]
fn range_various_prefixes_with_offset_2() {
let mem = make_memory();
let mut btree = BTreeMap::new(mem, 5, 5);
btree.insert(vec![0, 1], vec![]).unwrap();
btree.insert(vec![0, 2], vec![]).unwrap();
btree.insert(vec![0, 3], vec![]).unwrap();
btree.insert(vec![0, 4], vec![]).unwrap();
btree.insert(vec![1, 2], vec![]).unwrap();
btree.insert(vec![1, 4], vec![]).unwrap();
btree.insert(vec![1, 6], vec![]).unwrap();
btree.insert(vec![1, 8], vec![]).unwrap();
btree.insert(vec![1, 10], vec![]).unwrap();
btree.insert(vec![2, 1], vec![]).unwrap();
btree.insert(vec![2, 2], vec![]).unwrap();
btree.insert(vec![2, 3], vec![]).unwrap();
btree.insert(vec![2, 4], vec![]).unwrap();
btree.insert(vec![2, 5], vec![]).unwrap();
btree.insert(vec![2, 6], vec![]).unwrap();
btree.insert(vec![2, 7], vec![]).unwrap();
btree.insert(vec![2, 8], vec![]).unwrap();
btree.insert(vec![2, 9], vec![]).unwrap();
// The result should look like this:
// [(1, 4), (2, 3)]
// / | \
// [(0, 1), (0, 2), (0, 3), (0, 4), (1, 2)] | [(2, 4), (2, 5), (2, 6), (2, 7), (2, 8), (2, 9)]
// |
// [(1, 6), (1, 8), (1, 10), (2, 1), (2, 2)]
let root = btree.load_node(btree.root_addr);
assert_eq!(root.node_type, NodeType::Internal);
assert_eq!(
root.entries,
vec![(vec![1, 4], vec![]), (vec![2, 3], vec![])]
);
assert_eq!(root.children.len(), 3);
let child_0 = btree.load_node(root.children[0]);
assert_eq!(child_0.node_type, NodeType::Leaf);
assert_eq!(
child_0.entries,
vec![
(vec![0, 1], vec![]),
(vec![0, 2], vec![]),
(vec![0, 3], vec![]),
(vec![0, 4], vec![]),
(vec![1, 2], vec![]),
]
);
let child_1 = btree.load_node(root.children[1]);
assert_eq!(child_1.node_type, NodeType::Leaf);
assert_eq!(
child_1.entries,
vec![
(vec![1, 6], vec![]),
(vec![1, 8], vec![]),
(vec![1, 10], vec![]),
(vec![2, 1], vec![]),
(vec![2, 2], vec![]),
]
);
let child_2 = btree.load_node(root.children[2]);
assert_eq!(
child_2.entries,
vec![
(vec![2, 4], vec![]),
(vec![2, 5], vec![]),
(vec![2, 6], vec![]),
(vec![2, 7], vec![]),
(vec![2, 8], vec![]),
(vec![2, 9], vec![]),
]
);
// Tests a offset that crosses several nodes.
assert_eq!(
btree.range(vec![1], Some(vec![4])).collect::<Vec<_>>(),
vec![
(vec![1, 4], vec![]),
(vec![1, 6], vec![]),
(vec![1, 8], vec![]),
(vec![1, 10], vec![]),
]
);
// Tests a offset that starts from a leaf node, then iterates through the root and right
// sibling.
assert_eq!(
btree.range(vec![2], Some(vec![2])).collect::<Vec<_>>(),
vec![
(vec![2, 2], vec![]),
(vec![2, 3], vec![]),
(vec![2, 4], vec![]),
(vec![2, 5], vec![]),
(vec![2, 6], vec![]),
(vec![2, 7], vec![]),
(vec![2, 8], vec![]),
(vec![2, 9], vec![]),
]
);
}
}