github.com/cryptogateway/go-paymex@v0.0.0-20210204174735-96277fb1e602/trie/stacktrie.go (about) 1 // Copyright 2020 The go-ethereum Authors 2 // This file is part of the go-ethereum library. 3 // 4 // The go-ethereum library is free software: you can redistribute it and/or modify 5 // it under the terms of the GNU Lesser General Public License as published by 6 // the Free Software Foundation, either version 3 of the License, or 7 // (at your option) any later version. 8 // 9 // The go-ethereum library is distributed in the hope that it will be useful, 10 // but WITHOUT ANY WARRANTY; without even the implied warranty of 11 // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 12 // GNU Lesser General Public License for more details. 13 // 14 // You should have received a copy of the GNU Lesser General Public License 15 // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>. 16 17 package trie 18 19 import ( 20 "errors" 21 "fmt" 22 "sync" 23 24 "github.com/cryptogateway/go-paymex/common" 25 "github.com/cryptogateway/go-paymex/ethdb" 26 "github.com/cryptogateway/go-paymex/log" 27 "github.com/cryptogateway/go-paymex/rlp" 28 ) 29 30 var ErrCommitDisabled = errors.New("no database for committing") 31 32 var stPool = sync.Pool{ 33 New: func() interface{} { 34 return NewStackTrie(nil) 35 }, 36 } 37 38 func stackTrieFromPool(db ethdb.KeyValueStore) *StackTrie { 39 st := stPool.Get().(*StackTrie) 40 st.db = db 41 return st 42 } 43 44 func returnToPool(st *StackTrie) { 45 st.Reset() 46 stPool.Put(st) 47 } 48 49 // StackTrie is a trie implementation that expects keys to be inserted 50 // in order. Once it determines that a subtree will no longer be inserted 51 // into, it will hash it and free up the memory it uses. 52 type StackTrie struct { 53 nodeType uint8 // node type (as in branch, ext, leaf) 54 val []byte // value contained by this node if it's a leaf 55 key []byte // key chunk covered by this (full|ext) node 56 keyOffset int // offset of the key chunk inside a full key 57 children [16]*StackTrie // list of children (for fullnodes and exts) 58 59 db ethdb.KeyValueStore // Pointer to the commit db, can be nil 60 } 61 62 // NewStackTrie allocates and initializes an empty trie. 63 func NewStackTrie(db ethdb.KeyValueStore) *StackTrie { 64 return &StackTrie{ 65 nodeType: emptyNode, 66 db: db, 67 } 68 } 69 70 func newLeaf(ko int, key, val []byte, db ethdb.KeyValueStore) *StackTrie { 71 st := stackTrieFromPool(db) 72 st.nodeType = leafNode 73 st.keyOffset = ko 74 st.key = append(st.key, key[ko:]...) 75 st.val = val 76 return st 77 } 78 79 func newExt(ko int, key []byte, child *StackTrie, db ethdb.KeyValueStore) *StackTrie { 80 st := stackTrieFromPool(db) 81 st.nodeType = extNode 82 st.keyOffset = ko 83 st.key = append(st.key, key[ko:]...) 84 st.children[0] = child 85 return st 86 } 87 88 // List all values that StackTrie#nodeType can hold 89 const ( 90 emptyNode = iota 91 branchNode 92 extNode 93 leafNode 94 hashedNode 95 ) 96 97 // TryUpdate inserts a (key, value) pair into the stack trie 98 func (st *StackTrie) TryUpdate(key, value []byte) error { 99 k := keybytesToHex(key) 100 if len(value) == 0 { 101 panic("deletion not supported") 102 } 103 st.insert(k[:len(k)-1], value) 104 return nil 105 } 106 107 func (st *StackTrie) Update(key, value []byte) { 108 if err := st.TryUpdate(key, value); err != nil { 109 log.Error(fmt.Sprintf("Unhandled trie error: %v", err)) 110 } 111 } 112 113 func (st *StackTrie) Reset() { 114 st.db = nil 115 st.key = st.key[:0] 116 st.val = nil 117 for i := range st.children { 118 st.children[i] = nil 119 } 120 st.nodeType = emptyNode 121 st.keyOffset = 0 122 } 123 124 // Helper function that, given a full key, determines the index 125 // at which the chunk pointed by st.keyOffset is different from 126 // the same chunk in the full key. 127 func (st *StackTrie) getDiffIndex(key []byte) int { 128 diffindex := 0 129 for ; diffindex < len(st.key) && st.key[diffindex] == key[st.keyOffset+diffindex]; diffindex++ { 130 } 131 return diffindex 132 } 133 134 // Helper function to that inserts a (key, value) pair into 135 // the trie. 136 func (st *StackTrie) insert(key, value []byte) { 137 switch st.nodeType { 138 case branchNode: /* Branch */ 139 idx := int(key[st.keyOffset]) 140 // Unresolve elder siblings 141 for i := idx - 1; i >= 0; i-- { 142 if st.children[i] != nil { 143 if st.children[i].nodeType != hashedNode { 144 st.children[i].hash() 145 } 146 break 147 } 148 } 149 // Add new child 150 if st.children[idx] == nil { 151 st.children[idx] = stackTrieFromPool(st.db) 152 st.children[idx].keyOffset = st.keyOffset + 1 153 } 154 st.children[idx].insert(key, value) 155 case extNode: /* Ext */ 156 // Compare both key chunks and see where they differ 157 diffidx := st.getDiffIndex(key) 158 159 // Check if chunks are identical. If so, recurse into 160 // the child node. Otherwise, the key has to be split 161 // into 1) an optional common prefix, 2) the fullnode 162 // representing the two differing path, and 3) a leaf 163 // for each of the differentiated subtrees. 164 if diffidx == len(st.key) { 165 // Ext key and key segment are identical, recurse into 166 // the child node. 167 st.children[0].insert(key, value) 168 return 169 } 170 // Save the original part. Depending if the break is 171 // at the extension's last byte or not, create an 172 // intermediate extension or use the extension's child 173 // node directly. 174 var n *StackTrie 175 if diffidx < len(st.key)-1 { 176 n = newExt(diffidx+1, st.key, st.children[0], st.db) 177 } else { 178 // Break on the last byte, no need to insert 179 // an extension node: reuse the current node 180 n = st.children[0] 181 } 182 // Convert to hash 183 n.hash() 184 var p *StackTrie 185 if diffidx == 0 { 186 // the break is on the first byte, so 187 // the current node is converted into 188 // a branch node. 189 st.children[0] = nil 190 p = st 191 st.nodeType = branchNode 192 } else { 193 // the common prefix is at least one byte 194 // long, insert a new intermediate branch 195 // node. 196 st.children[0] = stackTrieFromPool(st.db) 197 st.children[0].nodeType = branchNode 198 st.children[0].keyOffset = st.keyOffset + diffidx 199 p = st.children[0] 200 } 201 // Create a leaf for the inserted part 202 o := newLeaf(st.keyOffset+diffidx+1, key, value, st.db) 203 204 // Insert both child leaves where they belong: 205 origIdx := st.key[diffidx] 206 newIdx := key[diffidx+st.keyOffset] 207 p.children[origIdx] = n 208 p.children[newIdx] = o 209 st.key = st.key[:diffidx] 210 211 case leafNode: /* Leaf */ 212 // Compare both key chunks and see where they differ 213 diffidx := st.getDiffIndex(key) 214 215 // Overwriting a key isn't supported, which means that 216 // the current leaf is expected to be split into 1) an 217 // optional extension for the common prefix of these 2 218 // keys, 2) a fullnode selecting the path on which the 219 // keys differ, and 3) one leaf for the differentiated 220 // component of each key. 221 if diffidx >= len(st.key) { 222 panic("Trying to insert into existing key") 223 } 224 225 // Check if the split occurs at the first nibble of the 226 // chunk. In that case, no prefix extnode is necessary. 227 // Otherwise, create that 228 var p *StackTrie 229 if diffidx == 0 { 230 // Convert current leaf into a branch 231 st.nodeType = branchNode 232 p = st 233 st.children[0] = nil 234 } else { 235 // Convert current node into an ext, 236 // and insert a child branch node. 237 st.nodeType = extNode 238 st.children[0] = NewStackTrie(st.db) 239 st.children[0].nodeType = branchNode 240 st.children[0].keyOffset = st.keyOffset + diffidx 241 p = st.children[0] 242 } 243 244 // Create the two child leaves: the one containing the 245 // original value and the one containing the new value 246 // The child leave will be hashed directly in order to 247 // free up some memory. 248 origIdx := st.key[diffidx] 249 p.children[origIdx] = newLeaf(diffidx+1, st.key, st.val, st.db) 250 p.children[origIdx].hash() 251 252 newIdx := key[diffidx+st.keyOffset] 253 p.children[newIdx] = newLeaf(p.keyOffset+1, key, value, st.db) 254 255 // Finally, cut off the key part that has been passed 256 // over to the children. 257 st.key = st.key[:diffidx] 258 st.val = nil 259 case emptyNode: /* Empty */ 260 st.nodeType = leafNode 261 st.key = key[st.keyOffset:] 262 st.val = value 263 case hashedNode: 264 panic("trying to insert into hash") 265 default: 266 panic("invalid type") 267 } 268 } 269 270 // hash() hashes the node 'st' and converts it into 'hashedNode', if possible. 271 // Possible outcomes: 272 // 1. The rlp-encoded value was >= 32 bytes: 273 // - Then the 32-byte `hash` will be accessible in `st.val`. 274 // - And the 'st.type' will be 'hashedNode' 275 // 2. The rlp-encoded value was < 32 bytes 276 // - Then the <32 byte rlp-encoded value will be accessible in 'st.val'. 277 // - And the 'st.type' will be 'hashedNode' AGAIN 278 // 279 // This method will also: 280 // set 'st.type' to hashedNode 281 // clear 'st.key' 282 func (st *StackTrie) hash() { 283 /* Shortcut if node is already hashed */ 284 if st.nodeType == hashedNode { 285 return 286 } 287 // The 'hasher' is taken from a pool, but we don't actually 288 // claim an instance until all children are done with their hashing, 289 // and we actually need one 290 var h *hasher 291 292 switch st.nodeType { 293 case branchNode: 294 var nodes [17]node 295 for i, child := range st.children { 296 if child == nil { 297 nodes[i] = nilValueNode 298 continue 299 } 300 child.hash() 301 if len(child.val) < 32 { 302 nodes[i] = rawNode(child.val) 303 } else { 304 nodes[i] = hashNode(child.val) 305 } 306 st.children[i] = nil // Reclaim mem from subtree 307 returnToPool(child) 308 } 309 nodes[16] = nilValueNode 310 h = newHasher(false) 311 defer returnHasherToPool(h) 312 h.tmp.Reset() 313 if err := rlp.Encode(&h.tmp, nodes); err != nil { 314 panic(err) 315 } 316 case extNode: 317 st.children[0].hash() 318 h = newHasher(false) 319 defer returnHasherToPool(h) 320 h.tmp.Reset() 321 var valuenode node 322 if len(st.children[0].val) < 32 { 323 valuenode = rawNode(st.children[0].val) 324 } else { 325 valuenode = hashNode(st.children[0].val) 326 } 327 n := struct { 328 Key []byte 329 Val node 330 }{ 331 Key: hexToCompact(st.key), 332 Val: valuenode, 333 } 334 if err := rlp.Encode(&h.tmp, n); err != nil { 335 panic(err) 336 } 337 returnToPool(st.children[0]) 338 st.children[0] = nil // Reclaim mem from subtree 339 case leafNode: 340 h = newHasher(false) 341 defer returnHasherToPool(h) 342 h.tmp.Reset() 343 st.key = append(st.key, byte(16)) 344 sz := hexToCompactInPlace(st.key) 345 n := [][]byte{st.key[:sz], st.val} 346 if err := rlp.Encode(&h.tmp, n); err != nil { 347 panic(err) 348 } 349 case emptyNode: 350 st.val = st.val[:0] 351 st.val = append(st.val, emptyRoot[:]...) 352 st.key = st.key[:0] 353 st.nodeType = hashedNode 354 return 355 default: 356 panic("Invalid node type") 357 } 358 st.key = st.key[:0] 359 st.nodeType = hashedNode 360 if len(h.tmp) < 32 { 361 st.val = st.val[:0] 362 st.val = append(st.val, h.tmp...) 363 return 364 } 365 // Going to write the hash to the 'val'. Need to ensure it's properly sized first 366 // Typically, 'branchNode's will have no 'val', and require this allocation 367 if required := 32 - len(st.val); required > 0 { 368 buf := make([]byte, required) 369 st.val = append(st.val, buf...) 370 } 371 st.val = st.val[:32] 372 h.sha.Reset() 373 h.sha.Write(h.tmp) 374 h.sha.Read(st.val) 375 if st.db != nil { 376 // TODO! Is it safe to Put the slice here? 377 // Do all db implementations copy the value provided? 378 st.db.Put(st.val, h.tmp) 379 } 380 } 381 382 // Hash returns the hash of the current node 383 func (st *StackTrie) Hash() (h common.Hash) { 384 st.hash() 385 if len(st.val) != 32 { 386 // If the node's RLP isn't 32 bytes long, the node will not 387 // be hashed, and instead contain the rlp-encoding of the 388 // node. For the top level node, we need to force the hashing. 389 ret := make([]byte, 32) 390 h := newHasher(false) 391 defer returnHasherToPool(h) 392 h.sha.Reset() 393 h.sha.Write(st.val) 394 h.sha.Read(ret) 395 return common.BytesToHash(ret) 396 } 397 return common.BytesToHash(st.val) 398 } 399 400 // Commit will firstly hash the entrie trie if it's still not hashed 401 // and then commit all nodes to the associated database. Actually most 402 // of the trie nodes MAY have been committed already. The main purpose 403 // here is to commit the root node. 404 // 405 // The associated database is expected, otherwise the whole commit 406 // functionality should be disabled. 407 func (st *StackTrie) Commit() (common.Hash, error) { 408 if st.db == nil { 409 return common.Hash{}, ErrCommitDisabled 410 } 411 st.hash() 412 if len(st.val) != 32 { 413 // If the node's RLP isn't 32 bytes long, the node will not 414 // be hashed (and committed), and instead contain the rlp-encoding of the 415 // node. For the top level node, we need to force the hashing+commit. 416 ret := make([]byte, 32) 417 h := newHasher(false) 418 defer returnHasherToPool(h) 419 h.sha.Reset() 420 h.sha.Write(st.val) 421 h.sha.Read(ret) 422 st.db.Put(ret, st.val) 423 return common.BytesToHash(ret), nil 424 } 425 return common.BytesToHash(st.val), nil 426 }