#6870

Committed 12 Feb 2026 01:11PM UTC coverage: 48.247% (-0.1%) from 48.357%

Build # #6870

Build Type

push

Committed by

Commit Message

Merge pull request #6460 from oasisprotocol/kostko/feature/rust-2026-02-11

Run Details

5 of 8 new or added lines in 3 files covered. (62.5%)

43 existing lines in 14 files now uncovered.

4651 of 9640 relevant lines covered (48.25%)

1.08 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

79.31

/runtime/src/storage/mkvs/tree/iterator.rs

//! Tree iterator.
use std::{collections::VecDeque, fmt};

use anyhow::{Error, Result};

use crate::storage::mkvs::{
    self,
    cache::{Cache, ReadSyncFetcher},
    sync::{IterateRequest, Proof, ReadSync, TreeID},
    tree::{Depth, Key, KeyTrait, NodeBox, NodeKind, NodePtrRef, Root, Tree},
};

pub(super) struct FetcherSyncIterate<'a> {
    key: &'a Key,
    prefetch: usize,
}

impl<'a> FetcherSyncIterate<'a> {
    pub(super) fn new(key: &'a Key, prefetch: usize) -> Self {
        Self { key, prefetch }
    }
}

impl ReadSyncFetcher for FetcherSyncIterate<'_> {
    fn fetch(&self, root: Root, ptr: NodePtrRef, rs: &mut Box<dyn ReadSync>) -> Result<Proof> {
        let rsp = rs.sync_iterate(IterateRequest {
            tree: TreeID {
                root,
                position: ptr.borrow().hash,
            },
            key: self.key.clone(),
            prefetch: self.prefetch as u16,
        })?;
        Ok(rsp.proof)
    }
}

/// Visit state of a node.
#[derive(Debug, PartialEq)]
enum VisitState {
    Before,
    At,
    AtLeft,
    After,
}

/// Atom in the current iterator path. Can be used to resume iteration
/// from a given position.
struct PathAtom {
    ptr: NodePtrRef,
    bit_depth: Depth,
    path: Key,
    state: VisitState,
}

impl fmt::Debug for PathAtom {
    fn fmt(&self, f: &mut fmt::Formatter) -> std::result::Result<(), fmt::Error> {
        f.debug_struct("PathAtom")
            .field("bit_depth", &self.bit_depth)
            .field("path", &self.path)
            .field("state", &self.state)
            .finish()
    }
}

/// Tree iterator.
pub struct TreeIterator<'tree> {
    tree: &'tree Tree,
    prefetch: usize,
    pos: VecDeque<PathAtom>,
    key: Option<Key>,
    value: Option<Vec<u8>>,
    error: Option<Error>,
}

impl<'tree> TreeIterator<'tree> {
    /// Create a new tree iterator.
    fn new(tree: &'tree Tree) -> Self {
        Self {
            tree,
            prefetch: 0,
            pos: VecDeque::new(),
            key: None,
            value: None,
            error: None,
        }
    }

    fn reset(&mut self) {
        self.pos.clear();
        self.key = None;
        self.value = None;
    }

    fn next(&mut self) {
        if self.error.is_some() {
            return;
        }

        while !self.pos.is_empty() {
            // Start where we left off.
            let atom = self.pos.pop_front().expect("not empty");
            let mut remainder = std::mem::take(&mut self.pos);

            // Remember where the path from root to target node ends (will end).
            let mut cache = self.tree.cache.borrow_mut();
            cache.mark_position();
            for atom in &remainder {
                cache.use_node(atom.ptr.clone());
            }
            drop(cache);

            // Try to proceed with the current node. If we don't succeed, proceed to the
            // next node.
            let key = self.key.take().expect("iterator is valid");
            self.reset();
            if let Err(error) =
                self._next(atom.ptr, atom.bit_depth, atom.path, key.clone(), atom.state)
            {
                self.error = Some(error);
                self.reset();
                return;
            }
            if self.key.is_some() {
                // Key has been found.
                self.pos.append(&mut remainder);
                return;
            }

            self.key = Some(key);
            self.pos = remainder;
        }

        // We have reached the end of the tree, make sure everything is reset.
        self.key = None;
        self.value = None;
    }

