1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
use crate::{BranchNodeCompact, Nibbles, StoredSubNode, CHILD_INDEX_RANGE};
use alloy_primitives::B256;

/// Cursor for iterating over a subtrie.
#[derive(Clone)]
pub struct CursorSubNode {
    /// The key of the current node.
    pub key: Nibbles,
    /// The index of the next child to visit.
    nibble: i8,
    /// The node itself.
    pub node: Option<BranchNodeCompact>,
    /// Full key
    full_key: Nibbles,
}

impl Default for CursorSubNode {
    fn default() -> Self {
        Self::new(Nibbles::default(), None)
    }
}

impl std::fmt::Debug for CursorSubNode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CursorSubNode")
            .field("key", &self.key)
            .field("nibble", &self.nibble)
            .field("state_flag", &self.state_flag())
            .field("tree_flag", &self.tree_flag())
            .field("hash_flag", &self.hash_flag())
            .field("hash", &self.hash())
            .finish()
    }
}

/// Implements conversion from `StoredSubNode` to `CursorSubNode`.
impl From<StoredSubNode> for CursorSubNode {
    /// Converts a `StoredSubNode` into a `CursorSubNode`.
    ///
    /// Extracts necessary values from the `StoredSubNode` and constructs
    /// a corresponding `CursorSubNode`.
    fn from(value: StoredSubNode) -> Self {
        let nibble = value.nibble.map_or(-1, |n| n as i8);
        let key = Nibbles::from_nibbles_unchecked(value.key);
        let full_key = full_key(key.clone(), nibble);
        Self { key, nibble, node: value.node, full_key }
    }
}

impl From<CursorSubNode> for StoredSubNode {
    fn from(value: CursorSubNode) -> Self {
        let nibble = if value.nibble >= 0 { Some(value.nibble as u8) } else { None };
        Self { key: value.key.to_vec(), nibble, node: value.node }
    }
}

impl CursorSubNode {
    /// Creates a new `CursorSubNode` from a key and an optional node.
    pub fn new(key: Nibbles, node: Option<BranchNodeCompact>) -> Self {
        // Find the first nibble that is set in the state mask of the node.
        let nibble = node.as_ref().filter(|n| n.root_hash.is_none()).map_or(-1, |n| {
            CHILD_INDEX_RANGE.clone().find(|i| n.state_mask.is_bit_set(*i)).unwrap() as i8
        });
        let full_key = full_key(key.clone(), nibble);
        Self { key, node, nibble, full_key }
    }

    /// Returns the full key of the current node.
    #[inline]
    pub const fn full_key(&self) -> &Nibbles {
        &self.full_key
    }

    /// Returns `true` if the state flag is set for the current nibble.
    #[inline]
    pub fn state_flag(&self) -> bool {
        self.node
            .as_ref()
            .map_or(true, |node| self.nibble < 0 || node.state_mask.is_bit_set(self.nibble as u8))
    }

    /// Returns `true` if the tree flag is set for the current nibble.
    #[inline]
    pub fn tree_flag(&self) -> bool {
        self.node
            .as_ref()
            .map_or(true, |node| self.nibble < 0 || node.tree_mask.is_bit_set(self.nibble as u8))
    }

    /// Returns `true` if the current nibble has a root hash.
    pub fn hash_flag(&self) -> bool {
        self.node.as_ref().map_or(false, |node| match self.nibble {
            // This guy has it
            -1 => node.root_hash.is_some(),
            // Or get it from the children
            _ => node.hash_mask.is_bit_set(self.nibble as u8),
        })
    }

    /// Returns the root hash of the current node, if it has one.
    pub fn hash(&self) -> Option<B256> {
        if self.hash_flag() {
            let node = self.node.as_ref().unwrap();
            match self.nibble {
                -1 => node.root_hash,
                _ => Some(node.hash_for_nibble(self.nibble as u8)),
            }
        } else {
            None
        }
    }

    /// Returns the next child index to visit.
    #[inline]
    pub const fn nibble(&self) -> i8 {
        self.nibble
    }

    /// Increments the nibble index.
    #[inline]
    pub fn inc_nibble(&mut self) {
        self.nibble += 1;
        update_full_key(&mut self.full_key, self.nibble - 1, self.nibble);
    }

    /// Sets the nibble index.
    #[inline]
    pub fn set_nibble(&mut self, nibble: i8) {
        let old_nibble = self.nibble;
        self.nibble = nibble;
        update_full_key(&mut self.full_key, old_nibble, self.nibble);
    }
}

/// Constructs a full key from the given `Nibbles` and `nibble`.
#[inline]
fn full_key(mut key: Nibbles, nibble: i8) -> Nibbles {
    if nibble >= 0 {
        key.push(nibble as u8);
    }
    key
}

/// Updates the key by replacing or appending a nibble based on the old and new nibble values.
#[inline]
fn update_full_key(key: &mut Nibbles, old_nibble: i8, new_nibble: i8) {
    if new_nibble >= 0 {
        if old_nibble >= 0 {
            let last_index = key.len() - 1;
            key.set_at(last_index, new_nibble as u8);
        } else {
            key.push(new_nibble as u8);
        }
    } else if old_nibble >= 0 {
        key.pop();
    }
}