Skip to content

refactor: improve naming of BTreeMap inner key/value storage #334

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jun 5, 2025
Merged

refactor: improve naming of BTreeMap inner key/value storage #334

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 22 additions & 32 deletions src/btreemap/node.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ pub struct Node<K: Storable + Ord + Clone> {
address: Address,
// List of tuples consisting of a key and the encoded value.
// INVARIANT: the list is sorted by key.
keys_and_encoded_values: Vec<LazyEntry<K>>,
entries: Vec<LazyEntry<K>>,
// For the key at position I, children[I] points to the left
// child of this key and children[I + 1] points to the right child.
children: Vec<Address>,
Expand Down Expand Up @@ -108,10 +108,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
pub fn get_max<M: Memory>(&self, memory: &M) -> Entry<K> {
match self.node_type {
NodeType::Leaf => {
let last_entry = self
.keys_and_encoded_values
.last()
.expect("A node can never be empty");
let last_entry = self.entries.last().expect("A node can never be empty");
(
self.get_key(last_entry, memory).clone(),
self.get_value(last_entry, memory).to_vec(),
Expand Down Expand Up @@ -153,15 +150,12 @@ impl<K: Storable + Ord + Clone> Node<K> {

/// Returns true if the node cannot store anymore entries, false otherwise.
pub fn is_full(&self) -> bool {
self.keys_and_encoded_values.len() >= CAPACITY
self.entries.len() >= CAPACITY
}

/// Replaces the value at `idx` and returns the old one.
pub fn swap_value<M: Memory>(&mut self, idx: usize, new: Vec<u8>, memory: &M) -> Vec<u8> {
let old = core::mem::replace(
&mut self.keys_and_encoded_values[idx].1,
LazyValue::by_value(new),
);
let old = core::mem::replace(&mut self.entries[idx].1, LazyValue::by_value(new));
self.extract_value(old, memory)
}

Expand All @@ -173,7 +167,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
memory: &M,
) -> Entry<K> {
let (old_key, old_value) = core::mem::replace(
&mut self.keys_and_encoded_values[idx],
&mut self.entries[idx],
(LazyKey::by_value(key), LazyValue::by_value(value)),
);
(
Expand Down Expand Up @@ -202,13 +196,13 @@ impl<K: Storable + Ord + Clone> Node<K> {
/// Returns a reference to the key at the specified index.
#[inline(always)]
pub fn key<M: Memory>(&self, idx: usize, memory: &M) -> &K {
self.get_key(&self.keys_and_encoded_values[idx], memory)
self.get_key(&self.entries[idx], memory)
}

/// Returns a reference to the encoded value at the specified index.
#[inline(always)]
pub fn value<M: Memory>(&self, idx: usize, memory: &M) -> &[u8] {
self.get_value(&self.keys_and_encoded_values[idx], memory)
self.get_value(&self.entries[idx], memory)
}

/// Extracts the contents of key (by loading it first if it's not loaded yet).
Expand Down Expand Up @@ -312,14 +306,14 @@ impl<K: Storable + Ord + Clone> Node<K> {

/// Inserts a new entry at the specified index.
pub fn insert_entry(&mut self, idx: usize, (key, value): Entry<K>) {
self.keys_and_encoded_values
self.entries
.insert(idx, (LazyKey::by_value(key), LazyValue::by_value(value)));
}

/// Returns the entry at the specified index while consuming this node.
pub fn into_entry<M: Memory>(mut self, idx: usize, memory: &M) -> Entry<K> {
let keys_and_encoded_values = core::mem::take(&mut self.keys_and_encoded_values);
let (key, value) = keys_and_encoded_values.into_iter().nth(idx).unwrap();
let entries = core::mem::take(&mut self.entries);
let (key, value) = entries.into_iter().nth(idx).unwrap();
(
self.extract_key(key, memory),
self.extract_value(value, memory),
Expand All @@ -328,7 +322,7 @@ impl<K: Storable + Ord + Clone> Node<K> {

/// Removes the entry at the specified index.
pub fn remove_entry<M: Memory>(&mut self, idx: usize, memory: &M) -> Entry<K> {
let (key, value) = self.keys_and_encoded_values.remove(idx);
let (key, value) = self.entries.remove(idx);
(
self.extract_key(key, memory),
self.extract_value(value, memory),
Expand All @@ -337,7 +331,7 @@ impl<K: Storable + Ord + Clone> Node<K> {

/// Adds a new entry at the back of the node.
pub fn push_entry(&mut self, (key, value): Entry<K>) {
self.keys_and_encoded_values
self.entries
.push((LazyKey::by_value(key), LazyValue::by_value(value)));
}

Expand All @@ -348,10 +342,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
return None;
}

let (key, value) = self
.keys_and_encoded_values
.pop()
.expect("node must not be empty");
let (key, value) = self.entries.pop().expect("node must not be empty");

Some((
self.extract_key(key, memory),
Expand Down Expand Up @@ -391,7 +382,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
Self::append(&mut source, self, median, allocator.memory());

// Move the entries and children into self.
self.keys_and_encoded_values = core::mem::take(&mut source.keys_and_encoded_values);
self.entries = core::mem::take(&mut source.entries);
self.children = core::mem::take(&mut source.children);
}

Expand Down Expand Up @@ -420,27 +411,26 @@ impl<K: Storable + Ord + Clone> Node<K> {

a.push_entry(median);

a.keys_and_encoded_values
.append(&mut b.keys_and_encoded_values);
a.entries.append(&mut b.entries);

// Move the children (if any exist).
a.children.append(&mut b.children);

// Assert postconditions.
assert_eq!(b.keys_and_encoded_values.len(), 0);
assert_eq!(b.entries.len(), 0);
assert_eq!(b.children.len(), 0);
}

#[cfg(test)]
pub fn entries<M: Memory>(&self, memory: &M) -> Vec<Entry<K>> {
(0..self.keys_and_encoded_values.len())
(0..self.entries.len())
.map(|i| (self.key(i, memory).clone(), self.value(i, memory).to_vec()))
.collect()
}

#[cfg(test)]
pub fn keys<M: Memory>(&self, memory: &M) -> Vec<&K> {
(0..self.keys_and_encoded_values.len())
(0..self.entries.len())
.map(|i| self.key(i, memory))
.collect()
}
Expand All @@ -452,7 +442,7 @@ impl<K: Storable + Ord + Clone> Node<K> {

/// Returns the number of entries in the node.
pub fn entries_len(&self) -> usize {
self.keys_and_encoded_values.len()
self.entries.len()
}

/// Searches for the key in the node's entries.
Expand All @@ -462,7 +452,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
/// returned, containing the index where a matching key could be inserted
/// while maintaining sorted order.
pub fn search<M: Memory>(&self, key: &K, memory: &M) -> Result<usize, usize> {
self.keys_and_encoded_values
self.entries
.binary_search_by_key(&key, |entry| self.get_key(entry, memory))
}

Expand All @@ -475,7 +465,7 @@ impl<K: Storable + Ord + Clone> Node<K> {

/// Returns true if the node is at the minimum required size, false otherwise.
pub fn at_minimum(&self) -> bool {
self.keys_and_encoded_values.len() < B
self.entries.len() < B
}

/// Returns true if an entry can be removed without having to merge it into another node
Expand All @@ -494,7 +484,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
}

// Move the entries and children above the median into the new sibling.
sibling.keys_and_encoded_values = self.keys_and_encoded_values.split_off(B);
sibling.entries = self.entries.split_off(B);
if self.node_type == NodeType::Internal {
sibling.children = self.children.split_off(B);
}
Expand Down
22 changes: 11 additions & 11 deletions src/btreemap/node/v1.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
Node {
address,
node_type,
keys_and_encoded_values: vec![],
entries: vec![],
children: vec![],
version: Version::V1(page_size),
overflows: Vec::with_capacity(0),
Expand All @@ -67,7 +67,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
let _p = canbench_rs::bench_scope("node_load_v1"); // May add significant overhead.

// Load the entries.
let mut keys_encoded_values = Vec::with_capacity(header.num_entries as usize);
let mut entries = Vec::with_capacity(header.num_entries as usize);
let mut offset = NodeHeader::size();
for _ in 0..header.num_entries {
let key_offset = offset;
Expand All @@ -80,7 +80,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
let value = LazyValue::by_ref(value_offset);
offset += Bytes::from(max_value_size);

keys_encoded_values.push((key, value));
entries.push((key, value));
}

// Load children if this is an internal node.
Expand All @@ -94,12 +94,12 @@ impl<K: Storable + Ord + Clone> Node<K> {
children.push(child);
}

assert_eq!(children.len(), keys_encoded_values.len() + 1);
assert_eq!(children.len(), entries.len() + 1);
}

Self {
address,
keys_and_encoded_values: keys_encoded_values,
entries,
children,
node_type: match header.node_type {
LEAF_NODE_TYPE => NodeType::Leaf,
Expand All @@ -123,16 +123,16 @@ impl<K: Storable + Ord + Clone> Node<K> {
assert!(self.children.is_empty());
}
NodeType::Internal => {
assert_eq!(self.children.len(), self.keys_and_encoded_values.len() + 1);
assert_eq!(self.children.len(), self.entries.len() + 1);
}
};

// We should never be saving an empty node.
assert!(!self.keys_and_encoded_values.is_empty() || !self.children.is_empty());
assert!(!self.entries.is_empty() || !self.children.is_empty());

// Assert entries are sorted in strictly increasing order.
assert!(self
.keys_and_encoded_values
.entries
.windows(2)
.all(|arr| self.get_key(&arr[0], memory) < self.get_key(&arr[1], memory)));

Expand All @@ -151,7 +151,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
NodeType::Leaf => LEAF_NODE_TYPE,
NodeType::Internal => INTERNAL_NODE_TYPE,
},
num_entries: self.keys_and_encoded_values.len() as u16,
num_entries: self.entries.len() as u16,
};

write_struct(&header, self.address, memory);
Expand All @@ -160,12 +160,12 @@ impl<K: Storable + Ord + Clone> Node<K> {

// Load all the entries. This is necessary so that we don't overwrite referenced
// entries when writing the entries to the node.
for i in 0..self.keys_and_encoded_values.len() {
for i in 0..self.entries.len() {
self.entry(i, memory);
}

// Write the entries.
for i in 0..self.keys_and_encoded_values.len() {
for i in 0..self.entries.len() {
// Write the size of the key.
let key_bytes = self.key(i, memory).to_bytes_checked();
write_u32(memory, self.address + offset, key_bytes.len() as u32);
Expand Down
18 changes: 9 additions & 9 deletions src/btreemap/node/v2.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,7 +98,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
address,
node_type,
version: Version::V2(page_size),
keys_and_encoded_values: vec![],
entries: vec![],
children: vec![],
overflows: Vec::with_capacity(0),
}
Expand Down Expand Up @@ -151,7 +151,7 @@ impl<K: Storable + Ord + Clone> Node<K> {

// Load the keys (eagerly if small).
const EAGER_LOAD_KEY_SIZE_THRESHOLD: u32 = 16;
let mut keys_encoded_values = Vec::with_capacity(num_entries);
let mut entries = Vec::with_capacity(num_entries);
let mut buf = vec![];

for _ in 0..num_entries {
Expand Down Expand Up @@ -180,11 +180,11 @@ impl<K: Storable + Ord + Clone> Node<K> {
};

offset += Bytes::from(key_size);
keys_encoded_values.push((key, LazyValue::by_ref(Bytes::from(0_u64))));
entries.push((key, LazyValue::by_ref(Bytes::from(0_u64))));
}

// Load the values
for (_key, value) in keys_encoded_values.iter_mut() {
for (_key, value) in entries.iter_mut() {
// Load the values lazily.
*value = LazyValue::by_ref(Bytes::from(offset.get()));
let value_size = read_u32(&reader, offset) as usize;
Expand All @@ -193,7 +193,7 @@ impl<K: Storable + Ord + Clone> Node<K> {

Self {
address,
keys_and_encoded_values: keys_encoded_values,
entries,
children,
node_type,
version: Version::V2(page_size),
Expand All @@ -211,7 +211,7 @@ impl<K: Storable + Ord + Clone> Node<K> {

// Load all the entries. One pass is required to load all entries;
// results are not stored to avoid unnecessary allocations.
for i in 0..self.keys_and_encoded_values.len() {
for i in 0..self.entries.len() {
self.entry(i, allocator.memory());
}

Expand All @@ -233,7 +233,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
NodeType::Leaf => LEAF_NODE_TYPE,
NodeType::Internal => INTERNAL_NODE_TYPE,
},
num_entries: self.keys_and_encoded_values.len() as u16,
num_entries: self.entries.len() as u16,
};

writer.write_struct(&header, offset);
Expand All @@ -251,7 +251,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
}

// Write the keys.
for i in 0..self.keys_and_encoded_values.len() {
for i in 0..self.entries.len() {
let key = self.key(i, writer.memory());
let key_bytes = key.to_bytes_checked();

Expand All @@ -267,7 +267,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
}

// Write the values.
for i in 0..self.keys_and_encoded_values.len() {
for i in 0..self.entries.len() {
// Write the size of the value.
let value = self.value(i, writer.memory());
writer.write_u32(offset, value.len() as u32);
Expand Down