eth2.0-specs/test_libs/pyspec/eth2spec/utils/ssz/ssz_partials.py

from ..merkle_minimal import hash, next_power_of_two
from .ssz_typing import (
    Container,
    infer_input_type,
    is_bool_type,
    is_bytes_type,
    is_bytesn_type,
    is_container_type,
    is_list_kind,
    is_list_type,
    is_uint_type,
    is_vector_kind,
    is_vector_type,
    read_elem_type,
    uint_byte_size,
)
from .ssz_impl import (
    chunkify,
    deserialize_basic,
    get_typed_values,
    is_basic_type,
    is_bottom_layer_kind,
    pack,
    serialize_basic,
)


ZERO_CHUNK = b'\x00' * 32


def last_power_of_two(x):
    return next_power_of_two(x + 1) // 2


def concat_generalized_indices(x, y):
    return x * last_power_of_two(y) + y - last_power_of_two(y)


def rebase(objs, new_root):
    return {concat_generalized_indices(new_root, k): v for k, v in objs.items()}


def constrict_generalized_index(x, q):
    depth = last_power_of_two(x // q)
    o = depth + x - q * depth
    if concat_generalized_indices(q, o) != x:
        return None
    return o


def unrebase(objs, q):
    o = {}
    for k, v in objs.items():
        new_k = constrict_generalized_index(k, q)
        if new_k is not None:
            o[new_k] = v
    return o


def filler(starting_position, chunk_count):
    at, skip, end = chunk_count, 1, next_power_of_two(chunk_count)
    value = ZERO_CHUNK
    o = {}
    while at < end:
        while at % (skip * 2) == 0:
            skip *= 2
            value = hash(value + value)
        o[(starting_position + at) // skip] = value
        at += skip
    return o


def merkle_tree_of_chunks(chunks, root):
    starting_index = root * next_power_of_two(len(chunks))
    o = {starting_index + i: chunk for i, chunk in enumerate(chunks)}
    o = {**o, **filler(starting_index, len(chunks))}
    return o


@infer_input_type
def ssz_leaves(obj, typ=None, root=1):
    if is_list_kind(typ):
        o = {root * 2 + 1: len(obj).to_bytes(32, 'little')}
        base = root * 2
    else:
        o = {}
        base = root
    if is_bottom_layer_kind(typ):
        data = serialize_basic(obj, typ) if is_basic_type(typ) else pack(obj, read_elem_type(typ))
        q = {**o, **merkle_tree_of_chunks(chunkify(data), base)}
        # print(obj, root, typ, base, list(q.keys()))
        return(q)
    else:
        fields = get_typed_values(obj, typ=typ)
        sub_base = base * next_power_of_two(len(fields))
        for i, (elem, elem_type) in enumerate(fields):
            o = {**o, **ssz_leaves(elem, typ=elem_type, root=sub_base + i)}
        q = {**o, **filler(sub_base, len(fields))}
        # print(obj, root, typ, base, list(q.keys()))
        return(q)


def fill(objects):
    objects = {k: v for k, v in objects.items()}
    keys = sorted(objects.keys())[::-1]
    pos = 0
    while pos < len(keys):
        k = keys[pos]
        if k in objects and k ^ 1 in objects and k // 2 not in objects:
            objects[k // 2] = hash(objects[k & - 2] + objects[k | 1])
            keys.append(k // 2)
        pos += 1
    return objects


@infer_input_type
def ssz_full(obj, typ=None):
    return fill(ssz_leaves(obj, typ=typ))


def get_basic_type_size(typ):
    if is_uint_type(typ):
        return uint_byte_size(typ)
    elif is_bool_type(typ):
        return 1
    else:
        raise Exception("Type not basic: {}".format(typ))


def get_bottom_layer_element_position(typ, base, length, index):
    """
    Returns the generalized index and the byte range of the index'th value
    in the list with the given base generalized index and given length
    """
    assert index < (1 if is_basic_type(typ) else length)
    elem_typ = typ if is_basic_type(typ) else read_elem_type(typ)
    elem_size = get_basic_type_size(elem_typ)
    chunk_index = index * elem_size // 32
    chunk_count = (1 if is_basic_type(typ) else length) * elem_size // 32
    generalized_index = base * next_power_of_two(chunk_count) + chunk_index
    start = elem_size * index % 32
    return generalized_index, start, start + elem_size


@infer_input_type
def get_generalized_indices(obj, path, typ=None, root=1):
    if len(path) == 0:
        return [root] if is_basic_type(typ) else list(ssz_leaves(obj, typ=typ, root=root).keys())
    if path[0] == '__len__':
        return [root * 2 + 1] if is_list_type(typ) else []
    base = root * 2 if is_list_kind(typ) else root
    if is_bottom_layer_kind(typ):
        length = 1 if is_basic_type(typ) else len(obj)
        index, _, _ = get_bottom_layer_element_position(typ, base, length, path[0])
        return [index]
    else:
        if is_container_type(typ):
            fields = typ.get_field_names()
            field_count, index = len(fields), fields.index(path[0])
            elem_type = typ.get_field_types()[index]
            child = obj.get_field_values()[index]
        else:
            field_count, index, elem_type, child = len(obj), path[0], read_elem_type(typ), obj[path[0]]
        return get_generalized_indices(
            child,
            path[1:],
            typ=elem_type,
            root=base * next_power_of_two(field_count) + index
        )


def get_branch_indices(tree_index):
    o = [tree_index, tree_index ^ 1]
    while o[-1] > 1:
        o.append((o[-1] // 2) ^ 1)
    return o[:-1]


def remove_redundant_indices(obj):
    return {k: v for k, v in obj.items() if not (k * 2 in obj and k * 2 + 1 in obj)}


def merge(*args):
    o = {}
    for arg in args:
        o = {**o, **arg}
    return fill(o)


@infer_input_type
def get_nodes_along_path(obj, path, typ=None):
    indices = get_generalized_indices(obj, path, typ=typ)
    return remove_redundant_indices(merge(
        *({i: obj.objects[i] for i in get_branch_indices(index)} for index in indices)
    ))


class OutOfRangeException(Exception):
    pass


class SSZPartial():
    def __init__(self, typ, objects, root=1):
        assert not is_basic_type(typ)
        self.objects = objects
        self.typ = typ
        self.root = root
        if is_container_type(self.typ):
            for field in self.typ.get_field_names():
                try:
                    setattr(self, field, self.getter(field))
                except OutOfRangeException:
                    pass

    def getter(self, index):
        base = self.root * 2 if is_list_kind(self.typ) else self.root
        if is_bottom_layer_kind(self.typ):
            tree_index, start, end = get_bottom_layer_element_position(
                self.typ, base, len(self), index
            )
            if tree_index not in self.objects:
                raise OutOfRangeException("Do not have required data")
            else:
                return deserialize_basic(
                    self.objects[tree_index][start:end],
                    self.typ if is_basic_type(self.typ) else read_elem_type(self.typ)
                )
        else:
            if is_container_type(self.typ):
                fields = self.typ.get_field_names()
                field_count, index = len(fields), fields.index(index)
                elem_type = self.typ.get_field_types()[index]
            else:
                field_count, index, elem_type = len(self), index, read_elem_type(self.typ)
            tree_index = base * next_power_of_two(field_count) + index
            if tree_index not in self.objects:
                raise OutOfRangeException("Do not have required data")
            if is_basic_type(elem_type):
                return deserialize_basic(self.objects[tree_index][:get_basic_type_size(elem_type)], elem_type)
            else:
                return ssz_partial(elem_type, self.objects, root=tree_index)

    def __getitem__(self, index):
        return self.getter(index)

    def __iter__(self):
        return (self[i] for i in range(len(self)))

    def __len__(self):
        if is_list_kind(self.typ):
            if self.root * 2 + 1 not in self.objects:
                raise OutOfRangeException("Do not have required data: {}".format(self.root * 2 + 1))
            return int.from_bytes(self.objects[self.root * 2 + 1], 'little')
        elif is_vector_kind(self.typ):
            return self.typ.length
        elif is_container_type(self.typ):
            return len(self.typ.get_fields())
        else:
            raise Exception("Unsupported type: {}".format(self.typ))

    def full_value(self):
        if is_bytes_type(self.typ) or is_bytesn_type(self.typ):
            return bytes([self.getter(i) for i in range(len(self))])
        elif is_list_kind(self.typ):
            return [self[i] for i in range(len(self))]
        elif is_vector_kind(self.typ):
            return self.typ(*(self[i] for i in range(len(self))))
        elif is_container_type(self.typ):
            def full_value(x):
                return x.full_value() if hasattr(x, 'full_value') else x

            return self.typ(**{field: full_value(self.getter(field)) for field in self.typ.get_field_names()})
        elif is_basic_type(self.typ):
            return self.getter(0)
        else:
            raise Exception("Unsupported type: {}".format(self.typ))

    def hash_tree_root(self):
        o = {**self.objects}
        keys = sorted(o.keys())[::-1]
        pos = 0
        while pos < len(keys):
            k = keys[pos]
            if k in o and k ^ 1 in o and k // 2 not in o:
                o[k // 2] = hash(o[k & - 2] + o[k | 1])
                keys.append(k // 2)
            pos += 1
        return o[self.root]

    def __str__(self):
        return str(self.full_value())


def ssz_partial(typ, objects, root=1):
    ssz_type = (
        Container if is_container_type(typ) else
        typ if (is_vector_type(typ) or is_bytesn_type(typ)) else object
    )

    class Partial(SSZPartial, ssz_type):
        pass

    if is_container_type(typ):
        Partial.__annotations__ = typ.__annotations__
    o = Partial(typ, objects, root=root)
    return o