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

Jump to bottom

wip remove empty case from trie node #12478

Draft
wants to merge 10 commits into
base: master
Choose a base branch
from
Draft

wip remove empty case from trie node #12478

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
e7b7cfb
wip, issue with crumbs
stedfn Nov 19, 2024
6e18a1e
fix increment fn
stedfn Nov 19, 2024
6ceb9a6
consistency tests failing
stedfn Nov 19, 2024
08c9d6b
remove comments
stedfn Nov 19, 2024
5ab8ae7
remove disk iter unwrap
stedfn Nov 19, 2024
fe4c78e
remove unwraps
stedfn Nov 20, 2024
257f9eb
remove trienode
stedfn Nov 20, 2024
0f476ac
remove comment modify iter_step
stedfn Nov 20, 2024
7c16ef5
fix test
stedfn Nov 20, 2024
5f5c047
Merge branch 'master' into stefan/remove_empty_case
stedfn Nov 27, 2024
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
208 changes: 115 additions & 93 deletions core/store/src/trie/iterator.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,15 +1,15 @@
use near_primitives::hash::CryptoHash;

use crate::trie::nibble_slice::NibbleSlice;
use crate::trie::{TrieNode, TrieNodeWithSize, ValueHandle};
use crate::{MissingTrieValueContext, StorageError, Trie};
use crate::{StorageError, Trie};

use super::mem::iter::STMemTrieIterator;
use super::{RawTrieNode, RawTrieNodeWithSize};

