utils.py

from typing import List, Tuple

from wake.testing import keccak256


class MerkleTree:
    _is_ready: bool
    _leaves: List[bytes]
    _levels: List[List[bytes]]

    def __init__(self):
        self._is_ready = False
        self._leaves = []
        self._levels = []

    @property
    def root(self) -> bytes:
        if not self._is_ready:
            self._build_tree()
        return self._levels[-1][0]

    @property
    def values(self) -> Tuple[bytes, ...]:
        return tuple(self._leaves)

    def get_proof(self, index: int) -> List[bytes]:
        if not self._is_ready:
            self._build_tree()

        proof = []
        for level in self._levels[:-1]:
            if index % 2 == 0:
                proof.append(level[index + 1])
            else:
                proof.append(level[index - 1])
            index //= 2
        return proof

    def get_multiproof(self, indexes: List[int]) -> Tuple[List[bytes], List[bool]]:
        if not self._is_ready:
            self._build_tree()

        proof = []
        flags = []
        known = indexes
        assert known == sorted(known), "Leaves must be sorted"

        for level in self._levels[:-1]:
            new_known = []
            for i in known:
                if i % 2 == 0:
                    if i + 1 in known:
                        flags.append(True)
                    else:
                        flags.append(False)
                        if i + 1 < len(level):
                            proof.append(level[i + 1])
                        else:
                            proof.append(level[i])
                else:
                    if i - 1 in known:
                        pass  # already processed
                    else:
                        flags.append(False)
                        proof.append(level[i - 1])
                if len(new_known) == 0 or new_known[-1] != i // 2:
                    new_known.append(i // 2)
            known = new_known

        return proof, flags

    def add_leaf(self, leaf: bytes):
        self._leaves.append(leaf)
        self._is_ready = False

    def _build_tree(self) -> None:
        self._levels.append([keccak256(leaf) for leaf in self._leaves])
        while len(self._levels[-1]) > 1:
            self._levels.append(self._build_level(self._levels[-1]))
        self._is_ready = True

    def _build_level(self, level: List[bytes]) -> List[bytes]:
        if len(level) % 2 == 1:
            level.append(level[-1])
        return [
            keccak256(level[i] + level[i + 1]) if level[i] < level[i + 1]
            else keccak256(level[i + 1] + level[i])
            for i in range(0, len(level), 2)
        ]