    fn _next(
        &mut self,
        ptr: NodePtrRef,
        bit_depth: Depth,
        path: Key,
        mut key: Key,
        mut state: VisitState,
    ) -> Result<()> {
        let node_ref = self.tree.cache.borrow_mut().deref_node_ptr(
            ptr.clone(),
            Some(FetcherSyncIterate::new(&key, self.prefetch)),
        )?;

        match classify_noderef!(?node_ref) {
            NodeKind::None => {
                // Reached a nil node, there is nothing here.
                Ok(())
            }
            NodeKind::Internal => {
                let node_ref = node_ref.unwrap();
                if let NodeBox::Internal(ref n) = *node_ref.borrow() {
                    // Internal node.
                    let bit_length = bit_depth + n.label_bit_length;
                    let new_path = path.merge(bit_depth, &n.label, n.label_bit_length);

                    // Check if the key is longer than the current path but lexicographically smaller. In this
                    // case everything in this subtree will be larger so we need to take the first value.
                    let take_first =
                        bit_length > 0 && key.bit_length() >= bit_length && key < new_path;

                    // Does lookup key end here? Look into LeafNode.
                    if (state == VisitState::Before
                        && (key.bit_length() <= bit_length || take_first))
                        || state == VisitState::At
                    {
                        if state == VisitState::Before {
                            self._next(
                                n.leaf_node.clone(),
                                bit_length,
                                path.clone(),
                                key.clone(),
                                VisitState::Before,
                            )?;
                            if self.key.is_some() {
                                // Key has been found.
                                self.pos.push_back(PathAtom {
                                    ptr,
                                    bit_depth,
                                    path,
                                    state: VisitState::At,
                                });
                                return Ok(());
                            }
                        }
                        // Key has not been found, continue with search for next key.
                        if key.bit_length() <= bit_length {
                            key = key.append_bit(bit_length, false);
                        }
                    }

                    if state == VisitState::Before {
                        state = VisitState::At;
                    }

                    // Continue recursively based on a bit value.
                    if (state == VisitState::At && (!key.get_bit(bit_length) || take_first))
                        || state == VisitState::AtLeft
                    {
                        if state == VisitState::At {
                            self._next(
                                n.left.clone(),
                                bit_length,
                                new_path.append_bit(bit_length, false),
                                key.clone(),
                                VisitState::Before,
                            )?;
                            if self.key.is_some() {
                                // Key has been found.
                                self.pos.push_back(PathAtom {
                                    ptr,
                                    bit_depth,
                                    path,
                                    state: VisitState::AtLeft,
                                });
                                return Ok(());
                            }
                        }
                        // Key has not been found, continue with search for next key.
                        key = key.split(bit_length, key.bit_length()).0;
                        key = key.append_bit(bit_length, true);
                    }

                    if state == VisitState::At || state == VisitState::AtLeft {
                        self._next(
                            n.right.clone(),
                            bit_length,
                            new_path.append_bit(bit_length, true),
                            key,
                            VisitState::Before,
                        )?;
                        if self.key.is_some() {
                            // Key has been found.
                            self.pos.push_back(PathAtom {
                                ptr,
                                bit_depth,
                                path,
                                state: VisitState::After,
                            });
                            return Ok(());
                        }
                    }

                    return Ok(());
                }

                unreachable!("node kind is internal node");
            }
            NodeKind::Leaf => {
                // Reached a leaf node.
                let node_ref = node_ref.unwrap();
                if let NodeBox::Leaf(ref n) = *node_ref.borrow() {
                    if n.key >= key {
                        self.key = Some(n.key.clone());
                        self.value = Some(n.value.clone());
                    }
                } else {
                    unreachable!("node kind is leaf node");
                }

                Ok(())
            }
        }
    }
}

impl Iterator for TreeIterator<'_> {
    type Item = (Vec<u8>, Vec<u8>);

    fn next(&mut self) -> Option<Self::Item> {
        use mkvs::Iterator;

        if !self.is_valid() {
            return None;
        }

        let key = self.key.as_ref().expect("iterator is valid").clone();
        let value = self.value.as_ref().expect("iterator is valid").clone();
        TreeIterator::next(self);

        Some((key, value))
    }
}

impl mkvs::Iterator for TreeIterator<'_> {
    fn set_prefetch(&mut self, prefetch: usize) {
        self.prefetch = prefetch;
    }

    fn is_valid(&self) -> bool {
        self.key.is_some()
    }

    fn error(&self) -> &Option<Error> {
        &self.error
    }

    fn rewind(&mut self) {
        self.seek(&[])
    }

    fn seek(&mut self, key: &[u8]) {
        if self.error.is_some() {
            return;
        }

        self.reset();
        let pending_root = self.tree.cache.borrow().get_pending_root();
        if let Err(error) = self._next(
            pending_root,
            0,
            Key::new(),
            key.to_vec(),
            VisitState::Before,
        ) {
            self.error = Some(error);
            self.reset();
        }
    }

    fn get_key(&self) -> &Option<Key> {
        &self.key
    }

    fn get_value(&self) -> &Option<Vec<u8>> {
        &self.value
    }

    fn next(&mut self) {
        TreeIterator::next(self)
    }
}

impl Tree {
    /// Return an iterator over the tree.
    pub fn iter(&self) -> TreeIterator<'_> {
        TreeIterator::new(self)
    }
}

#[cfg(test)]
pub(super) mod test {
    use std::iter;

    use rustc_hex::FromHex;

    use super::{super::tree_test::generate_key_value_pairs_ex, *};
    use crate::storage::mkvs::{
        interop::{Driver, ProtocolServer},
        sync::{NoopReadSyncer, StatsCollector},
        Iterator, OverlayTree, RootType,
    };

