go-merkletree-iden3/merkletree.go
merkletree.goRaw
1package merkletree
2
3import (
4 "bytes"
5 "encoding/hex"
6 "errors"
7 "fmt"
8 "io"
9 "math/big"
10 "sync"
11
12 "github.com/iden3/go-iden3-core/common"
13 "github.com/iden3/go-iden3-core/db"
14 cryptoUtils "github.com/iden3/go-iden3-crypto/utils"
15)
16
17const (
18 // proofFlagsLen is the byte length of the flags in the proof header (first 32
19 // bytes).
20 proofFlagsLen = 2
21 // ElemBytesLen is the length of the Hash byte array
22 ElemBytesLen = 32
23)
24
25var (
26 // ErrNodeKeyAlreadyExists is used when a node key already exists.
27 ErrNodeKeyAlreadyExists = errors.New("node already exists")
28 // ErrEntryIndexNotFound is used when no entry is found for an index.
29 ErrEntryIndexNotFound = errors.New("node index not found in the DB")
30 // ErrNodeBytesBadSize is used when the data of a node has an incorrect
31 // size and can't be parsed.
32 ErrNodeBytesBadSize = errors.New("node data has incorrect size in the DB")
33 // ErrReachedMaxLevel is used when a traversal of the MT reaches the
34 // maximum level.
35 ErrReachedMaxLevel = errors.New("reached maximum level of the merkle tree")
36 // ErrInvalidNodeFound is used when an invalid node is found and can't
37 // be parsed.
38 ErrInvalidNodeFound = errors.New("found an invalid node in the DB")
39 // ErrInvalidProofBytes is used when a serialized proof is invalid.
40 ErrInvalidProofBytes = errors.New("the serialized proof is invalid")
41 // ErrInvalidDBValue is used when a value in the key value DB is
42 // invalid (for example, it doen't contain a byte header and a []byte
43 // body of at least len=1.
44 ErrInvalidDBValue = errors.New("the value in the DB is invalid")
45 // ErrEntryIndexAlreadyExists is used when the entry index already
46 // exists in the tree.
47 ErrEntryIndexAlreadyExists = errors.New("the entry index already exists in the tree")
48 // ErrNotWritable is used when the MerkleTree is not writable and a write function is called
49 ErrNotWritable = errors.New("Merkle Tree not writable")
50 // ErrKeyNotFound is used when a key is not found in the MerkleTree.
51 ErrKeyNotFound = errors.New("Key not found in the tree")
52 rootNodeValue = []byte("currentroot")
53 // HashZero is used at Empty nodes
54 HashZero = Hash{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}
55)
56
57// Hash is the generic type stored in the MerkleTree
58type Hash [32]byte
59
60// String returns decimal representation in string format of the Hash
61func (h Hash) String() string {
62 s := h.BigInt().String()
63 if len(s) < 8 {
64 return s
65 }
66 return s[0:8] + "..."
67}
68
69// Hex returns the hexadecimal representation of the Hash
70func (h Hash) Hex() string {
71 return hex.EncodeToString(h.BigInt().Bytes())
72}
73
74// BigInt returns the *big.Int representation of the *Hash
75func (h *Hash) BigInt() *big.Int {
76 return new(big.Int).SetBytes(common.SwapEndianness(h[:]))
77}
78
79// NewHashFromBigInt returns a *Hash representation of the given *big.Int
80func NewHashFromBigInt(b *big.Int) *Hash {
81 r := &Hash{}
82 copy(r[:], common.SwapEndianness(b.Bytes()))
83 return r
84}
85
86// MerkleTree is the struct with the main elements of the MerkleTree
87type MerkleTree struct {
88 sync.RWMutex
89 db db.Storage
90 rootKey *Hash
91 writable bool
92 maxLevels int
93}
94
95// NewMerkleTree loads a new Merkletree. If in the sotrage already exists one will open that one, if not, will create a new one.
96func NewMerkleTree(storage db.Storage, maxLevels int) (*MerkleTree, error) {
97 mt := MerkleTree{db: storage, maxLevels: maxLevels, writable: true}
98
99 v, err := mt.db.Get(rootNodeValue)
100 if err != nil {
101 tx, err := mt.db.NewTx()
102 if err != nil {
103 return nil, err
104 }
105 mt.rootKey = &HashZero
106 tx.Put(rootNodeValue, mt.rootKey[:])
107 err = tx.Commit()
108 if err != nil {
109 return nil, err
110 }
111 return &mt, nil
112 }
113 mt.rootKey = &Hash{}
114 copy(mt.rootKey[:], v)
115 return &mt, nil
116}
117
118// DB returns the MerkleTree.DB()
119func (mt *MerkleTree) DB() db.Storage {
120 return mt.db
121}
122
123// Root returns the MerkleRoot
124func (mt *MerkleTree) Root() *Hash {
125 return mt.rootKey
126}
127
128// Add adds a Key & Value into the MerkleTree. Where the `k` determines the path from the Root to the Leaf.
129func (mt *MerkleTree) Add(k, v *big.Int) error {
130 // verify that the MerkleTree is writable
131 if !mt.writable {
132 return ErrNotWritable
133 }
134
135 // verfy that the ElemBytes are valid and fit inside the Finite Field.
136 if !cryptoUtils.CheckBigIntInField(k) {
137 return errors.New("Key not inside the Finite Field")
138 }
139 if !cryptoUtils.CheckBigIntInField(v) {
140 return errors.New("Value not inside the Finite Field")
141 }
142
143 tx, err := mt.db.NewTx()
144 if err != nil {
145 return err
146 }
147 mt.Lock()
148 defer mt.Unlock()
149
150 kHash := NewHashFromBigInt(k)
151 vHash := NewHashFromBigInt(v)
152 newNodeLeaf := NewNodeLeaf(kHash, vHash)
153 path := getPath(mt.maxLevels, kHash[:])
154
155 newRootKey, err := mt.addLeaf(tx, newNodeLeaf, mt.rootKey, 0, path)
156 if err != nil {
157 return err
158 }
159 mt.rootKey = newRootKey
160 mt.dbInsert(tx, rootNodeValue, DBEntryTypeRoot, mt.rootKey[:])
161
162 if err := tx.Commit(); err != nil {
163 return err
164 }
165
166 return nil
167}
168
169// pushLeaf recursively pushes an existing oldLeaf down until its path diverges
170// from newLeaf, at which point both leafs are stored, all while updating the
171// path.
172func (mt *MerkleTree) pushLeaf(tx db.Tx, newLeaf *Node, oldLeaf *Node,
173 lvl int, pathNewLeaf []bool, pathOldLeaf []bool) (*Hash, error) {
174 if lvl > mt.maxLevels-2 {
175 return nil, ErrReachedMaxLevel
176 }
177 var newNodeMiddle *Node
178 if pathNewLeaf[lvl] == pathOldLeaf[lvl] { // We need to go deeper!
179 nextKey, err := mt.pushLeaf(tx, newLeaf, oldLeaf, lvl+1, pathNewLeaf, pathOldLeaf)
180 if err != nil {
181 return nil, err
182 }
183 if pathNewLeaf[lvl] {
184 newNodeMiddle = NewNodeMiddle(&HashZero, nextKey) // go right
185 } else {
186 newNodeMiddle = NewNodeMiddle(nextKey, &HashZero) // go left
187 }
188 return mt.addNode(tx, newNodeMiddle)
189 } else {
190 oldLeafKey, err := oldLeaf.Key()
191 if err != nil {
192 return nil, err
193 }
194 newLeafKey, err := newLeaf.Key()
195 if err != nil {
196 return nil, err
197 }
198
199 if pathNewLeaf[lvl] {
200 newNodeMiddle = NewNodeMiddle(oldLeafKey, newLeafKey)
201 } else {
202 newNodeMiddle = NewNodeMiddle(newLeafKey, oldLeafKey)
203 }
204 // We can add newLeaf now. We don't need to add oldLeaf because it's already in the tree.
205 _, err = mt.addNode(tx, newLeaf)
206 if err != nil {
207 return nil, err
208 }
209 return mt.addNode(tx, newNodeMiddle)
210 }
211}
212
213// addLeaf recursively adds a newLeaf in the MT while updating the path.
214func (mt *MerkleTree) addLeaf(tx db.Tx, newLeaf *Node, key *Hash,
215 lvl int, path []bool) (*Hash, error) {
216 var err error
217 var nextKey *Hash
218 if lvl > mt.maxLevels-1 {
219 return nil, ErrReachedMaxLevel
220 }
221 n, err := mt.GetNode(key)
222 if err != nil {
223 return nil, err
224 }
225 switch n.Type {
226 case NodeTypeEmpty:
227 // We can add newLeaf now
228 return mt.addNode(tx, newLeaf)
229 case NodeTypeLeaf:
230 nKey := n.Entry[0]
231 // Check if leaf node found contains the leaf node we are trying to add
232 newLeafKey := newLeaf.Entry[0]
233 if bytes.Equal(nKey[:], newLeafKey[:]) {
234 return nil, ErrEntryIndexAlreadyExists
235 }
236 pathOldLeaf := getPath(mt.maxLevels, nKey[:])
237 // We need to push newLeaf down until its path diverges from n's path
238 return mt.pushLeaf(tx, newLeaf, n, lvl, path, pathOldLeaf)
239 case NodeTypeMiddle:
240 // We need to go deeper, continue traversing the tree, left or right depending on path
241 var newNodeMiddle *Node
242 if path[lvl] {
243 nextKey, err = mt.addLeaf(tx, newLeaf, n.ChildR, lvl+1, path) // go right
244 newNodeMiddle = NewNodeMiddle(n.ChildL, nextKey)
245 } else {
246 nextKey, err = mt.addLeaf(tx, newLeaf, n.ChildL, lvl+1, path) // go left
247 newNodeMiddle = NewNodeMiddle(nextKey, n.ChildR)
248 }
249 if err != nil {
250 return nil, err
251 }
252 // Update the node to reflect the modified child
253 return mt.addNode(tx, newNodeMiddle)
254 default:
255 return nil, ErrInvalidNodeFound
256 }
257}
258
259// addNode adds a node into the MT. Empty nodes are not stored in the tree;
260// they are all the same and assumed to always exist.
261func (mt *MerkleTree) addNode(tx db.Tx, n *Node) (*Hash, error) {
262 // verify that the MerkleTree is writable
263 if !mt.writable {
264 return nil, ErrNotWritable
265 }
266 if n.Type == NodeTypeEmpty {
267 return n.Key()
268 }
269 k, err := n.Key()
270 if err != nil {
271 return nil, err
272 }
273 v := n.Value()
274 // Check that the node key doesn't already exist
275 if _, err := tx.Get(k[:]); err == nil {
276 return nil, ErrNodeKeyAlreadyExists
277 }
278 tx.Put(k[:], v)
279 return k, nil
280}
281
282// Delete removes the specified Key from the MerkleTree, and updates the pad from the delted key to the Root with the new values.
283// This method removes the key from the MerkleTree, but does not remove the old nodes from the key-value database; this means that if the tree is accessed by an old Root where the key was not deleted yet, the key will still exist. If is desired to remove the key-values from the database that are not under the current Root, an option could be to dump all the claims and import them in a new MerkleTree in a new database, but this will loose all the Root history of the MerkleTree
284func (mt *MerkleTree) Delete(k *big.Int) error {
285 // verify that the MerkleTree is writable
286 if !mt.writable {
287 return ErrNotWritable
288 }
289
290 // verfy that the ElemBytes are valid and fit inside the Finite Field.
291 if !cryptoUtils.CheckBigIntInField(k) {
292 return errors.New("Key not inside the Finite Field")
293 }
294 tx, err := mt.db.NewTx()
295 if err != nil {
296 return err
297 }
298 mt.Lock()
299 defer mt.Unlock()
300
301 kHash := NewHashFromBigInt(k)
302 path := getPath(mt.maxLevels, kHash[:])
303
304 nextKey := mt.rootKey
305 var siblings []*Hash
306 for i := 0; i < mt.maxLevels; i++ {
307 n, err := mt.GetNode(nextKey)
308 if err != nil {
309 return err
310 }
311 switch n.Type {
312 case NodeTypeEmpty:
313 return nil
314 case NodeTypeLeaf:
315 if bytes.Equal(kHash[:], n.Entry[0][:]) {
316 // remove and go up with the sibling
317 err = mt.rmAndUpload(tx, path, kHash, siblings)
318 return err
319 } else {
320 return ErrKeyNotFound
321 }
322 case NodeTypeMiddle:
323 if path[i] {
324 nextKey = n.ChildR
325 siblings = append(siblings, n.ChildL)
326 } else {
327 nextKey = n.ChildL
328 siblings = append(siblings, n.ChildR)
329 }
330 default:
331 return ErrInvalidNodeFound
332 }
333 }
334
335 return nil
336}
337
338// rmAndUpload removes the key, and goes up until the root updating all the nodes with the new values.
339func (mt *MerkleTree) rmAndUpload(tx db.Tx, path []bool, kHash *Hash, siblings []*Hash) error {
340 toUpload := siblings[len(siblings)-1]
341 if len(siblings) < 2 {
342 mt.rootKey = siblings[0]
343 mt.dbInsert(tx, rootNodeValue, DBEntryTypeRoot, mt.rootKey[:])
344 return tx.Commit()
345 }
346 for i := len(siblings) - 2; i >= 0; i-- {
347 if !bytes.Equal(siblings[i][:], HashZero[:]) {
348 var newNode *Node
349 if path[i] {
350 newNode = NewNodeMiddle(siblings[i], toUpload)
351 } else {
352 newNode = NewNodeMiddle(toUpload, siblings[i])
353 }
354 _, err := mt.addNode(tx, newNode)
355 if err != ErrNodeKeyAlreadyExists && err != nil {
356 return err
357 }
358 // go up until the root
359 newRootKey, err := mt.recalculatePathUntilRoot(tx, path, newNode, siblings[:i])
360 if err != nil {
361 return err
362 }
363 mt.rootKey = newRootKey
364 mt.dbInsert(tx, rootNodeValue, DBEntryTypeRoot, mt.rootKey[:])
365 break
366 }
367 // if i==0 (root position), stop and store the sibling of the deleted leaf as root
368 if i == 0 {
369 mt.rootKey = toUpload
370 mt.dbInsert(tx, rootNodeValue, DBEntryTypeRoot, mt.rootKey[:])
371 break
372 }
373 }
374 if err := tx.Commit(); err != nil {
375 return err
376 }
377
378 return nil
379}
380
381// recalculatePathUntilRoot recalculates the nodes until the Root
382func (mt *MerkleTree) recalculatePathUntilRoot(tx db.Tx, path []bool, node *Node, siblings []*Hash) (*Hash, error) {
383 for i := len(siblings) - 1; i >= 0; i-- {
384 nodeKey, err := node.Key()
385 if err != nil {
386 return nil, err
387 }
388 if path[i] {
389 node = NewNodeMiddle(siblings[i], nodeKey)
390 } else {
391 node = NewNodeMiddle(nodeKey, siblings[i])
392 }
393 _, err = mt.addNode(tx, node)
394 if err != ErrNodeKeyAlreadyExists && err != nil {
395 return nil, err
396 }
397 }
398
399 // return last node added, which is the root
400 nodeKey, err := node.Key()
401 return nodeKey, err
402}
403
404// dbGet is a helper function to get the node of a key from the internal
405// storage.
406func (mt *MerkleTree) dbGet(k []byte) (NodeType, []byte, error) {
407 if bytes.Equal(k, HashZero[:]) {
408 return 0, nil, nil
409 }
410
411 value, err := mt.db.Get(k)
412 if err != nil {
413 return 0, nil, err
414 }
415
416 if len(value) < 2 {
417 return 0, nil, ErrInvalidDBValue
418 }
419 nodeType := value[0]
420 nodeBytes := value[1:]
421
422 return NodeType(nodeType), nodeBytes, nil
423}
424
425// dbInsert is a helper function to insert a node into a key in an open db
426// transaction.
427func (mt *MerkleTree) dbInsert(tx db.Tx, k []byte, t NodeType, data []byte) {
428 v := append([]byte{byte(t)}, data...)
429 tx.Put(k, v)
430}
431
432// GetNode gets a node by key from the MT. Empty nodes are not stored in the
433// tree; they are all the same and assumed to always exist.
434func (mt *MerkleTree) GetNode(key *Hash) (*Node, error) {
435 if bytes.Equal(key[:], HashZero[:]) {
436 return NewNodeEmpty(), nil
437 }
438 nBytes, err := mt.db.Get(key[:])
439 if err != nil {
440 return nil, err
441 }
442 return NewNodeFromBytes(nBytes)
443}
444
445// getPath returns the binary path, from the root to the leaf.
446func getPath(numLevels int, k []byte) []bool {
447 path := make([]bool, numLevels)
448 for n := 0; n < numLevels; n++ {
449 path[n] = common.TestBit(k[:], uint(n))
450 }
451 return path
452}
453
454// NodeAux contains the auxiliary node used in a non-existence proof.
455type NodeAux struct {
456 Key *Hash
457 Value *Hash
458}
459
460// Proof defines the required elements for a MT proof of existence or non-existence.
461type Proof struct {
462 // existence indicates wether this is a proof of existence or non-existence.
463 Existence bool
464 // depth indicates how deep in the tree the proof goes.
465 depth uint
466 // notempties is a bitmap of non-empty Siblings found in Siblings.
467 notempties [ElemBytesLen - proofFlagsLen]byte
468 // Siblings is a list of non-empty sibling keys.
469 Siblings []*Hash
470 NodeAux *NodeAux
471}
472
473// NewProofFromBytes parses a byte array into a Proof.
474func NewProofFromBytes(bs []byte) (*Proof, error) {
475 if len(bs) < ElemBytesLen {
476 return nil, ErrInvalidProofBytes
477 }
478 p := &Proof{}
479 if (bs[0] & 0x01) == 0 {
480 p.Existence = true
481 }
482 p.depth = uint(bs[1])
483 copy(p.notempties[:], bs[proofFlagsLen:ElemBytesLen])
484 siblingBytes := bs[ElemBytesLen:]
485 sibIdx := 0
486 for i := uint(0); i < p.depth; i++ {
487 if common.TestBitBigEndian(p.notempties[:], i) {
488 if len(siblingBytes) < (sibIdx+1)*ElemBytesLen {
489 return nil, ErrInvalidProofBytes
490 }
491 var sib Hash
492 copy(sib[:], siblingBytes[sibIdx*ElemBytesLen:(sibIdx+1)*ElemBytesLen])
493 p.Siblings = append(p.Siblings, &sib)
494 sibIdx++
495 }
496 }
497
498 if !p.Existence && ((bs[0] & 0x02) != 0) {
499 p.NodeAux = &NodeAux{Key: &Hash{}, Value: &Hash{}}
500 nodeAuxBytes := siblingBytes[len(p.Siblings)*ElemBytesLen:]
501 if len(nodeAuxBytes) != 2*ElemBytesLen {
502 return nil, ErrInvalidProofBytes
503 }
504 copy(p.NodeAux.Key[:], nodeAuxBytes[:ElemBytesLen])
505 copy(p.NodeAux.Value[:], nodeAuxBytes[ElemBytesLen:2*ElemBytesLen])
506 }
507 return p, nil
508}
509
510// Bytes serializes a Proof into a byte array.
511func (p *Proof) Bytes() []byte {
512 bsLen := proofFlagsLen + len(p.notempties) + ElemBytesLen*len(p.Siblings)
513 if p.NodeAux != nil {
514 bsLen += 2 * ElemBytesLen
515 }
516 bs := make([]byte, bsLen)
517
518 if !p.Existence {
519 bs[0] |= 0x01
520 }
521 bs[1] = byte(p.depth)
522 copy(bs[proofFlagsLen:len(p.notempties)+proofFlagsLen], p.notempties[:])
523 siblingsBytes := bs[len(p.notempties)+proofFlagsLen:]
524 for i, k := range p.Siblings {
525 copy(siblingsBytes[i*ElemBytesLen:(i+1)*ElemBytesLen], k[:])
526 }
527 if p.NodeAux != nil {
528 bs[0] |= 0x02
529 copy(bs[len(bs)-2*ElemBytesLen:], p.NodeAux.Key[:])
530 copy(bs[len(bs)-1*ElemBytesLen:], p.NodeAux.Value[:])
531 }
532 return bs
533}
534
535// SiblingsFromProof returns all the siblings of the proof. This function is used to generate the siblings input for the circom circuits.
536func SiblingsFromProof(proof *Proof) []*Hash {
537 sibIdx := 0
538 var siblings []*Hash
539 for lvl := 0; lvl < int(proof.depth); lvl++ {
540 if common.TestBitBigEndian(proof.notempties[:], uint(lvl)) {
541 siblings = append(siblings, proof.Siblings[sibIdx])
542 sibIdx++
543 } else {
544 siblings = append(siblings, &HashZero)
545 }
546 }
547 return siblings
548}
549
550func (p *Proof) AllSiblings() []*Hash {
551 return SiblingsFromProof(p)
552}
553
554func (p *Proof) AllSiblingsCircom(levels int) []*big.Int {
555 siblings := p.AllSiblings()
556 // Add the rest of empty levels to the siblings
557 for i := len(siblings); i < levels; i++ {
558 siblings = append(siblings, &HashZero)
559 }
560 siblings = append(siblings, &HashZero) // add extra level for circom compatibility
561 siblingsBigInt := make([]*big.Int, len(siblings))
562 for i, sibling := range siblings {
563 siblingsBigInt[i] = sibling.BigInt()
564 }
565 return siblingsBigInt
566}
567
568// GenerateProof generates the proof of existence (or non-existence) of an
569// Entry's hash Index for a Merkle Tree given the root.
570// If the rootKey is nil, the current merkletree root is used
571func (mt *MerkleTree) GenerateProof(k *big.Int, rootKey *Hash) (*Proof, error) {
572 p := &Proof{}
573 var siblingKey *Hash
574
575 kHash := NewHashFromBigInt(k)
576 path := getPath(mt.maxLevels, kHash[:])
577 if rootKey == nil {
578 rootKey = mt.Root()
579 }
580 nextKey := rootKey
581 for p.depth = 0; p.depth < uint(mt.maxLevels); p.depth++ {
582 n, err := mt.GetNode(nextKey)
583 if err != nil {
584 return nil, err
585 }
586 switch n.Type {
587 case NodeTypeEmpty:
588 return p, nil
589 case NodeTypeLeaf:
590 if bytes.Equal(kHash[:], n.Entry[0][:]) {
591 p.Existence = true
592 return p, nil
593 } else {
594 // We found a leaf whose entry didn't match hIndex
595 p.NodeAux = &NodeAux{Key: n.Entry[0], Value: n.Entry[1]}
596 return p, nil
597 }
598 case NodeTypeMiddle:
599 if path[p.depth] {
600 nextKey = n.ChildR
601 siblingKey = n.ChildL
602 } else {
603 nextKey = n.ChildL
604 siblingKey = n.ChildR
605 }
606 default:
607 return nil, ErrInvalidNodeFound
608 }
609 if !bytes.Equal(siblingKey[:], HashZero[:]) {
610 common.SetBitBigEndian(p.notempties[:], uint(p.depth))
611 p.Siblings = append(p.Siblings, siblingKey)
612 }
613 }
614 return nil, ErrEntryIndexNotFound
615}
616
617// VerifyProof verifies the Merkle Proof for the entry and root.
618func VerifyProof(rootKey *Hash, proof *Proof, k, v *big.Int) bool {
619 rootFromProof, err := RootFromProof(proof, k, v)
620 if err != nil {
621 return false
622 }
623 return bytes.Equal(rootKey[:], rootFromProof[:])
624}
625
626// RootFromProof calculates the root that would correspond to a tree whose
627// siblings are the ones in the proof with the claim hashing to hIndex and
628// hValue.
629func RootFromProof(proof *Proof, k, v *big.Int) (*Hash, error) {
630 kHash := NewHashFromBigInt(k)
631 vHash := NewHashFromBigInt(v)
632 sibIdx := len(proof.Siblings) - 1
633 var err error
634 var midKey *Hash
635 if proof.Existence {
636 midKey, err = LeafKey(kHash, vHash)
637 if err != nil {
638 return nil, err
639 }
640 } else {
641 if proof.NodeAux == nil {
642 midKey = &HashZero
643 } else {
644 if bytes.Equal(kHash[:], proof.NodeAux.Key[:]) {
645 return nil, fmt.Errorf("Non-existence proof being checked against hIndex equal to nodeAux")
646 }
647 midKey, err = LeafKey(proof.NodeAux.Key, proof.NodeAux.Value)
648 if err != nil {
649 return nil, err
650 }
651 }
652 }
653 path := getPath(int(proof.depth), kHash[:])
654 var siblingKey *Hash
655 for lvl := int(proof.depth) - 1; lvl >= 0; lvl-- {
656 if common.TestBitBigEndian(proof.notempties[:], uint(lvl)) {
657 siblingKey = proof.Siblings[sibIdx]
658 sibIdx--
659 } else {
660 siblingKey = &HashZero
661 }
662 if path[lvl] {
663 midKey, err = NewNodeMiddle(siblingKey, midKey).Key()
664 if err != nil {
665 return nil, err
666 }
667 } else {
668 midKey, err = NewNodeMiddle(midKey, siblingKey).Key()
669 if err != nil {
670 return nil, err
671 }
672 }
673 }
674 return midKey, nil
675}
676
677// walk is a helper recursive function to iterate over all tree branches
678func (mt *MerkleTree) walk(key *Hash, f func(*Node)) error {
679 n, err := mt.GetNode(key)
680 if err != nil {
681 return err
682 }
683 switch n.Type {
684 case NodeTypeEmpty:
685 f(n)
686 case NodeTypeLeaf:
687 f(n)
688 case NodeTypeMiddle:
689 f(n)
690 if err := mt.walk(n.ChildL, f); err != nil {
691 return err
692 }
693 if err := mt.walk(n.ChildR, f); err != nil {
694 return err
695 }
696 default:
697 return ErrInvalidNodeFound
698 }
699 return nil
700}
701
702// Walk iterates over all the branches of a MerkleTree with the given rootKey
703// if rootKey is nil, it will get the current RootKey of the current state of the MerkleTree.
704// For each node, it calls the f function given in the parameters.
705// See some examples of the Walk function usage in the merkletree_test.go
706// test functions: TestMTWalk, TestMTWalkGraphViz, TestMTWalkDumpClaims
707func (mt *MerkleTree) Walk(rootKey *Hash, f func(*Node)) error {
708 if rootKey == nil {
709 rootKey = mt.Root()
710 }
711 err := mt.walk(rootKey, f)
712 return err
713}
714
715// GraphViz uses Walk function to generate a string GraphViz representation of the
716// tree and writes it to w
717func (mt *MerkleTree) GraphViz(w io.Writer, rootKey *Hash) error {
718 fmt.Fprintf(w, `digraph hierarchy {
719node [fontname=Monospace,fontsize=10,shape=box]
720`)
721 cnt := 0
722 var errIn error
723 err := mt.Walk(rootKey, func(n *Node) {
724 k, err := n.Key()
725 if err != nil {
726 errIn = err
727 }
728 switch n.Type {
729 case NodeTypeEmpty:
730 case NodeTypeLeaf:
731 fmt.Fprintf(w, "\"%v\" [style=filled];\n", k.String())
732 case NodeTypeMiddle:
733 lr := [2]string{n.ChildL.String(), n.ChildR.String()}
734 emptyNodes := ""
735 for i := range lr {
736 if lr[i] == "0" {
737 lr[i] = fmt.Sprintf("empty%v", cnt)
738 emptyNodes += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n", lr[i])
739 cnt++
740 }
741 }
742 fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", k.String(), lr[0], lr[1])
743 fmt.Fprint(w, emptyNodes)
744 default:
745 }
746 })
747 fmt.Fprintf(w, "}\n")
748 if errIn != nil {
749 return errIn
750 }
751 return err
752}
753
754// PrintGraphViz prints directly the GraphViz() output
755func (mt *MerkleTree) PrintGraphViz(rootKey *Hash) error {
756 if rootKey == nil {
757 rootKey = mt.Root()
758 }
759 w := bytes.NewBufferString("")
760 fmt.Fprintf(w, "--------\nGraphViz of the MerkleTree with RootKey "+rootKey.BigInt().String()+"\n")
761 err := mt.GraphViz(w, nil)
762 if err != nil {
763 return err
764 }
765 fmt.Fprintf(w, "End of GraphViz of the MerkleTree with RootKey "+rootKey.BigInt().String()+"\n--------\n")
766
767 fmt.Println(w)
768 return nil
769}