-
Notifications
You must be signed in to change notification settings - Fork 8
/
Copy pathmerkletree.py
164 lines (144 loc) · 6.23 KB
/
merkletree.py
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
158
159
160
161
162
163
164
import math
import util
def merkle_parent(hash1, hash2):
"""Takes the binary hashes and calculates the hash256"""
# return the hash256 of hash1 + hash2
return util.hash256(hash1 + hash2)
def merkle_parent_level(hashes):
"""Takes a list of binary hashes and returns a list that's half
the length"""
# if the list has exactly 1 element raise an error
if len(hashes) == 1:
raise RuntimeError('Cannot take a parent level with only 1 item')
# if the list has an odd number of elements, duplicate the last one
# and put it at the end so it has an even number of elements
if len(hashes) % 2 == 1:
hashes.append(hashes[-1])
# initialize next level
parent_level = []
# loop over every pair (use: for i in range(0, len(hashes), 2))
for i in range(0, len(hashes), 2):
# get the merkle parent of the hashes at index i and i+1
parent = merkle_parent(hashes[i], hashes[i + 1])
# append parent to parent level
parent_level.append(parent)
# return parent level
return parent_level
def merkle_root(hashes):
"""Takes a list of binary hashes and returns the merkle root
"""
# current level starts as hashes
current_level = hashes
# loop until there's exactly 1 element
while len(current_level) > 1:
# current level becomes the merkle parent level
current_level = merkle_parent_level(current_level)
# return the 1st item of the current level
return current_level[0]
class MerkleTree:
def __init__(self, total):
self.total = total
# compute max depth math.ceil(math.log(self.total, 2))
self.max_depth = math.ceil(math.log(self.total, 2))
# initialize the nodes property to hold the actual tree
self.nodes = []
# loop over the number of levels (max_depth+1)
for depth in range(self.max_depth + 1):
# the number of items at this depth is
# math.ceil(self.total / 2**(self.max_depth - depth))
num_items = math.ceil(self.total / 2 ** (self.max_depth - depth))
# create this level's hashes list with the right number of items
level_hashes = [None] * num_items
# append this level's hashes to the merkle tree
self.nodes.append(level_hashes)
# set the pointer to the root (depth=0, index=0)
self.current_depth = 0
self.current_index = 0
def __repr__(self):
result = []
for depth, level in enumerate(self.nodes):
items = []
for index, h in enumerate(level):
if h is None:
short = 'None'
else:
short = '{}...'.format(h.hex()[:8])
if depth == self.current_depth and index == self.current_index:
items.append('*{}*'.format(short[:-2]))
else:
items.append('{}'.format(short))
result.append(', '.join(items))
return '\n'.join(result)
def up(self):
# reduce depth by 1 and halve the index
self.current_depth -= 1
self.current_index //= 2
def left(self):
# increase depth by 1 and double the index
self.current_depth += 1
self.current_index *= 2
def right(self):
# increase depth by 1 and double the index + 1
self.current_depth += 1
self.current_index = self.current_index * 2 + 1
def root(self):
return self.nodes[0][0]
def set_current_node(self, value):
self.nodes[self.current_depth][self.current_index] = value
def get_current_node(self):
return self.nodes[self.current_depth][self.current_index]
def get_left_node(self):
return self.nodes[self.current_depth + 1][self.current_index * 2]
def get_right_node(self):
return self.nodes[self.current_depth + 1][self.current_index * 2 + 1]
def is_leaf(self):
return self.current_depth == self.max_depth
def right_exists(self):
return len(self.nodes[self.current_depth + 1]) > self.current_index * 2 + 1
def populate_tree(self, flag_bits, hashes):
# populate until we have the root
while self.root() is None:
# if we are a leaf, we know this position's hash
if self.is_leaf():
# get the next bit from flag_bits: flag_bits.pop(0)
flag_bits.pop(0)
# set the current node in the merkle tree to the next hash: hashes.pop(0)
self.set_current_node(hashes.pop(0))
# go up a level
self.up()
else:
# get the left hash
left_hash = self.get_left_node()
# if we don't have the left hash
if left_hash is None:
# if the next flag bit is 0, the next hash is our current node
if flag_bits.pop(0) == 0:
# set the current node to be the next hash
self.set_current_node(hashes.pop(0))
# sub-tree doesn't need calculation, go up
self.up()
else:
# go to the left node
self.left()
elif self.right_exists():
# get the right hash
right_hash = self.get_right_node()
# if we don't have the right hash
if right_hash is None:
# go to the right node
self.right()
else:
# combine the left and right hashes
self.set_current_node(merkle_parent(left_hash, right_hash))
# we've completed this sub-tree, go up
self.up()
else:
# combine the left hash twice
self.set_current_node(merkle_parent(left_hash, left_hash))
# we've completed this sub-tree, go up
self.up()
if len(hashes) != 0:
raise RuntimeError('hashes not all consumed {}'.format(len(hashes)))
for flag_bit in flag_bits:
if flag_bit != 0:
raise RuntimeError('flag bits not all consumed')