Skip to content
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

Refactor iterators #78

Merged
merged 7 commits into from
Jan 7, 2025
Merged

Refactor iterators #78

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
5996d0b
convert vec to box
arriqaaq Jan 6, 2025
e7890e0
feat: add query iterator
arriqaaq Jan 6, 2025
c0a2154
chore: dont use box for keys
arriqaaq Jan 6, 2025
7e6b665
fix: clippy
arriqaaq Jan 6, 2025
5b7347e
fix: iter
arriqaaq Jan 6, 2025
c17c692
fix: iter
arriqaaq Jan 6, 2025
88d291e
refactor: unify return type for all iters
arriqaaq Jan 6, 2025
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
84 changes: 41 additions & 43 deletions src/art.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ use std::cmp::Ordering;
use std::ops::RangeBounds;
use std::sync::Arc;

use crate::iter::IterItem;
use crate::iter::{query_keys_at_node, scan_node, Iter, Range};
use crate::node::{FlatNode, LeafValue, Node256, Node48, NodeTrait, TwigNode};
use crate::{KeyTrait, TrieError};
Expand Down Expand Up @@ -690,7 +691,7 @@ impl<P: KeyTrait, V: Clone> Node<P, V> {
}

#[inline]
pub(crate) fn get_leaf_by_query(&self, query_type: QueryType) -> Option<Arc<LeafValue<V>>> {
pub(crate) fn get_leaf_by_query(&self, query_type: QueryType) -> Option<&Arc<LeafValue<V>>> {
let twig = if let NodeType::Twig(twig) = &self.node_type {
// For a Twig node simply use its inner value.
twig
Expand Down Expand Up @@ -1589,10 +1590,7 @@ impl<P: KeyTrait, V: Clone> Tree<P, V> {
/// Returns a `Range` iterator instance that iterates over the key-value pairs within the given range.
/// If the Trie is empty, an empty `Range` iterator is returned.
///
pub fn range<'a, R>(
&'a self,
range: R,
) -> impl Iterator<Item = (Vec<u8>, &'a V, &'a u64, &'a u64)>
pub fn range<'a, R>(&'a self, range: R) -> impl Iterator<Item = IterItem<'a, V>>
where
R: RangeBounds<P> + 'a,
{
Expand All @@ -1618,10 +1616,7 @@ impl<P: KeyTrait, V: Clone> Tree<P, V> {
/// # Returns
///
/// Returns an iterator over the key-value pairs, versions, and timestamps within the specified range.
pub fn range_with_versions<'a, R>(
&'a self,
range: R,
) -> impl Iterator<Item = (Vec<u8>, &'a V, &'a u64, &'a u64)>
pub fn range_with_versions<'a, R>(&'a self, range: R) -> impl Iterator<Item = IterItem<'a, V>>
where
R: RangeBounds<P> + 'a,
{
Expand Down Expand Up @@ -1713,16 +1708,20 @@ impl<P: KeyTrait, V: Clone> Tree<P, V> {
self.size == 0
}

pub fn scan_at_ts<R>(&self, range: R, ts: u64) -> Vec<(Vec<u8>, V)>
pub fn scan_at_ts<'a, R>(
&'a self,
range: R,
ts: u64,
) -> impl Iterator<Item = IterItem<'a, V>> + 'a
where
R: RangeBounds<P>,
R: RangeBounds<P> + 'a,
{
scan_node(self.root.as_ref(), range, QueryType::LatestByTs(ts))
}

pub fn keys_at_ts<R>(&self, range: R, ts: u64) -> Vec<Vec<u8>>
pub fn keys_at_ts<'a, R>(&'a self, range: R, ts: u64) -> impl Iterator<Item = &'a [u8]> + 'a
where
R: RangeBounds<P>,
R: RangeBounds<P> + 'a,
{
query_keys_at_node(self.root.as_ref(), range, QueryType::LatestByTs(ts))
}
Expand All @@ -1733,7 +1732,7 @@ impl<P: KeyTrait, V: Clone> Tree<P, V> {
/// all the nodes. This version could be lesser than the current incremental version.
pub fn max_version_in_trie(&self) -> u64 {
self.iter()
.map(|(_, _, version, _)| *version)
.map(|(_, _, version, _)| version)
.max()
.unwrap_or(0)
}
Expand Down Expand Up @@ -2419,10 +2418,10 @@ mod tests {

let tree_iter = tree.iter();
for tree_entry in tree_iter {
let k = from_be_bytes_key(&tree_entry.0);
let k = from_be_bytes_key(tree_entry.0);
assert_eq!(expected as u64, k);
let ts = tree_entry.3;
assert_eq!(expected as u64, *ts);
assert_eq!(expected as u64, ts);
expected = expected.wrapping_add(1);
len += 1;
}
Expand All @@ -2447,7 +2446,7 @@ mod tests {

let tree_iter = tree.iter();
for tree_entry in tree_iter {
let k = from_be_bytes_key(&tree_entry.0);
let k = from_be_bytes_key(tree_entry.0);
assert_eq!(expected as u64, k);
expected = expected.wrapping_add(1);
len += 1;
Expand Down Expand Up @@ -2869,7 +2868,7 @@ mod tests {
let mut iter = tree.iter();
let (_, value, version, _) = iter.next().unwrap();
assert_eq!(value, &2);
assert_eq!(version, &2);
assert_eq!(version, 2);

// Verify get at version 0 gives the latest version
let (value, version, _) = tree.get(&key, 0).unwrap();
Expand Down Expand Up @@ -2933,7 +2932,7 @@ mod tests {

// Assert that results are in expected order
for (i, result) in results.iter().enumerate() {
let result_str = std::str::from_utf8(result.0.as_slice()).expect("Invalid UTF-8");
let result_str = std::str::from_utf8(result.0).expect("Invalid UTF-8");
assert_eq!(result_str, expected_order[i]);
}
}
Expand Down Expand Up @@ -3019,7 +3018,7 @@ mod tests {
let mut count = 0;
let tree_iter = tree.iter();
for (key, value, version, ts) in tree_iter {
assert!(inserted_data.contains(&(key, *value, *version, *ts)));
assert!(inserted_data.contains(&(key.to_vec(), *value, version, ts)));
count += 1;
}

Expand Down Expand Up @@ -3333,7 +3332,7 @@ mod tests {
#[test]
fn scan_empty_range() {
let tree: Tree<VariableSizeKey, i32> = Tree::new();
let result = tree.scan_at_ts(RangeFull {}, 0);
let result: Vec<_> = tree.scan_at_ts(RangeFull {}, 0).collect();
assert!(result.is_empty());
}

Expand All @@ -3347,7 +3346,7 @@ mod tests {
}
let range = VariableSizeKey::from_slice("key_1".as_bytes())
..=VariableSizeKey::from_slice("key_2".as_bytes());
let values = tree.scan_at_ts(range, 0);
let values: Vec<_> = tree.scan_at_ts(range, 0).collect();
assert_eq!(values.len(), 2);
}

Expand All @@ -3359,7 +3358,7 @@ mod tests {
tree.insert(&VariableSizeKey::from_str(key).unwrap(), 1, 0, 0)
.unwrap();
}
let values = tree.scan_at_ts(RangeFull {}, 0);
let values: Vec<_> = tree.scan_at_ts(RangeFull {}, 0).collect();
assert_eq!(values.len(), keys.len());
}

Expand All @@ -3373,7 +3372,7 @@ mod tests {
}
let range = VariableSizeKey::from_slice("key_4".as_bytes())
..VariableSizeKey::from_slice("key_5".as_bytes());
let values = tree.scan_at_ts(range, 0);
let values: Vec<_> = tree.scan_at_ts(range, 0).collect();
assert!(values.is_empty());
}

Expand All @@ -3391,7 +3390,7 @@ mod tests {
.unwrap();
}
for (i, _) in keys.iter().enumerate() {
let values = tree.scan_at_ts(RangeFull {}, i as u64);
let values: Vec<_> = tree.scan_at_ts(RangeFull {}, i as u64).collect();
assert_eq!(values.len(), i + 1);
}
}
Expand All @@ -3407,7 +3406,7 @@ mod tests {
.unwrap();
}

let values = tree.scan_at_ts(RangeFull {}, num_keys);
let values: Vec<_> = tree.scan_at_ts(RangeFull {}, num_keys).collect();
assert_eq!(values.len(), num_keys as usize); // Expect all keys to be visible
}

Expand All @@ -3428,15 +3427,15 @@ mod tests {
}

// Scan at a timestamp before any insertions
let result_before = tree.scan_at_ts(RangeFull {}, 0);
let result_before: Vec<_> = tree.scan_at_ts(RangeFull {}, 0).collect();
assert!(result_before.is_empty());

// Scan between insertions
let result_mid = tree.scan_at_ts(RangeFull {}, 4);
let result_mid: Vec<_> = tree.scan_at_ts(RangeFull {}, 4).collect();
assert_eq!(result_mid.len(), 2); // Expect first two keys to be visible

// Scan after all insertions
let result_after = tree.scan_at_ts(RangeFull {}, 7);
let result_after: Vec<_> = tree.scan_at_ts(RangeFull {}, 7).collect();
assert_eq!(result_after.len(), keys.len()); // Expect all keys to be visible
}

Expand All @@ -3447,16 +3446,15 @@ mod tests {
tree.insert(&VariableSizeKey::from_str("key_1").unwrap(), 42, 0, 0)
.unwrap();

let values = tree.scan_at_ts(RangeFull {}, 0);

let values: Vec<_> = tree.scan_at_ts(RangeFull {}, 0).collect();
assert_eq!(values.len(), 1);
assert_eq!(values[0].1, 42);
assert_eq!(values[0].1, &42);
}

#[test]
fn keys_at_empty_range() {
let tree: Tree<VariableSizeKey, i32> = Tree::new();
let keys = tree.keys_at_ts(RangeFull {}, 0);
let keys: Vec<_> = tree.keys_at_ts(RangeFull {}, 0).collect();
assert!(keys.is_empty());
}

Expand All @@ -3470,7 +3468,7 @@ mod tests {
}
let range = VariableSizeKey::from_slice("key_1".as_bytes())
..=VariableSizeKey::from_slice("key_2".as_bytes());
let keys = tree.keys_at_ts(range, 0);
let keys: Vec<_> = tree.keys_at_ts(range, 0).collect();
assert_eq!(keys.len(), 2);
}

Expand All @@ -3482,7 +3480,7 @@ mod tests {
tree.insert(&VariableSizeKey::from_str(key).unwrap(), 1, 0, 0)
.unwrap();
}
let keys = tree.keys_at_ts(RangeFull {}, 0);
let keys: Vec<_> = tree.keys_at_ts(RangeFull {}, 0).collect();
assert_eq!(keys.len(), keys_to_insert.len());
}

Expand All @@ -3496,7 +3494,7 @@ mod tests {
}
let range = VariableSizeKey::from("key_4".as_bytes().to_vec())
..VariableSizeKey::from("key_5".as_bytes().to_vec());
let keys = tree.keys_at_ts(range, 0);
let keys: Vec<_> = tree.keys_at_ts(range, 0).collect();
assert!(keys.is_empty());
}

Expand All @@ -3514,7 +3512,7 @@ mod tests {
.unwrap();
}
for (i, _) in keys_to_insert.iter().enumerate() {
let keys = tree.keys_at_ts(RangeFull {}, i as u64);
let keys: Vec<_> = tree.keys_at_ts(RangeFull {}, i as u64).collect();
assert_eq!(keys.len(), i + 1);
}
}
Expand All @@ -3531,15 +3529,15 @@ mod tests {
.unwrap();
}

