code_merkilization.py

code_merkilization.py

_kv_store: Dict[Hash32, bytes] = {}
_trie = BinaryTrie(_kv_store, BLANK_ROOT_HASH)


class MerklizedCode(NamedTuple):
    """
    Standalone implementation of Code Merkleization
    """
    code_hash: Hash32 # keccak(...)
    code_root: Hash32 # the actual merkle root
    code_length: int
    version: int
    chunks: Tuple[Tuple[int, bytes], ...]


def merkleize_code(code: bytes, chunk_size=CHUNK_SIZE) -> MerklizedCode:
    chunks = _chunkify(code, chunk_size)
    _merkleize_chunks(chunks)

    code_length = _get_big_endian(len(code), 4)
    code_hash = keccak(code)

    _trie.set(VERSION_KEY, VERSION)
    _trie.set(CODE_LEN_KEY, code_length)
    _trie.set(CODE_HASH_KEY, code_hash)

    return MerklizedCode(
        code_hash=code_hash,
        code_root=_trie.root_hash,
        code_length=code_length,
        chunks=chunks
    )


def _chunkify(code: bytes, chunk_size: int) -> Tuple[Tuple[int, bytes], ...]:
    chunks = []
    for i in range(0, len(code), chunk_size):
        chunks.append((0, code[i: i + chunk_size]))

    if len(chunks) >= MAX_CHUNK_COUNT:
        raise ValidationError(f"Number of chunk must be smaller than {MAX_CHUNK_COUNT}, got {chunk_size}")

    # set FIO
    # TODO: Add test cases
    for i, chunk in enumerate(chunks):
        for j, code_chunk in enumerate(chunk[Chunk.FIO]):
            if not PUSH1 <= code_chunk <= PUSH32:
                continue

            size = (code_chunk - PUSH1) + 1
            if j + size < chunk_size:
                continue

            next_fio = (j + size + 1) - chunk_size
            chunks[i + 1] = (int(next_fio), code_chunk)

    return tuple(chunks)


def _merkleize_chunks(chunks) -> MerklizedCode:
    for i, chunk in enumerate(chunks):
        key = _get_big_endian(i, KEY_LENGTH)
        # BE(FIO_i, 1) || code_i
        val = _get_big_endian(chunk[0], 1) + chunk[1]

        _trie.set(key, val)


def _get_big_endian(x, n: int) -> bytes:
    # casts x to an unsigned integer of n bytes and returns its big-endian representation
    # TODO: give title a better name
    # validate_uint32(x, title="Binary Trie")
    # return pack('>I', x)[-n:]
    return x.to_bytes(n, byteorder='big', signed=False)