/// Crumb is a piece of trie iteration state. It describes a node on the trail and processing status of that node.
#[derive(Debug)]
struct Crumb {
node: TrieNodeWithSize,
node: Option<RawTrieNodeWithSize>,
status: CrumbStatus,
prefix_boundary: bool,
}
Expand All @@ -31,13 +31,19 @@ impl Crumb {
self.status = CrumbStatus::Exiting;
return;
}
self.status = match (&self.status, &self.node.node) {
(_, &TrieNode::Empty) => CrumbStatus::Exiting,
self.status = match (&self.status, &self.node) {
(_, None) => CrumbStatus::Exiting,
(&CrumbStatus::Entering, _) => CrumbStatus::At,
(&CrumbStatus::At, &TrieNode::Branch(_, _)) => CrumbStatus::AtChild(0),
(&CrumbStatus::AtChild(x), &TrieNode::Branch(_, _)) if x < 15 => {
CrumbStatus::AtChild(x + 1)
}
(&CrumbStatus::At, Some(node)) => match node.node {
RawTrieNode::BranchNoValue(_) => CrumbStatus::AtChild(0),
RawTrieNode::BranchWithValue(_, _) => CrumbStatus::AtChild(0),
_ => CrumbStatus::Exiting,
},
(&CrumbStatus::AtChild(x), Some(node)) if x < 15 => match node.node {
RawTrieNode::BranchNoValue(_) => CrumbStatus::AtChild(x + 1),
RawTrieNode::BranchWithValue(_, _) => CrumbStatus::AtChild(x + 1),
_ => CrumbStatus::Exiting,
},
_ => CrumbStatus::Exiting,
}
}
Expand Down Expand Up @@ -154,46 +160,49 @@ impl<'a> DiskTrieIterator<'a> {
self.descend_into_node(&hash)?;
let Crumb { status, node, prefix_boundary } = self.trail.last_mut().unwrap();
prev_prefix_boundary = prefix_boundary;
match &node.node {
TrieNode::Empty => break,
TrieNode::Leaf(leaf_key, _) => {
let existing_key = NibbleSlice::from_encoded(leaf_key).0;
if !check_ext_key(&key, &existing_key) {
self.key_nibbles.extend(existing_key.iter());
*status = CrumbStatus::Exiting;
}
break;
}
TrieNode::Branch(children, _) => {
if key.is_empty() {
match &node {
None => break,
Some(node) => match &node.node {
RawTrieNode::Leaf(leaf_key, _) => {
let existing_key = NibbleSlice::from_encoded(leaf_key).0;
if !check_ext_key(&key, &existing_key) {
self.key_nibbles.extend(existing_key.iter());
*status = CrumbStatus::Exiting;
}
break;
}
let idx = key.at(0);
self.key_nibbles.push(idx);
*status = CrumbStatus::AtChild(idx);
if let Some(ref child) = children[idx] {
hash = *child.unwrap_hash();
key = key.mid(1);
} else {
*prefix_boundary = is_prefix_seek;
break;
RawTrieNode::BranchNoValue(children)
| RawTrieNode::BranchWithValue(_, children) => {
if key.is_empty() {
break;
}
let idx = key.at(0);
self.key_nibbles.push(idx);
*status = CrumbStatus::AtChild(idx);
if let Some(ref child) = children[idx] {
hash = *child;
key = key.mid(1);
} else {
*prefix_boundary = is_prefix_seek;
break;
}
}
}
TrieNode::Extension(ext_key, child) => {
let existing_key = NibbleSlice::from_encoded(ext_key).0;
if key.starts_with(&existing_key) {
key = key.mid(existing_key.len());
hash = *child.unwrap_hash();
*status = CrumbStatus::At;
self.key_nibbles.extend(existing_key.iter());
} else {
if !check_ext_key(&key, &existing_key) {
*status = CrumbStatus::Exiting;
RawTrieNode::Extension(ext_key, child) => {
let existing_key = NibbleSlice::from_encoded(ext_key).0;
if key.starts_with(&existing_key) {
key = key.mid(existing_key.len());
hash = *child;
*status = CrumbStatus::At;
self.key_nibbles.extend(existing_key.iter());
} else {
if !check_ext_key(&key, &existing_key) {
*status = CrumbStatus::Exiting;
self.key_nibbles.extend(existing_key.iter());
}
break;
}
break;
}
}
},
}
}
Ok(hash)
Expand All @@ -206,13 +215,22 @@ impl<'a> DiskTrieIterator<'a> {
/// with [`Self::remember_visited_nodes`]), the node will be added to the
/// list.
fn descend_into_node(&mut self, hash: &CryptoHash) -> Result<(), StorageError> {
let (bytes, node) = self.trie.retrieve_node(hash)?;
if let Some(ref mut visited) = self.visited_nodes {
visited.push(bytes.ok_or({
StorageError::MissingTrieValue(MissingTrieValueContext::TrieIterator, *hash)
})?);
if let Some((bytes, node)) = self.trie.retrieve_raw_node(hash, true, true)? {
if let Some(ref mut visited) = self.visited_nodes {
visited.push(bytes);
}
self.trail.push(Crumb {
status: CrumbStatus::Entering,
node: Some(node),
prefix_boundary: false,
});
} else {
self.trail.push(Crumb {
status: CrumbStatus::Entering,
node: None,
prefix_boundary: false,
});
}
self.trail.push(Crumb { status: CrumbStatus::Entering, node, prefix_boundary: false });
Ok(())
}

Expand All @@ -227,7 +245,13 @@ impl<'a> DiskTrieIterator<'a> {
fn has_value(&self) -> bool {
match self.trail.last() {
Some(b) => match &b.status {
CrumbStatus::At => b.node.node.has_value(),
CrumbStatus::At => match &b.node {
Some(node) => match node.node {
RawTrieNode::Leaf(_, _) | RawTrieNode::BranchWithValue(_, _) => true,
_ => false,
},
None => false,
},
_ => false,
},
None => false, // Trail finished
Expand All @@ -237,57 +261,55 @@ impl<'a> DiskTrieIterator<'a> {
fn iter_step(&mut self) -> Option<IterStep> {
let last = self.trail.last_mut()?;
last.increment();
Some(match (last.status, &last.node.node) {
Some(match (last.status, &last.node) {
(CrumbStatus::Exiting, n) => {
match n {
TrieNode::Leaf(ref key, _) | TrieNode::Extension(ref key, _) => {
let existing_key = NibbleSlice::from_encoded(key).0;
let l = self.key_nibbles.len();
self.key_nibbles.truncate(l - existing_key.len());
}
TrieNode::Branch(_, _) => {
self.key_nibbles.pop();
if let Some(node) = n {
match node.node {
RawTrieNode::Leaf(ref key, _) | RawTrieNode::Extension(ref key, _) => {
let existing_key = NibbleSlice::from_encoded(key).0;
let l = self.key_nibbles.len();
self.key_nibbles.truncate(l - existing_key.len());
}
RawTrieNode::BranchNoValue(_) | RawTrieNode::BranchWithValue(_, _) => {
self.key_nibbles.pop();
}
}
_ => {}
}
IterStep::PopTrail
}
(CrumbStatus::At, TrieNode::Branch(_, Some(value))) => {
let hash = match value {
ValueHandle::HashAndSize(value) => value.hash,
ValueHandle::InMemory(_node) => unreachable!(),
};
IterStep::Value(hash)
}
(CrumbStatus::At, TrieNode::Branch(_, None)) => IterStep::Continue,
(CrumbStatus::At, TrieNode::Leaf(key, value)) => {
let hash = match value {
ValueHandle::HashAndSize(value) => value.hash,
ValueHandle::InMemory(_node) => unreachable!(),
};
let key = NibbleSlice::from_encoded(key).0;
self.key_nibbles.extend(key.iter());
IterStep::Value(hash)
}
(CrumbStatus::At, TrieNode::Extension(key, child)) => {
let hash = *child.unwrap_hash();
let key = NibbleSlice::from_encoded(key).0;
self.key_nibbles.extend(key.iter());
IterStep::Descend(hash)
}
(CrumbStatus::AtChild(i), TrieNode::Branch(children, _)) => {
if i == 0 {
self.key_nibbles.push(0);
(CrumbStatus::At, Some(node)) => match &node.node {
RawTrieNode::BranchWithValue(value, _) => IterStep::Value(value.hash),
RawTrieNode::BranchNoValue(_) => IterStep::Continue,
RawTrieNode::Leaf(key, value) => {
let key = NibbleSlice::from_encoded(key).0;
self.key_nibbles.extend(key.iter());
IterStep::Value(value.hash)
}
RawTrieNode::Extension(key, child) => {
let key = NibbleSlice::from_encoded(key).0;
self.key_nibbles.extend(key.iter());
IterStep::Descend(*child)
}
if let Some(ref child) = children[i] {
if i != 0 {
*self.key_nibbles.last_mut().expect("Pushed child value before") = i;
},
(CrumbStatus::AtChild(i), Some(node)) => match &node.node {
RawTrieNode::BranchNoValue(children)
| RawTrieNode::BranchWithValue(_, children) => {
if i == 0 {
self.key_nibbles.push(0);
}
if let Some(ref child) = children[i] {
if i != 0 {
*self.key_nibbles.last_mut().expect("Pushed child value before") = i;
}
IterStep::Descend(*child)
} else {
IterStep::Continue
}
IterStep::Descend(*child.unwrap_hash())
} else {
IterStep::Continue
}
}
_ => {
panic!("Should never see AtChild without a Branch here.");
}
},
_ => panic!("Should never see Entering or AtChild without a Branch here."),
})
}
Expand Down
Loading
Loading