    #[test]
    fn test_iterator() {
        let server = ProtocolServer::new(None);

        let mut tree = Tree::builder()
            .with_root_type(RootType::State)
            .build(Box::new(NoopReadSyncer));

        // Test with an empty tree.
        let mut it = tree.iter();
        it.rewind();
        assert!(
            !it.is_valid(),
            "iterator should be invalid on an empty tree"
        );

        // Test with one item.
        tree.insert(b"key", b"first").unwrap();
        let mut it = tree.iter();
        it.rewind();
        assert!(
            it.is_valid(),
            "iterator should be valid on a non-empty tree"
        );

        // Insert some items.
        let items = vec![
            (b"key".to_vec(), b"first".to_vec()),
            (b"key 1".to_vec(), b"one".to_vec()),
            (b"key 2".to_vec(), b"two".to_vec()),
            (b"key 5".to_vec(), b"five".to_vec()),
            (b"key 8".to_vec(), b"eight".to_vec()),
            (b"key 9".to_vec(), b"nine".to_vec()),
        ];
        for (key, value) in items.iter() {
            tree.insert(key, value).unwrap();
        }

        let tests = vec![
            (b"k".to_vec(), 0),
            (b"key 1".to_vec(), 1),
            (b"key 3".to_vec(), 3),
            (b"key 4".to_vec(), 3),
            (b"key 5".to_vec(), 3),
            (b"key 6".to_vec(), 4),
            (b"key 7".to_vec(), 4),
            (b"key 8".to_vec(), 4),
            (b"key 9".to_vec(), 5),
            (b"key A".to_vec(), -1),
        ];

        // Direct.
        let it = tree.iter();
        test_iterator_with(&items, it, &tests);

        // Remote.
        let hash = tree.commit(Default::default(), 0).expect("commit");
        let write_log = items
            .iter()
            .cloned()
            .map(|(key, value)| mkvs::LogEntry {
                key,
                value: Some(value),
            })
            .collect();
        server.apply(&write_log, hash, Default::default(), 0);

        let remote_tree = Tree::builder()
            .with_capacity(0, 0)
            .with_root(Root {
                root_type: RootType::State,
                hash,
                ..Default::default()
            })
            .build(server.read_sync());

        let it = remote_tree.iter();
        test_iterator_with(&items, it, &tests);

        // Remote with prefetch (10).
        let stats = StatsCollector::new(server.read_sync());
        let remote_tree = Tree::builder()
            .with_capacity(0, 0)
            .with_root(Root {
                root_type: RootType::State,
                hash,
                ..Default::default()
            })
            .build(Box::new(stats));

        let mut it = remote_tree.iter();
        it.set_prefetch(10);
        test_iterator_with(&items, it, &tests);

        let cache = remote_tree.cache.borrow();
        let stats = cache
            .get_read_syncer()
            .as_any()
            .downcast_ref::<StatsCollector>()
            .expect("stats");
        assert_eq!(0, stats.sync_get_count, "sync_get_count");
        assert_eq!(0, stats.sync_get_prefixes_count, "sync_get_prefixes_count");
        assert_eq!(1, stats.sync_iterate_count, "sync_iterate_count");

        // Remote with prefetch (3).
        let stats = StatsCollector::new(server.read_sync());
        let remote_tree = Tree::builder()
            .with_capacity(0, 0)
            .with_root(Root {
                root_type: RootType::State,
                hash,
                ..Default::default()
            })
            .build(Box::new(stats));

        let mut it = remote_tree.iter();
        it.set_prefetch(3);
        test_iterator_with(&items, it, &tests);

        let cache = remote_tree.cache.borrow();
        let stats = cache
            .get_read_syncer()
            .as_any()
            .downcast_ref::<StatsCollector>()
            .expect("stats");
        assert_eq!(0, stats.sync_get_count, "sync_get_count");
        assert_eq!(0, stats.sync_get_prefixes_count, "sync_get_prefixes_count");
        assert_eq!(2, stats.sync_iterate_count, "sync_iterate_count");
    }

    #[test]
    fn test_iterator_case1() {
        let mut tree = Tree::builder()
            .with_root_type(RootType::State)
            .build(Box::new(NoopReadSyncer));

        let items = vec![
            (b"key 5".to_vec(), b"fivey".to_vec()),
            (b"key 7".to_vec(), b"seven".to_vec()),
        ];
        for (key, value) in items.iter() {
            tree.insert(key, value).unwrap();
        }

        let tests = vec![(b"key 3".to_vec(), 0)];

        let it = tree.iter();
        test_iterator_with(&items, it, &tests);
    }

