refactor(trie): extract set_child/remove_child methods on ArenaSparseNodeBranch (#22936)

This commit is contained in:
Sergei Shulepov
2026-03-11 18:01:09 +07:00
committed by GitHub
parent ddc3ecaca6
commit fddf94c166
2 changed files with 29 additions and 16 deletions

View File

@@ -1377,9 +1377,6 @@ impl ArenaParallelSparseTrie {
}
SeekResult::NoChild { child_nibble } => {
let head_idx = head.index;
let head_branch = arena[head_idx].branch_ref();
let state_mask = head_branch.state_mask;
let insert_pos = BranchChildIdx::insertion_point(state_mask, child_nibble);
let head_branch_logical_path = cursor.head_logical_branch_path(arena);
let leaf_key = full_path.slice(head_branch_logical_path.len() + 1..);
@@ -1390,11 +1387,7 @@ impl ArenaParallelSparseTrie {
});
let branch = arena[head_idx].branch_mut();
branch.state_mask.set_bit(child_nibble);
branch
.children
.insert(insert_pos.get(), ArenaSparseNodeBranchChild::Revealed(new_leaf));
branch.state = ArenaSparseNodeState::Dirty;
branch.set_child(child_nibble, ArenaSparseNodeBranchChild::Revealed(new_leaf));
// Re-seek to position the cursor on the newly inserted leaf.
cursor.seek(arena, full_path);
@@ -1514,16 +1507,10 @@ impl ArenaParallelSparseTrie {
let parent_idx = parent_entry.index;
let child_nibble = head_path.last().expect("non-root leaf");
let parent_branch = arena[parent_idx].branch_ref();
let child_idx = BranchChildIdx::new(parent_branch.state_mask, child_nibble)
.expect("leaf nibble not found in parent state_mask");
// Remove the leaf from the arena and from the parent's children.
arena.remove(head_idx);
let parent_branch = arena[parent_idx].branch_mut();
parent_branch.children.remove(child_idx.get());
parent_branch.unset_child_bit(child_nibble);
parent_branch.state = ArenaSparseNodeState::Dirty;
parent_branch.remove_child(child_nibble);
// If the branch now has only one child, collapse it. The blinded sibling
// case was already handled above before any mutations.

View File

@@ -1,4 +1,7 @@
use super::{branch_child_idx::BranchChildIter, ArenaSparseSubtrie};
use super::{
branch_child_idx::{BranchChildIdx, BranchChildIter},
ArenaSparseSubtrie,
};
use alloc::{boxed::Box, vec::Vec};
use alloy_primitives::{keccak256, B256};
use alloy_trie::{BranchNodeCompact, TrieMask};
@@ -80,6 +83,29 @@ impl ArenaSparseNodeBranch {
self.state_mask.unset_bit(nibble);
}
/// Inserts a child at `nibble`, updating the state mask, children array, and marking the
/// branch as dirty.
pub(super) fn set_child(&mut self, nibble: u8, child: ArenaSparseNodeBranchChild) {
let insert_pos = BranchChildIdx::insertion_point(self.state_mask, nibble);
self.state_mask.set_bit(nibble);
self.children.insert(insert_pos.get(), child);
self.state = ArenaSparseNodeState::Dirty;
}
/// Removes the child at `nibble`, updating the state mask, children array, and marking the
/// branch as dirty.
///
/// # Panics
///
/// Panics if `nibble` is not set in the state mask.
pub(super) fn remove_child(&mut self, nibble: u8) {
let child_idx =
BranchChildIdx::new(self.state_mask, nibble).expect("nibble not found in state_mask");
self.children.remove(child_idx.get());
self.unset_child_bit(nibble);
self.state = ArenaSparseNodeState::Dirty;
}
/// Iterates over `(nibble, &ArenaSparseNodeBranchChild)` pairs in nibble order.
pub(super) fn child_iter(
&self,