|
| 1 | +// Copyright (C) 2025, Ava Labs, Inc. All rights reserved. |
| 2 | +// See the file LICENSE.md for licensing terms. |
| 3 | + |
| 4 | +use crate::{ |
| 5 | + Children, HashType, Hashable, IntoSplitPath, PathBuf, PathComponent, SplitPath, TrieEdgeState, |
| 6 | + TrieNode, TriePath, ValueDigest, |
| 7 | +}; |
| 8 | + |
| 9 | +/// An error indicating that a slice of proof nodes is invalid. |
| 10 | +#[derive(Debug, Clone, PartialEq, Eq, Hash, thiserror::Error)] |
| 11 | +pub enum FromKeyProofError { |
| 12 | + /// The parent node's path is not a strict prefix the node that follows it. |
| 13 | + #[error( |
| 14 | + "parent node {parent_path} precedes child node {child_path} but is not a strict prefix of it", |
| 15 | + parent_path = parent_path.display(), |
| 16 | + child_path = child_path.display(), |
| 17 | + )] |
| 18 | + InvalidChildPath { |
| 19 | + /// The path of the parent node. |
| 20 | + parent_path: PathBuf, |
| 21 | + /// The path of the following child node. |
| 22 | + child_path: PathBuf, |
| 23 | + }, |
| 24 | + /// The parent node does not reference the child node at the path component |
| 25 | + /// leading to the child node. |
| 26 | + #[error( |
| 27 | + "child node {child_path} is not reachable from parent node {parent_path}", |
| 28 | + parent_path = parent_path.display(), |
| 29 | + child_path = child_path.display(), |
| 30 | + )] |
| 31 | + MissingChild { |
| 32 | + /// The path of the parent node. |
| 33 | + parent_path: PathBuf, |
| 34 | + /// The path of the following child node. |
| 35 | + child_path: PathBuf, |
| 36 | + }, |
| 37 | +} |
| 38 | + |
| 39 | +/// A root node in a trie formed from a [`ProofNode`]. |
| 40 | +/// |
| 41 | +/// A proof trie follows a linear path from the root to a terminal node, and |
| 42 | +/// includes the necessary information to calculate the hash of each node along |
| 43 | +/// that path. |
| 44 | +/// |
| 45 | +/// In the proof, each node will include the value or value digest at that node, |
| 46 | +/// depending on what is required by the hasher. Additionally, the hashes of each |
| 47 | +/// child node that branches off the node along the path are included. |
| 48 | +#[derive(Debug)] |
| 49 | +pub struct KeyProofTrieRoot<'a, P> { |
| 50 | + partial_path: P, |
| 51 | + value_digest: Option<ValueDigest<&'a [u8]>>, |
| 52 | + children: Children<Option<KeyProofTrieNode<'a, P>>>, |
| 53 | +} |
| 54 | + |
| 55 | +#[derive(Debug)] |
| 56 | +enum KeyProofTrieNode<'a, P> { |
| 57 | + /// Described nodes are proof nodes where we have the data necessary to |
| 58 | + /// reconstruct the hash. The value digest may be a value or a digest. We can |
| 59 | + /// verify the hash of theses nodes using the value or digest, but may not |
| 60 | + /// have the full value. |
| 61 | + Described { |
| 62 | + node: Box<KeyProofTrieRoot<'a, P>>, |
| 63 | + hash: HashType, |
| 64 | + }, |
| 65 | + /// Remote nodes are the nodes where we only know the ID, as discovered |
| 66 | + /// from a proof node. If we only have the child, we can't infer anything |
| 67 | + /// else about the node. |
| 68 | + Remote { hash: HashType }, |
| 69 | +} |
| 70 | + |
| 71 | +impl<'a, P: SplitPath> KeyProofTrieRoot<'a, P> { |
| 72 | + /// Constructs a trie root from a slice of proof nodes. |
| 73 | + /// |
| 74 | + /// Each node in the slice must be a strict prefix of the following node. And, |
| 75 | + /// each child node must be referenced by its parent (i.e., the parent must |
| 76 | + /// indicate a child at the path component leading to the child). The hash |
| 77 | + /// is not verified here. |
| 78 | + /// |
| 79 | + /// # Errors |
| 80 | + /// |
| 81 | + /// - [`FromKeyProofError::InvalidChildPath`] if any node's path is not a strict |
| 82 | + /// prefix of the following node's path. |
| 83 | + /// - [`FromKeyProofError::MissingChild`] if any parent node does not reference |
| 84 | + /// the following child node at the path component leading to the child. |
| 85 | + pub fn new<T, N>(proof: &'a T) -> Result<Option<Box<Self>>, FromKeyProofError> |
| 86 | + where |
| 87 | + T: AsRef<[N]> + ?Sized, |
| 88 | + N: Hashable<FullPath<'a>: IntoSplitPath<Path = P>> + 'a, |
| 89 | + { |
| 90 | + proof |
| 91 | + .as_ref() |
| 92 | + .iter() |
| 93 | + .rev() |
| 94 | + .try_fold(None::<Box<Self>>, |parent, node| match parent { |
| 95 | + None => Ok(Some(Self::new_tail_node(node))), |
| 96 | + Some(p) => p.new_parent_node(node).map(Some), |
| 97 | + }) |
| 98 | + } |
| 99 | + |
| 100 | + /// Creates a new trie root from the tail node of a proof. |
| 101 | + fn new_tail_node<N>(node: &'a N) -> Box<Self> |
| 102 | + where |
| 103 | + N: Hashable<FullPath<'a>: IntoSplitPath<Path = P>>, |
| 104 | + { |
| 105 | + Box::new(Self { |
| 106 | + partial_path: node.full_path().into_split_path(), |
| 107 | + value_digest: node.value_digest(), |
| 108 | + children: node |
| 109 | + .children() |
| 110 | + .map(|_, child| child.map(|hash| KeyProofTrieNode::Remote { hash })), |
| 111 | + }) |
| 112 | + } |
| 113 | + |
| 114 | + /// Creates a new trie root by making this node a child of the given parent. |
| 115 | + /// |
| 116 | + /// The parent key must be a strict prefix of this node's key, and the parent |
| 117 | + /// must reference this node in its children by hash (the hash is not verified |
| 118 | + /// here). |
| 119 | + fn new_parent_node<N>( |
| 120 | + mut self: Box<Self>, |
| 121 | + parent: &'a N, |
| 122 | + ) -> Result<Box<Self>, FromKeyProofError> |
| 123 | + where |
| 124 | + N: Hashable<FullPath<'a>: IntoSplitPath<Path = P>>, |
| 125 | + { |
| 126 | + match parent |
| 127 | + .full_path() |
| 128 | + .into_split_path() |
| 129 | + .longest_common_prefix(self.partial_path) |
| 130 | + .split_first_parts() |
| 131 | + { |
| 132 | + (None, Some((pc, child_path)), parent_path) => { |
| 133 | + let mut parent = Self::new_tail_node(parent); |
| 134 | + if let Some(KeyProofTrieNode::Remote { hash }) = parent.children.take(pc) { |
| 135 | + self.partial_path = child_path; |
| 136 | + parent.partial_path = parent_path; |
| 137 | + parent.children[pc] = Some(KeyProofTrieNode::Described { node: self, hash }); |
| 138 | + Ok(parent) |
| 139 | + } else { |
| 140 | + Err(FromKeyProofError::MissingChild { |
| 141 | + parent_path: parent.partial_path.as_component_slice().into_owned(), |
| 142 | + child_path: self.partial_path.as_component_slice().into_owned(), |
| 143 | + }) |
| 144 | + } |
| 145 | + } |
| 146 | + _ => Err(FromKeyProofError::InvalidChildPath { |
| 147 | + parent_path: parent.full_path().as_component_slice().into_owned(), |
| 148 | + child_path: self.partial_path.as_component_slice().into_owned(), |
| 149 | + }), |
| 150 | + } |
| 151 | + } |
| 152 | +} |
| 153 | + |
| 154 | +impl<'a, P: IntoSplitPath + 'a> KeyProofTrieNode<'a, P> { |
| 155 | + const fn hash(&self) -> &HashType { |
| 156 | + match self { |
| 157 | + KeyProofTrieNode::Described { hash, .. } | KeyProofTrieNode::Remote { hash } => hash, |
| 158 | + } |
| 159 | + } |
| 160 | + |
| 161 | + const fn node(&self) -> Option<&KeyProofTrieRoot<'a, P>> { |
| 162 | + match self { |
| 163 | + KeyProofTrieNode::Described { node, .. } => Some(node), |
| 164 | + KeyProofTrieNode::Remote { .. } => None, |
| 165 | + } |
| 166 | + } |
| 167 | + |
| 168 | + const fn as_edge_state(&self) -> TrieEdgeState<'_, KeyProofTrieRoot<'a, P>> { |
| 169 | + match self { |
| 170 | + KeyProofTrieNode::Described { node, hash } => TrieEdgeState::LocalChild { node, hash }, |
| 171 | + KeyProofTrieNode::Remote { hash } => TrieEdgeState::RemoteChild { hash }, |
| 172 | + } |
| 173 | + } |
| 174 | +} |
| 175 | + |
| 176 | +impl<'a, P: SplitPath + 'a> TrieNode<ValueDigest<&'a [u8]>> for KeyProofTrieRoot<'a, P> { |
| 177 | + type PartialPath<'b> |
| 178 | + = P |
| 179 | + where |
| 180 | + Self: 'b; |
| 181 | + |
| 182 | + fn partial_path(&self) -> Self::PartialPath<'_> { |
| 183 | + self.partial_path |
| 184 | + } |
| 185 | + |
| 186 | + fn value(&self) -> Option<&ValueDigest<&'a [u8]>> { |
| 187 | + self.value_digest.as_ref() |
| 188 | + } |
| 189 | + |
| 190 | + fn child_hash(&self, pc: PathComponent) -> Option<&HashType> { |
| 191 | + self.children[pc].as_ref().map(KeyProofTrieNode::hash) |
| 192 | + } |
| 193 | + |
| 194 | + fn child_node(&self, pc: PathComponent) -> Option<&Self> { |
| 195 | + self.children[pc].as_ref().and_then(KeyProofTrieNode::node) |
| 196 | + } |
| 197 | + |
| 198 | + fn child_state(&self, pc: PathComponent) -> Option<super::TrieEdgeState<'_, Self>> { |
| 199 | + self.children[pc] |
| 200 | + .as_ref() |
| 201 | + .map(KeyProofTrieNode::as_edge_state) |
| 202 | + } |
| 203 | +} |
0 commit comments