    #[test]
    fn test_iterator_case2() {
        let mut tree = Tree::builder()
            .with_root_type(RootType::State)
            .build(Box::new(NoopReadSyncer));

        let items: Vec<(Vec<u8>, Vec<u8>)> = vec![
            (
                "54dcb497eb46bc7cb1a1a29d143d5d41f1a684c97e12f2ae536eceb828c15fc34c02"
                    .from_hex()
                    .unwrap(),
                b"value".to_vec(),
            ),
            (
                "54dcb497eb46bc7cb1a1a29d143d5d41f1a684c97e12f2ae536eceb828c15fc34c02"
                    .from_hex()
                    .unwrap(),
                b"value".to_vec(),
            ),
        ];
        for (key, value) in items.iter() {
            tree.insert(key, value).unwrap();
        }

        let mut it = tree.iter();
        let missing_key: Vec<u8> =
            "54da85be3251772db943cba67341d402117c87ada2a9e8aad7171d40b6b4dc9fbc"
                .from_hex()
                .unwrap();
        it.seek(&missing_key);
        assert!(it.is_valid(), "iterator should be valid");
        let item = iter::Iterator::next(&mut it);
        assert_eq!(
            Some((items[0].0.clone(), b"value".to_vec())),
            item,
            "value should be correct"
        );
    }

    #[test]
    fn test_iterator_eviction() {
        let server = ProtocolServer::new(None);

        let mut tree = OverlayTree::new(
            Tree::builder()
                .with_capacity(0, 0)
                .with_root_type(RootType::State)
                .build(Box::new(NoopReadSyncer)),
        );

        let (keys, values) = generate_key_value_pairs_ex("T".to_owned(), 100);
        let items: Vec<(Vec<u8>, Vec<u8>)> = keys.into_iter().zip(values.into_iter()).collect();
        for (key, value) in &items {
            tree.insert(&key, &value).unwrap();
        }

        let (write_log, hash) = tree.commit_both(Default::default(), 0).expect("commit");
        server.apply(&write_log, hash, Default::default(), 0);

        // Create a remote tree with limited cache capacity so that nodes will
        // be evicted while iterating.
        let stats = StatsCollector::new(server.read_sync());
        let remote_tree = Tree::builder()
            .with_capacity(50, 16 * 1024 * 1024)
            .with_root(Root {
                root_type: RootType::State,
                hash,
                ..Default::default()
            })
            .build(Box::new(stats));

        let mut it = remote_tree.iter();
        it.set_prefetch(1000);
        test_iterator_with(&items, it, &vec![]);

        let cache = remote_tree.cache.borrow();
        let stats = cache
            .get_read_syncer()
            .as_any()
            .downcast_ref::<StatsCollector>()
            .expect("stats");
        assert_eq!(0, stats.sync_get_count, "sync_get_count");
        assert_eq!(0, stats.sync_get_prefixes_count, "sync_get_prefixes_count");
        // We require multiple fetches as we can only store a limited amount of
        // results per fetch due to the cache being too small.
        assert_eq!(2, stats.sync_iterate_count, "sync_iterate_count");
    }

    pub(in super::super) fn test_iterator_with<I: mkvs::Iterator>(
        items: &[(Vec<u8>, Vec<u8>)],
        mut it: I,
        tests: &[(Vec<u8>, isize)],
    ) {
        // Iterate through the whole tree.
        let mut iterations = 0;
        it.rewind();
        for (idx, (key, value)) in it.by_ref().enumerate() {
            if !tests.is_empty() {
                assert_eq!(items[idx].0, key, "iterator should have the correct key");
                assert_eq!(
                    items[idx].1, value,
                    "iterator should have the correct value"
                );
            }
            iterations += 1;
        }
        assert!(it.error().is_none(), "iterator should not error");
        assert_eq!(iterations, items.len(), "iterator should go over all items");

        for (seek, pos) in tests {
            it.seek(&seek);
            if *pos == -1 {
                assert!(!it.is_valid(), "iterator should not be valid after seek");
                continue;
            }

            for expected in &items[*pos as usize..] {
                let item = iter::Iterator::next(&mut it);
                assert_eq!(
                    Some(expected.clone()),
                    item,
                    "iterator should have the correct item"
                );
            }
        }
    }
}

