Another pass: rough shape of tree and tests

It isn't working yet, the tests don't pass.

But it is on the way!
This commit is contained in:
Luke Murphy 2019-08-01 13:06:32 +02:00
parent c89cafe022
commit 45e24073d2
No known key found for this signature in database
GPG Key ID: 5E2EF5A63E3718CC
4 changed files with 154 additions and 25 deletions

View File

@ -5,11 +5,10 @@ from typing import Any, Callable, List, Optional
import attr
from flat_tree import FlatTreeAccessor
__all__ = ['MerkleTreeGenerator', 'MerkleTreeNode']
Hash = str
__all__ = ['MerkleTreeGenerator', 'MerkleTreeNode']
flat_tree = FlatTreeAccessor()
@ -17,17 +16,17 @@ flat_tree = FlatTreeAccessor()
class MerkleTreeNode:
"""A node in a merkle tree.
:param index: TODO
:param parent: TODO
:param size: TODO
:param data: TODO
:param hash: TODO
:param index: The index of node
:param parent: The parent of the node
:param size: The size of the data
:param data: The data of the node
:param hash: The hash of the data
"""
index: int
parent: Optional[int]
parent: int
size: int
data: bytes
data: Optional[bytes]
hash: Optional[str] = None
def __attrs_post_init__(self) -> Any:
@ -39,30 +38,69 @@ class MerkleTreeNode:
class MerkleTreeGenerator:
"""A stream that generates a merkle tree based on the incoming data.
:param leaf: TODO
:param parent: TODO
:param roots: TODO
:param blocks: TODO
:param leaf: The leaf hash generation function
:param parent: The parent hash generation function
:param roots: The tree roots
"""
leaf: Callable[[bytes], Hash]
parent: Callable[[bytes], Hash]
blocks: int
roots: Optional[List[MerkleTreeNode]] = attr.Factory(list)
leaf: Callable[[MerkleTreeNode, List[MerkleTreeNode]], Hash]
parent: Callable[[MerkleTreeNode, List[MerkleTreeNode]], Hash]
roots: List[MerkleTreeNode] = attr.Factory(list)
def next(self, data: bytes) -> List[MerkleTreeNode]:
"""Further generate the treem based on the incoming data.
def next(
self, data: bytes, nodes: Optional[List[MerkleTreeNode]] = None
) -> List[MerkleTreeNode]:
"""Further generate the tree based on the incoming data.
:param data: Incoming data
:param nodes: Pre-existing nodes
"""
pass
nodes = nodes or []
index = 2 * (self.blocks + 1)
leaf_node = MerkleTreeNode(
index=index,
parent=flat_tree.parent(index),
hash=None,
data=data,
size=len(data),
)
leaf_node.hash = self.leaf(leaf_node, self.roots)
self.roots.append(leaf_node)
nodes.append(leaf_node)
while len(self.roots) > 1:
left = self.roots[len(self.roots) - 2]
right = self.roots[len(self.roots) - 1]
if left.parent != right.parent:
break
self.roots.pop()
new_node = MerkleTreeNode(
index=left.parent,
parent=flat_tree.parent(left.parent),
hash=self.parent(left, [right]),
size=left.size + right.size,
data=None,
)
self.roots[len(self.roots) - 1] = new_node
nodes.append(new_node)
return nodes
def __attrs_post_init__(self) -> Any:
"""Initialise parent and block defaults."""
index = self.roots[len(self.roots) - 1].index
right_span = flat_tree.right_span(index)
self.blocks = (1 + (right_span / 2)) if self.roots else 0
for root in self.roots:
if not root.parent:
root.parent = flat_tree.parent(root.index)
# https://github.com/mafintosh/merkle-tree-stream/blob/master/generator.js#L14
# self.roots[self.roots.length] ...
# self.blocks = (1 + (flat_tree.right_span(...) / 2)) if self.roots else 0

View File

@ -45,6 +45,9 @@ package_dir =
= .
packages = find:
zip_safe = False
install_requires =
attrs >= 19.1.0, < 20.0
flat-tree == 0.0.1a3 # TODO(decentral1se): use bounds when 0.0.1 lands
[options.packages.find]
where = .

22
test/conftest.py Normal file
View File

@ -0,0 +1,22 @@
import hashlib
import pytest
@pytest.fixture
def leaf():
def _leaf(node):
return hashlib.sha256(leaf.data).hexdigest()
return _leaf
@pytest.fixture
def parent():
def _parent(left, right):
sha256 = hashlib.sha256()
sha256.update(left)
sha256.update(right)
return sha256.hexdigest()
return _parent

66
test/test_generator.py Normal file
View File

@ -0,0 +1,66 @@
"""Generator test module."""
import hashlib
from merkle_tree_stream import MerkleTreeGenerator, MerkleTreeNode
def test_hashes(leaf, parent):
stream = MerkleTreeGenerator(leaf=leaf, parent=parent)
stream.next(b'a')
first_node = (
MerkleTreeNode(
index=0,
parent=1,
hash=hashlib.sha256(b'a').hexdigest(),
size=1,
data=b'a',
),
)
stream.next(b'b')
second_node = (
MerkleTreeNode(
index=2,
parent=1,
hash=hashlib.sha256(b'b').hexdigest(),
size=1,
data=b'a',
),
)
stream.next(b'c')
third = hashlib.sha256(b'a')
third.update(b'b')
third_hash = third.hexdigest()
third_node = (
MerkleTreeNode(index=1, parent=3, hash=third_hash, size=2, data=b'a'),
)
assert stream.nodes == [first_node, second_node, third_node]
def test_single_root(leaf, parent):
stream = MerkleTreeGenerator(leaf=leaf, parent=parent)
stream.next('a')
stream.next('b')
stream.next('c')
stream.next('d')
assert stream.roots.length == 1
def multiple_roots(leaf, parent):
stream = MerkleTreeGenerator(leaf=leaf, parent=parent)
stream.next('a')
stream.next('b')
stream.next('c')
assert stream.roots.length > 1