let keys = tree.keys_at_ts(RangeFull {}, num_keys);
let keys: Vec<_> = tree.keys_at_ts(RangeFull {}, num_keys).collect();
assert_eq!(keys.len(), num_keys as usize); // Expect all keys to be visible

// Sort the expected keys lexicographically
expected_keys.sort();

// Verify each key is proper
for (expected_key, key) in expected_keys.iter().zip(keys.iter()) {
assert_eq!(key, expected_key.to_slice());
assert_eq!(*key, expected_key.to_slice());
}
}

Expand All @@ -3558,13 +3556,13 @@ mod tests {
.unwrap();
}

let keys_before = tree.keys_at_ts(RangeFull {}, 0);
let keys_before: Vec<_> = tree.keys_at_ts(RangeFull {}, 0).collect();
assert!(keys_before.is_empty());

let keys_mid = tree.keys_at_ts(RangeFull {}, 4);
let keys_mid: Vec<_> = tree.keys_at_ts(RangeFull {}, 4).collect();
assert_eq!(keys_mid.len(), 2); // Expect first two keys to be visible

let keys_after = tree.keys_at_ts(RangeFull {}, 7);
let keys_after: Vec<_> = tree.keys_at_ts(RangeFull {}, 7).collect();
assert_eq!(keys_after.len(), keys_to_insert.len()); // Expect all keys to be visible
}

Expand All @@ -3574,7 +3572,7 @@ mod tests {
tree.insert(&VariableSizeKey::from_str("key_1").unwrap(), 42, 0, 0)
.unwrap();

let keys = tree.keys_at_ts(RangeFull {}, 0);
let keys: Vec<_> = tree.keys_at_ts(RangeFull {}, 0).collect();
assert_eq!(keys.len(), 1);
}

Expand Down
Loading
Loading