1	//! Tree iterator.
2	use std::{collections::VecDeque, fmt};
3
4	use anyhow::{Error, Result};
5
6	use crate::storage::mkvs::{
7	self,
8	cache::{Cache, ReadSyncFetcher},
9	sync::{IterateRequest, Proof, ReadSync, TreeID},
10	tree::{Depth, Key, KeyTrait, NodeBox, NodeKind, NodePtrRef, Root, Tree},
11	};
12
13	pub(super) struct FetcherSyncIterate<'a> {
14	key: &'a Key,
15	prefetch: usize,
16	}
17
18	impl<'a> FetcherSyncIterate<'a> {
19	pub(super) fn new(key: &'a Key, prefetch: usize) -> Self {	1✔
20	Self { key, prefetch }
21	}
22	}
23
24	impl ReadSyncFetcher for FetcherSyncIterate<'_> {
25	fn fetch(&self, root: Root, ptr: NodePtrRef, rs: &mut Box<dyn ReadSync>) -> Result<Proof> {	1✔
26	let rsp = rs.sync_iterate(IterateRequest {	2✔
27	tree: TreeID {	1✔
28	root,
29	position: ptr.borrow().hash,	3✔
30	},
31	key: self.key.clone(),	1✔
32	prefetch: self.prefetch as u16,	1✔
33	})?;
34	Ok(rsp.proof)	1✔
35	}
36	}
37
38	/// Visit state of a node.
39	#[derive(Debug, PartialEq)]
40	enum VisitState {
41	Before,
42	At,
43	AtLeft,
44	After,
45	}
46
47	/// Atom in the current iterator path. Can be used to resume iteration
48	/// from a given position.
49	struct PathAtom {
50	ptr: NodePtrRef,
51	bit_depth: Depth,
52	path: Key,
53	state: VisitState,
54	}
55
56	impl fmt::Debug for PathAtom {
57	fn fmt(&self, f: &mut fmt::Formatter) -> std::result::Result<(), fmt::Error> {	×
58	f.debug_struct("PathAtom")	×
59	.field("bit_depth", &self.bit_depth)	×
60	.field("path", &self.path)	×
61	.field("state", &self.state)	×
62	.finish()
63	}
64	}
65
66	/// Tree iterator.
67	pub struct TreeIterator<'tree> {
68	tree: &'tree Tree,
69	prefetch: usize,
70	pos: VecDeque<PathAtom>,
71	key: Option<Key>,
72	value: Option<Vec<u8>>,
73	error: Option<Error>,
74	}
75
76	impl<'tree> TreeIterator<'tree> {
77	/// Create a new tree iterator.
78	fn new(tree: &'tree Tree) -> Self {	1✔
79	Self {
80	tree,
81	prefetch: 0,
82	pos: VecDeque::new(),	1✔
83	key: None,
84	value: None,
85	error: None,
86	}
87	}
88
89	fn reset(&mut self) {	1✔
90	self.pos.clear();	1✔
91	self.key = None;	1✔
92	self.value = None;	1✔
93	}
94
95	fn next(&mut self) {	1✔
96	if self.error.is_some() {	1✔
97	return;	×
98	}
99
100	while !self.pos.is_empty() {	2✔
101	// Start where we left off.
102	let atom = self.pos.pop_front().expect("not empty");	1✔
103	let mut remainder = std::mem::take(&mut self.pos);	2✔
104
105	// Remember where the path from root to target node ends (will end).
106	let mut cache = self.tree.cache.borrow_mut();	2✔
107	cache.mark_position();	2✔
108	for atom in &remainder {	1✔
109	cache.use_node(atom.ptr.clone());	2✔
110	}
111	drop(cache);	1✔
112
113	// Try to proceed with the current node. If we don't succeed, proceed to the
114	// next node.
115	let key = self.key.take().expect("iterator is valid");	1✔
116	self.reset();	1✔
117	if let Err(error) =	1✔
118	self._next(atom.ptr, atom.bit_depth, atom.path, key.clone(), atom.state)	×
119	{
120	self.error = Some(error);	×
121	self.reset();	×
122	return;	×
123	}
124	if self.key.is_some() {	2✔
125	// Key has been found.
126	self.pos.append(&mut remainder);	2✔
127	return;	×
128	}
129
130	self.key = Some(key);	1✔
131	self.pos = remainder;	1✔
132	}
133
134	// We have reached the end of the tree, make sure everything is reset.
135	self.key = None;	1✔
136	self.value = None;	1✔
137	}
138
139	fn _next(	1✔
140	&mut self,
141	ptr: NodePtrRef,
142	bit_depth: Depth,
143	path: Key,
144	mut key: Key,
145	mut state: VisitState,
146	) -> Result<()> {
147	let node_ref = self.tree.cache.borrow_mut().deref_node_ptr(	5✔
148	ptr.clone(),	2✔
149	Some(FetcherSyncIterate::new(&key, self.prefetch)),	1✔
150	)?;
151
152	match classify_noderef!(?node_ref) {	1✔
153	NodeKind::None => {	×
154	// Reached a nil node, there is nothing here.
155	Ok(())	1✔
156	}
157	NodeKind::Internal => {	×
158	let node_ref = node_ref.unwrap();	2✔
159	if let NodeBox::Internal(ref n) = *node_ref.borrow() {	4✔
160	// Internal node.
161	let bit_length = bit_depth + n.label_bit_length;	1✔
162	let new_path = path.merge(bit_depth, &n.label, n.label_bit_length);	1✔
163
164	// Check if the key is longer than the current path but lexicographically smaller. In this
165	// case everything in this subtree will be larger so we need to take the first value.
166	let take_first =	3✔
167	bit_length > 0 && key.bit_length() >= bit_length && key < new_path;	×
168
169	// Does lookup key end here? Look into LeafNode.
170	if (state == VisitState::Before	2✔
171	&& (key.bit_length() <= bit_length \|\| take_first))	1✔
172	\|\| state == VisitState::At	2✔
173	{
174	if state == VisitState::Before {	2✔
175	self._next(	2✔
176	n.leaf_node.clone(),	1✔
177	bit_length,	×
178	path.clone(),	2✔
179	key.clone(),	1✔
UNCOV 180	VisitState::Before,	×
181	)?;
182	if self.key.is_some() {	1✔
183	// Key has been found.
184	self.pos.push_back(PathAtom {	1✔
185	ptr,	1✔
186	bit_depth,	×
187	path,	1✔
UNCOV 188	state: VisitState::At,	×
189	});
190	return Ok(());	1✔
191	}
192	}
193	// Key has not been found, continue with search for next key.
194	if key.bit_length() <= bit_length {	3✔
195	key = key.append_bit(bit_length, false);	1✔
196	}
197	}
198
199	if state == VisitState::Before {	3✔
200	state = VisitState::At;	1✔
201	}
202
203	// Continue recursively based on a bit value.
204	if (state == VisitState::At && (!key.get_bit(bit_length) \|\| take_first))	5✔
205	\|\| state == VisitState::AtLeft	2✔
206	{
207	if state == VisitState::At {	2✔
208	self._next(	2✔
209	n.left.clone(),	1✔
210	bit_length,	×
211	new_path.append_bit(bit_length, false),	2✔
212	key.clone(),	1✔
UNCOV 213	VisitState::Before,	×
214	)?;
215	if self.key.is_some() {	1✔
216	// Key has been found.
217	self.pos.push_back(PathAtom {	1✔
218	ptr,	1✔
219	bit_depth,	×
220	path,	1✔
UNCOV 221	state: VisitState::AtLeft,	×
222	});
223	return Ok(());	1✔
224	}
225	}
226	// Key has not been found, continue with search for next key.
227	key = key.split(bit_length, key.bit_length()).0;	2✔
228	key = key.append_bit(bit_length, true);	1✔
229	}
230
231	if state == VisitState::At \|\| state == VisitState::AtLeft {	3✔
232	self._next(	2✔
233	n.right.clone(),	2✔
234	bit_length,	×
235	new_path.append_bit(bit_length, true),	1✔
236	key,	1✔
UNCOV 237	VisitState::Before,	×
238	)?;
239	if self.key.is_some() {	1✔
240	// Key has been found.
241	self.pos.push_back(PathAtom {	1✔
242	ptr,	1✔
243	bit_depth,	×
244	path,	1✔
UNCOV 245	state: VisitState::After,	×
246	});
247	return Ok(());	1✔
248	}
249	}
250
251	return Ok(());	1✔
252	}
253
254	unreachable!("node kind is internal node");
255	}
256	NodeKind::Leaf => {
257	// Reached a leaf node.
258	let node_ref = node_ref.unwrap();
259	if let NodeBox::Leaf(ref n) = *node_ref.borrow() {
260	if n.key >= key {
261	self.key = Some(n.key.clone());
262	self.value = Some(n.value.clone());
263	}
264	} else {
265	unreachable!("node kind is leaf node");
266	}
267
268	Ok(())
269	}
270	}
271	}
272	}
273
274	impl Iterator for TreeIterator<'_> {
275	type Item = (Vec<u8>, Vec<u8>);
276
277	fn next(&mut self) -> Option<Self::Item> {	1✔
278	use mkvs::Iterator;
279
280	if !self.is_valid() {	1✔
281	return None;	1✔
282	}
283
284	let key = self.key.as_ref().expect("iterator is valid").clone();	1✔
285	let value = self.value.as_ref().expect("iterator is valid").clone();	2✔
286	TreeIterator::next(self);	1✔
287
288	Some((key, value))	1✔
289	}
290	}
291
292	impl mkvs::Iterator for TreeIterator<'_> {
293	fn set_prefetch(&mut self, prefetch: usize) {	1✔
294	self.prefetch = prefetch;	1✔
295	}
296
297	fn is_valid(&self) -> bool {	1✔
298	self.key.is_some()	1✔
299	}
300
301	fn error(&self) -> &Option<Error> {	1✔
302	&self.error	1✔
303	}
304
305	fn rewind(&mut self) {	1✔
306	self.seek(&[])	1✔
307	}
308
309	fn seek(&mut self, key: &[u8]) {	1✔
310	if self.error.is_some() {	1✔
311	return;
312	}
313
314	self.reset();	1✔
315	let pending_root = self.tree.cache.borrow().get_pending_root();	1✔
316	if let Err(error) = self._next(	2✔
317	pending_root,	1✔
318	0,
319	Key::new(),	1✔
320	key.to_vec(),	1✔
321	VisitState::Before,
322	) {
323	self.error = Some(error);	×
324	self.reset();	×
325	}
326	}
327
328	fn get_key(&self) -> &Option<Key> {	1✔
329	&self.key	1✔
330	}
331
332	fn get_value(&self) -> &Option<Vec<u8>> {	1✔
333	&self.value	1✔
334	}
335
336	fn next(&mut self) {	×
337	TreeIterator::next(self)	×
338	}
339	}
340
341	impl Tree {
342	/// Return an iterator over the tree.
343	pub fn iter(&self) -> TreeIterator<'_> {	1✔
344	TreeIterator::new(self)	1✔
345	}
346	}
347
348	#[cfg(test)]
349	pub(super) mod test {
350	use std::iter;
351
352	use rustc_hex::FromHex;
353
354	use super::{super::tree_test::generate_key_value_pairs_ex, *};
355	use crate::storage::mkvs::{
356	interop::{Driver, ProtocolServer},
357	sync::{NoopReadSyncer, StatsCollector},
358	Iterator, OverlayTree, RootType,
359	};
360
361	#[test]
362	fn test_iterator() {
363	let server = ProtocolServer::new(None);
364
365	let mut tree = Tree::builder()
366	.with_root_type(RootType::State)
367	.build(Box::new(NoopReadSyncer));
368
369	// Test with an empty tree.
370	let mut it = tree.iter();
371	it.rewind();
372	assert!(
373	!it.is_valid(),
374	"iterator should be invalid on an empty tree"
375	);
376
377	// Test with one item.
378	tree.insert(b"key", b"first").unwrap();
379	let mut it = tree.iter();
380	it.rewind();
381	assert!(
382	it.is_valid(),
383	"iterator should be valid on a non-empty tree"
384	);
385
386	// Insert some items.
387	let items = vec![
388	(b"key".to_vec(), b"first".to_vec()),
389	(b"key 1".to_vec(), b"one".to_vec()),
390	(b"key 2".to_vec(), b"two".to_vec()),
391	(b"key 5".to_vec(), b"five".to_vec()),
392	(b"key 8".to_vec(), b"eight".to_vec()),
393	(b"key 9".to_vec(), b"nine".to_vec()),
394	];
395	for (key, value) in items.iter() {
396	tree.insert(key, value).unwrap();
397	}
398
399	let tests = vec![
400	(b"k".to_vec(), 0),
401	(b"key 1".to_vec(), 1),
402	(b"key 3".to_vec(), 3),
403	(b"key 4".to_vec(), 3),
404	(b"key 5".to_vec(), 3),
405	(b"key 6".to_vec(), 4),
406	(b"key 7".to_vec(), 4),
407	(b"key 8".to_vec(), 4),
408	(b"key 9".to_vec(), 5),
409	(b"key A".to_vec(), -1),
410	];
411
412	// Direct.
413	let it = tree.iter();
414	test_iterator_with(&items, it, &tests);
415
416	// Remote.
417	let hash = tree.commit(Default::default(), 0).expect("commit");
418	let write_log = items
419	.iter()
420	.cloned()
421	.map(\|(key, value)\| mkvs::LogEntry {
422	key,
423	value: Some(value),
424	})
425	.collect();
426	server.apply(&write_log, hash, Default::default(), 0);
427
428	let remote_tree = Tree::builder()
429	.with_capacity(0, 0)
430	.with_root(Root {
431	root_type: RootType::State,
432	hash,
433	..Default::default()
434	})
435	.build(server.read_sync());
436
437	let it = remote_tree.iter();
438	test_iterator_with(&items, it, &tests);
439
440	// Remote with prefetch (10).
441	let stats = StatsCollector::new(server.read_sync());
442	let remote_tree = Tree::builder()
443	.with_capacity(0, 0)
444	.with_root(Root {
445	root_type: RootType::State,
446	hash,
447	..Default::default()
448	})
449	.build(Box::new(stats));
450
451	let mut it = remote_tree.iter();
452	it.set_prefetch(10);
453	test_iterator_with(&items, it, &tests);
454
455	let cache = remote_tree.cache.borrow();
456	let stats = cache
457	.get_read_syncer()
458	.as_any()
459	.downcast_ref::<StatsCollector>()
460	.expect("stats");
461	assert_eq!(0, stats.sync_get_count, "sync_get_count");
462	assert_eq!(0, stats.sync_get_prefixes_count, "sync_get_prefixes_count");
463	assert_eq!(1, stats.sync_iterate_count, "sync_iterate_count");
464
465	// Remote with prefetch (3).
466	let stats = StatsCollector::new(server.read_sync());
467	let remote_tree = Tree::builder()
468	.with_capacity(0, 0)
469	.with_root(Root {
470	root_type: RootType::State,
471	hash,
472	..Default::default()
473	})
474	.build(Box::new(stats));
475
476	let mut it = remote_tree.iter();
477	it.set_prefetch(3);
478	test_iterator_with(&items, it, &tests);
479
480	let cache = remote_tree.cache.borrow();
481	let stats = cache
482	.get_read_syncer()
483	.as_any()
484	.downcast_ref::<StatsCollector>()
485	.expect("stats");
486	assert_eq!(0, stats.sync_get_count, "sync_get_count");
487	assert_eq!(0, stats.sync_get_prefixes_count, "sync_get_prefixes_count");
488	assert_eq!(2, stats.sync_iterate_count, "sync_iterate_count");
489	}
490
491	#[test]
492	fn test_iterator_case1() {
493	let mut tree = Tree::builder()
494	.with_root_type(RootType::State)
495	.build(Box::new(NoopReadSyncer));
496
497	let items = vec![
498	(b"key 5".to_vec(), b"fivey".to_vec()),
499	(b"key 7".to_vec(), b"seven".to_vec()),
500	];
501	for (key, value) in items.iter() {
502	tree.insert(key, value).unwrap();
503	}
504
505	let tests = vec![(b"key 3".to_vec(), 0)];
506
507	let it = tree.iter();
508	test_iterator_with(&items, it, &tests);
509	}
510
511	#[test]
512	fn test_iterator_case2() {
513	let mut tree = Tree::builder()
514	.with_root_type(RootType::State)
515	.build(Box::new(NoopReadSyncer));
516
517	let items: Vec<(Vec<u8>, Vec<u8>)> = vec![
518	(
519	"54dcb497eb46bc7cb1a1a29d143d5d41f1a684c97e12f2ae536eceb828c15fc34c02"
520	.from_hex()
521	.unwrap(),
522	b"value".to_vec(),
523	),
524	(
525	"54dcb497eb46bc7cb1a1a29d143d5d41f1a684c97e12f2ae536eceb828c15fc34c02"
526	.from_hex()
527	.unwrap(),
528	b"value".to_vec(),
529	),
530	];
531	for (key, value) in items.iter() {
532	tree.insert(key, value).unwrap();
533	}
534
535	let mut it = tree.iter();
536	let missing_key: Vec<u8> =
537	"54da85be3251772db943cba67341d402117c87ada2a9e8aad7171d40b6b4dc9fbc"
538	.from_hex()
539	.unwrap();
540	it.seek(&missing_key);
541	assert!(it.is_valid(), "iterator should be valid");
542	let item = iter::Iterator::next(&mut it);
543	assert_eq!(
544	Some((items[0].0.clone(), b"value".to_vec())),
545	item,
546	"value should be correct"
547	);
548	}
549
550	#[test]
551	fn test_iterator_eviction() {
552	let server = ProtocolServer::new(None);
553
554	let mut tree = OverlayTree::new(
555	Tree::builder()
556	.with_capacity(0, 0)
557	.with_root_type(RootType::State)
558	.build(Box::new(NoopReadSyncer)),
559	);
560
561	let (keys, values) = generate_key_value_pairs_ex("T".to_owned(), 100);
562	let items: Vec<(Vec<u8>, Vec<u8>)> = keys.into_iter().zip(values.into_iter()).collect();
563	for (key, value) in &items {
564	tree.insert(&key, &value).unwrap();
565	}
566
567	let (write_log, hash) = tree.commit_both(Default::default(), 0).expect("commit");
568	server.apply(&write_log, hash, Default::default(), 0);
569
570	// Create a remote tree with limited cache capacity so that nodes will
571	// be evicted while iterating.
572	let stats = StatsCollector::new(server.read_sync());
573	let remote_tree = Tree::builder()
574	.with_capacity(50, 16 * 1024 * 1024)
575	.with_root(Root {
576	root_type: RootType::State,
577	hash,
578	..Default::default()
579	})
580	.build(Box::new(stats));
581
582	let mut it = remote_tree.iter();
583	it.set_prefetch(1000);
584	test_iterator_with(&items, it, &vec![]);
585
586	let cache = remote_tree.cache.borrow();
587	let stats = cache
588	.get_read_syncer()
589	.as_any()
590	.downcast_ref::<StatsCollector>()
591	.expect("stats");
592	assert_eq!(0, stats.sync_get_count, "sync_get_count");
593	assert_eq!(0, stats.sync_get_prefixes_count, "sync_get_prefixes_count");
594	// We require multiple fetches as we can only store a limited amount of
595	// results per fetch due to the cache being too small.
596	assert_eq!(2, stats.sync_iterate_count, "sync_iterate_count");
597	}
598
599	pub(in super::super) fn test_iterator_with<I: mkvs::Iterator>(
600	items: &[(Vec<u8>, Vec<u8>)],
601	mut it: I,
602	tests: &[(Vec<u8>, isize)],
603	) {
604	// Iterate through the whole tree.
605	let mut iterations = 0;
606	it.rewind();
607	for (idx, (key, value)) in it.by_ref().enumerate() {
608	if !tests.is_empty() {
609	assert_eq!(items[idx].0, key, "iterator should have the correct key");
610	assert_eq!(
611	items[idx].1, value,
612	"iterator should have the correct value"
613	);
614	}
615	iterations += 1;
616	}
617	assert!(it.error().is_none(), "iterator should not error");
618	assert_eq!(iterations, items.len(), "iterator should go over all items");
619
620	for (seek, pos) in tests {
621	it.seek(&seek);
622	if *pos == -1 {
623	assert!(!it.is_valid(), "iterator should not be valid after seek");
624	continue;
625	}
626
627	for expected in &items[*pos as usize..] {
628	let item = iter::Iterator::next(&mut it);
629	assert_eq!(
630	Some(expected.clone()),
631	item,
632	"iterator should have the correct item"
633	);
634	}
635	}
636	}
637	}

oasisprotocol / oasis-core / #6870

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous