github.com/0chain/gosdk@v1.17.11/core/util/merkle_tree.go (about)

     1  package util
     2  
     3  import (
     4  	"fmt"
     5  )
     6  
     7  /*MerkleTree - A data structure that implements MerkleTreeI interface */
     8  type MerkleTree struct {
     9  	tree        []string
    10  	leavesCount int
    11  	levels      int
    12  }
    13  
    14  func VerifyMerklePath(hash string, path *MTPath, root string) bool {
    15  	mthash := hash
    16  	pathNodes := path.Nodes
    17  	pl := len(pathNodes)
    18  	idx := path.LeafIndex
    19  	for i := 0; i < pl; i++ {
    20  		if idx&1 == 1 {
    21  			mthash = MHash(pathNodes[i], mthash)
    22  		} else {
    23  			mthash = MHash(mthash, pathNodes[i])
    24  		}
    25  		idx = (idx - idx&1) / 2
    26  	}
    27  	return mthash == root
    28  }
    29  
    30  func (mt *MerkleTree) computeSize(leaves int) (int, int) {
    31  	if leaves == 1 {
    32  		return 2, 2
    33  	}
    34  	var tsize int
    35  	var levels int
    36  	for ll := leaves; ll > 1; ll = (ll + 1) / 2 {
    37  		tsize += ll
    38  		levels++
    39  	}
    40  	tsize++
    41  	levels++
    42  	return tsize, levels
    43  }
    44  
    45  /*ComputeTree - given the leaf nodes, compute the merkle tree */
    46  func (mt *MerkleTree) ComputeTree(hashes []Hashable) {
    47  	var tsize int
    48  	tsize, mt.levels = mt.computeSize(len(hashes))
    49  	mt.leavesCount = len(hashes)
    50  	mt.tree = make([]string, tsize)
    51  	for idx, hashable := range hashes {
    52  		mt.tree[idx] = hashable.GetHash()
    53  	}
    54  	if len(hashes) == 1 {
    55  		mt.tree[1] = DecodeAndMHash(mt.tree[0], mt.tree[0])
    56  		return
    57  	}
    58  	for pl0, plsize := 0, mt.leavesCount; plsize > 1; pl0, plsize = pl0+plsize, (plsize+1)/2 {
    59  		l0 := pl0 + plsize
    60  		for i, j := 0, 0; i < plsize; i, j = i+2, j+1 {
    61  			mt.tree[pl0+plsize+j] = DecodeAndMHash(mt.tree[pl0+i], mt.tree[pl0+i+1])
    62  		}
    63  		if plsize&1 == 1 {
    64  			mt.tree[l0+plsize/2] = DecodeAndMHash(mt.tree[pl0+plsize-1], mt.tree[pl0+plsize-1])
    65  		}
    66  	}
    67  }
    68  
    69  /*GetRoot - get the root of the merkle tree */
    70  func (mt *MerkleTree) GetRoot() string {
    71  	return mt.tree[len(mt.tree)-1]
    72  }
    73  
    74  /*GetTree - get the entire merkle tree */
    75  func (mt *MerkleTree) GetTree() []string {
    76  	return mt.tree
    77  }
    78  
    79  /*SetTree - set the entire merkle tree */
    80  func (mt *MerkleTree) SetTree(leavesCount int, tree []string) error {
    81  	size, levels := mt.computeSize(leavesCount)
    82  	if size != len(tree) {
    83  		return fmt.Errorf("Merkle tree with leaves %v should have size %v but only %v is given", leavesCount, size, len(tree))
    84  	}
    85  	mt.levels = levels
    86  	mt.tree = tree
    87  	mt.leavesCount = leavesCount
    88  	return nil
    89  }
    90  
    91  /*GetLeafIndex - Get the index of the leaf node in the tree */
    92  func (mt *MerkleTree) GetLeafIndex(hash Hashable) int {
    93  	hs := hash.GetHash()
    94  	for i := 0; i < mt.leavesCount; i++ {
    95  		if mt.tree[i] == hs {
    96  			return i
    97  		}
    98  	}
    99  	return -1
   100  }
   101  
   102  /*GetPath - get the path that can be used to verify the merkle tree */
   103  func (mt *MerkleTree) GetPath(hash Hashable) *MTPath {
   104  	hidx := mt.GetLeafIndex(hash)
   105  	if hidx < 0 {
   106  		return &MTPath{}
   107  	}
   108  	return mt.GetPathByIndex(hidx)
   109  }
   110  
   111  /*VerifyPath - given a leaf node and the path, verify that the node is part of the tree */
   112  func (mt *MerkleTree) VerifyPath(hash Hashable, path *MTPath) bool {
   113  	hs := hash.GetHash()
   114  	return VerifyMerklePath(hs, path, mt.GetRoot())
   115  }
   116  
   117  /*GetPathByIndex - get the path of a leaf node at index i */
   118  func (mt *MerkleTree) GetPathByIndex(idx int) *MTPath {
   119  	path := make([]string, mt.levels-1)
   120  	mpath := &MTPath{LeafIndex: idx}
   121  	if idx&1 == 1 {
   122  		path[0] = mt.tree[idx-1]
   123  	} else {
   124  		if idx+1 < mt.leavesCount {
   125  			path[0] = mt.tree[idx+1]
   126  		} else {
   127  			path[0] = mt.tree[idx]
   128  		}
   129  	}
   130  	for pl0, plsize, pi := 0, mt.leavesCount, 1; plsize > 2; pl0, plsize, pi = pl0+plsize, (plsize+1)/2, pi+1 {
   131  		l0 := pl0 + plsize
   132  		idx = (idx - idx&1) / 2
   133  		if idx&1 == 1 {
   134  			//path = append(path, mt.tree[l0+idx-1])
   135  			path[pi] = mt.tree[l0+idx-1]
   136  		} else {
   137  			if l0+idx+1 < l0+(plsize+1)/2 {
   138  				//path = append(path, mt.tree[l0+idx+1])
   139  				path[pi] = mt.tree[l0+idx+1]
   140  			} else {
   141  				//path = append(path, mt.tree[l0+idx])
   142  				path[pi] = mt.tree[l0+idx]
   143  			}
   144  		}
   145  	}
   146  	mpath.Nodes = path
   147  	return mpath
   148  }