pbcabi/pbcabi/model.py

'''
Specifies the ABI Spec model and how to deserialize the ABI, and additionally
how specific type specifications can deserialize their elements.

Follows the ABI specification quite closely. The full ABI specification can be found at:

    https://partisiablockchain.gitlab.io/documentation/abiv.html
'''

from frozendict import frozendict
import dataclasses
from typing import Mapping
from enum import Enum
import enforce_typing
from .binaryreader import BinaryReader
import pbcabi.data as data

Identifier = str


class SerializeMode(Enum):
    RPC = 0x01
    STATE = 0x02
    ZK = 0x03


class SimpleType(Enum):
    '''
    The specific simple type.
    '''

    U8 = 0X01
    U16 = 0X02
    U32 = 0X03
    U64 = 0X04
    U128 = 0X05
    U256 = 0X18
    I8 = 0X06
    I16 = 0X07
    I32 = 0X08
    I64 = 0X09
    I128 = 0X0a
    STRING = 0X0b
    BOOL = 0X0c
    ADDRESS = 0X0d
    HASH = 0X13
    PUBLIC_KEY = 0X14
    SIGNATURE = 0X15
    BLS_PUBLIC_KEY = 0X16
    BLS_SIGNATURE = 0X17


@enforce_typing.enforce_types
@dataclasses.dataclass(frozen=True, slots=True)
class TypeSpec:
    '''
    Supertype of all 'TypeSpec'.
    '''

    @staticmethod
    def read_from(reader: BinaryReader):
        '''
        Deserialize this 'TypeSpec' type from 'BinaryReader'.
        '''
        discriminant = reader.readUInt8()
        if type_spec_read_from := TYPE_SPEC_SUBTYPE_READ_FROM_BY_DISCRIMINANT.get(
                discriminant):
            type_spec = type_spec_read_from(reader)
            return type_spec
        assert False, "Unknown discriminant: 0x{:02x}".format(discriminant)


@enforce_typing.enforce_types
@dataclasses.dataclass(frozen=True, slots=True)
class NamedTypeRef(TypeSpec):
    '''
    Reference to a named type.
    '''

    idx: int

    DISCRIMINANT = 0x00

    @staticmethod
    def read_from(reader: BinaryReader):
        '''
        Deserialize this 'TypeSpec' type from 'BinaryReader'.
        '''
        assert reader.readUInt8() == NamedTypeRef.DISCRIMINANT
        return NamedTypeRef.read_from_inner(reader)

    @staticmethod
    def read_from_inner(reader: BinaryReader):
        '''
        Deserialize this 'TypeSpec' type from 'BinaryReader'.
        '''
        return NamedTypeRef(reader.readUInt8())

    def read_element_from(self, reader: BinaryReader, type_env,
                          mode: SerializeMode):
        '''
        Deserialize elements for this 'TypeSpec' from RPC 'BinaryReader'.
        '''
        return type_env[self.idx].read_element_from(reader, type_env, mode)


READ_SIMPLE_TYPE = {}

READ_SIMPLE_TYPE[SerializeMode.RPC] = {
    (SimpleType.BOOL): lambda r: r.readUInt8() != 0,
    (SimpleType.U8): BinaryReader.readUInt8,
    (SimpleType.U16): lambda r: r.readUIntBigEndian(2),
    (SimpleType.U32): BinaryReader.readUInt32BigEndian,
    (SimpleType.U64): lambda r: r.readUIntBigEndian(8),
    (SimpleType.U128): lambda r: r.readUIntBigEndian(16),
    (SimpleType.U256): lambda r: r.readUIntBigEndian(32),
    (SimpleType.I8): lambda r: r.readSignedIntBigEndian(1),
    (SimpleType.I16): lambda r: r.readSignedIntBigEndian(2),
    (SimpleType.I32): lambda r: r.readSignedIntBigEndian(4),
    (SimpleType.I64): lambda r: r.readSignedIntBigEndian(8),
    (SimpleType.I128): lambda r: r.readSignedIntBigEndian(16),
    (SimpleType.STRING): lambda r: r.readString(BinaryReader.readUInt32BigEndian),
    (SimpleType.ADDRESS): data.BlockchainAddress.read_from,
    (SimpleType.HASH): data.Hash.read_from,
    (SimpleType.PUBLIC_KEY): data.PublicKey.read_from,
    (SimpleType.SIGNATURE): data.Signature.read_from,
    (SimpleType.BLS_PUBLIC_KEY): data.BlsPublicKey.read_from,
    (SimpleType.BLS_SIGNATURE): data.BlsSignature.read_from,
}

READ_SIMPLE_TYPE[SerializeMode.STATE] = {
    (SimpleType.BOOL): lambda r: r.readUInt8() != 0,
    (SimpleType.U8): BinaryReader.readUInt8,
    (SimpleType.U16): lambda r: r.readUIntLittleEndian(2),
    (SimpleType.U32): BinaryReader.readUInt32LittleEndian,
    (SimpleType.U64): lambda r: r.readUIntLittleEndian(8),
    (SimpleType.U128): lambda r: r.readUIntLittleEndian(16),
    (SimpleType.U256): lambda r: r.readUIntLittleEndian(32),
    (SimpleType.I8): lambda r: r.readSignedIntLittleEndian(1),
    (SimpleType.I16): lambda r: r.readSignedIntLittleEndian(2),
    (SimpleType.I32): lambda r: r.readSignedIntLittleEndian(4),
    (SimpleType.I64): lambda r: r.readSignedIntLittleEndian(8),
    (SimpleType.I128): lambda r: r.readSignedIntLittleEndian(16),
    (SimpleType.STRING): lambda r: r.readString(BinaryReader.readUInt32LittleEndian),
    (SimpleType.ADDRESS): data.BlockchainAddress.read_from,
    (SimpleType.HASH): data.Hash.read_from,
    (SimpleType.PUBLIC_KEY): data.PublicKey.read_from,
    (SimpleType.SIGNATURE): data.Signature.read_from,
    (SimpleType.BLS_PUBLIC_KEY): data.BlsPublicKey.read_from,
    (SimpleType.BLS_SIGNATURE): data.BlsSignature.read_from,
}


@enforce_typing.enforce_types
@dataclasses.dataclass(frozen=True, slots=True)
class SimpleTypeSpec(TypeSpec):
    '''
    Simple type spec.
    '''

    type: SimpleType

    @staticmethod
    def read_from(reader: BinaryReader):
        '''
        Deserialize this 'TypeSpec' type from 'BinaryReader'.
        '''
        return SimpleTypeSpec(SimpleType[reader.readUInt8()])

    def read_element_from(self, reader: BinaryReader, type_env,
                          mode: SerializeMode):
        '''
        Deserialize elements for this 'TypeSpec' from RPC 'BinaryReader'.
        '''
        return READ_SIMPLE_TYPE[mode][self.type](reader)


SIZE_SIMPLE_TYPE_SPEC = SimpleTypeSpec(SimpleType.U32)


@enforce_typing.enforce_types
@dataclasses.dataclass(frozen=True, slots=True)
class CompositeTypeSpec(TypeSpec):
    '''
    Supertype of all composite type specs.
    '''


@enforce_typing.enforce_types
@dataclasses.dataclass(frozen=True, slots=True)
class VecTypeSpec(TypeSpec):
    '''
    Type spec for a dynamically-sized vector of some some type.
    '''

    type_elements: TypeSpec

    DISCRIMINANT = 0x0e

    @staticmethod
    def read_from(reader: BinaryReader):
        '''
        Deserialize this 'TypeSpec' type from 'BinaryReader'.
        '''
        type_elements = TypeSpec.read_from(reader)
        return VecTypeSpec(type_elements)

    def read_element_from(self, reader: BinaryReader, type_env,
                          mode: SerializeMode):
        '''
        Deserialize elements for this 'TypeSpec' from RPC 'BinaryReader'.
        '''
        length = SIZE_SIMPLE_TYPE_SPEC.read_element_from(reader, type_env, mode)

        if self.type_elements == SimpleTypeSpec(SimpleType.U8):
            return reader.readBytes(length)

        array = []
        for i in range(0, length):
            elem = self.type_elements.read_element_from(reader, type_env, mode)
            array.append(elem)
        return array


@enforce_typing.enforce_types
@dataclasses.dataclass(frozen=True, slots=True)
class MapTypeSpec(TypeSpec):
    '''
    Type spec for a dynamically-sized map from some type to another type.
    '''

    type_key: TypeSpec
    type_value: TypeSpec

    DISCRIMINANT = 0x0f

    @staticmethod
    def read_from(reader: BinaryReader):
        '''
        Deserialize this 'TypeSpec' type from 'BinaryReader'.
        '''
        type_key = TypeSpec.read_from(reader)
        type_value = TypeSpec.read_from(reader)
        return MapTypeSpec(type_key, type_value)

    def read_element_from(self, reader: BinaryReader, type_env,
                          mode: SerializeMode):
        '''
        Deserialize elements for this 'TypeSpec' from RPC 'BinaryReader'.
        '''
        out = {}
        length = SIZE_SIMPLE_TYPE_SPEC.read_element_from(reader, type_env, mode)
        for i in range(0, length):
            key = self.type_key.read_element_from(reader, type_env, mode)
            val = self.type_value.read_element_from(reader, type_env, mode)
            out[key] = val
        return out


@enforce_typing.enforce_types
@dataclasses.dataclass(frozen=True, slots=True)
class SetTypeSpec(TypeSpec):
    '''
    Type spec for a dynamically-sized set.
    '''

    type_elements: TypeSpec

    DISCRIMINANT = 0x10

    @staticmethod
    def read_from(reader: BinaryReader):
        '''
        Deserialize this 'TypeSpec' type from 'BinaryReader'.
        '''
        type_elements = TypeSpec.read_from(reader)
        return SetTypeSpec(type_elements)

    def read_element_from(self, reader: BinaryReader, type_env,
                          mode: SerializeMode):
        '''
        Deserialize elements for this 'TypeSpec' from RPC 'BinaryReader'.
        '''
        array = []
        length = SIZE_SIMPLE_TYPE_SPEC.read_element_from(reader, type_env, mode)
        for i in range(0, length):
            array.append(self.type_elements.read_element_from(reader, type_env),
                         mode)
        return frozenset(array)


@enforce_typing.enforce_types
@dataclasses.dataclass(frozen=True, slots=True)
class ArrayTypeSpec(TypeSpec):
    '''
    Type spec for byte-array.
    '''

    length: int

    DISCRIMINANT = 0x11

    @staticmethod
    def read_from(reader: BinaryReader):
        '''
        Deserialize this 'TypeSpec' type from 'BinaryReader'.
        '''
        length = reader.readUInt8()
        return ArrayTypeSpec(length)

    def read_element_from(self, reader: BinaryReader, type_env,
                          mode: SerializeMode):
        '''
        Deserialize elements for this 'TypeSpec' from RPC 'BinaryReader'.
        '''
        return reader.readBytes(self.length)

    def read_element_from(self, reader: BinaryReader, type_env,
                          mode: SerializeMode):
        '''
        Deserialize elements for this 'TypeSpec' from state 'BinaryReader'.
        '''
        return reader.readBytes(self.length)


@enforce_typing.enforce_types
@dataclasses.dataclass(frozen=True, slots=True)
class OptionTypeSpec(TypeSpec):
    '''
    Type spec for some option type containing a specific sub-type.
    '''

    type_elements: TypeSpec

    DISCRIMINANT = 0x12

    @staticmethod
    def read_from(reader: BinaryReader):
        '''
        Deserialize this 'TypeSpec' type from 'BinaryReader'.
        '''
        type_elements = TypeSpec.read_from(reader)
        return OptionTypeSpec(type_elements)

    def read_element_from(self, reader: BinaryReader, type_env,
                          mode: SerializeMode):
        '''
        Deserialize elements for this 'TypeSpec' from RPC 'BinaryReader'.
        '''
        discriminant = reader.readUInt8()
        if discriminant == 0:
            return None
        else:
            return self.type_elements.read_element_from(reader, type_env, mode)

    def read_element_from(self, reader: BinaryReader, type_env,
                          mode: SerializeMode):
        '''
        Deserialize elements for this 'TypeSpec' from state 'BinaryReader'.
        '''
        discriminant = reader.readUInt8()
        if discriminant == 0:
            return None
        else:
            return self.type_elements.read_element_from(reader, type_env, mode)


def simple_type_spec_read_from(variant):
    v = SimpleTypeSpec(variant)
    return lambda r: v


TYPE_SPEC_SUBTYPE_READ_FROM_BY_DISCRIMINANT = {
    0x00: NamedTypeRef.read_from_inner,
} | {
    v.value: simple_type_spec_read_from(v) for v in SimpleType
} | {
    t.DISCRIMINANT: t.read_from for t in
    [VecTypeSpec, MapTypeSpec, SetTypeSpec, ArrayTypeSpec, OptionTypeSpec]
}


@enforce_typing.enforce_types
@dataclasses.dataclass(frozen=True, slots = True, order=True)
class Version:
    '''
    Version field.
    '''
    major: int
    minor: int
    patch: int

    @staticmethod
    def read_from(reader: BinaryReader):
        '''
        Deserialize this struct from 'BinaryReader'.
        '''
        major = reader.readUInt8()
        minor = reader.readUInt8()
        patch = reader.readUInt8()
        return Version(major, minor, patch)


class FnKind(Enum):
    '''
    Function kinds that contract can be invoked with.
    '''

    INIT = 0x01
    ACTION = 0x02
    CALLBACK = 0x03
    ZK_SECRET_INPUT = 0x10
    ZK_VAR_INPUTTED = 0x11
    ZK_VAR_REJECTED = 0x12
    ZK_COMPUTE_COMPLETE = 0x13
    ZK_VAR_OPENED = 0x14
    ZK_USER_VAR_OPENED = 0x15
    ZK_ATTESTATION_COMPLETE = 0x16
    ZK_SECRET_INPUT_WITH_EXPLICIT_TYPE = 0x17

    @staticmethod
    def read_from(reader: BinaryReader):
        '''
        Deserialize this struct from 'BinaryReader'.
        '''
        discriminant = reader.readUInt8()
        if found := [kind for kind in FnKind if kind.value == discriminant]:
            assert len(found) == 1
            return found[0]
        assert False, "Unknown discriminant: 0x{:02x}".format(discriminant)


@enforce_typing.enforce_types
@dataclasses.dataclass(frozen=True, slots=True)
class FieldAbi:
    '''
    Field of a struct.
    '''

    name: Identifier
    type: TypeSpec

    @staticmethod
    def read_from(reader: BinaryReader):
        '''
        Deserialize this struct from 'BinaryReader'.
        '''
        name = reader.readString()
        type = TypeSpec.read_from(reader)
        return FieldAbi(name, type)


@enforce_typing.enforce_types
@dataclasses.dataclass(frozen=True, slots=True)
class ArgumentAbi:
    '''
    Argument of a function.
    '''
    name: Identifier
    type: TypeSpec

    @staticmethod
    def read_from(reader: BinaryReader):
        '''
        Deserialize this struct from 'BinaryReader'.
        '''
        name = reader.readString()
        type = TypeSpec.read_from(reader)
        return ArgumentAbi(name, type)


@enforce_typing.enforce_types
@dataclasses.dataclass(frozen=True, slots=True)
class NamedTypeSpec:
    '''
    Supertype of named types.
    '''

    name: Identifier
    type_index: int

    @staticmethod
    def read_from(reader: BinaryReader, type_index: int):
        '''
        Deserialize this 'TypeSpec' type from 'BinaryReader'.
        '''
        discriminant = reader.readUInt8()
        if discriminant == StructTypeSpec.DISCRIMINANT:
            return StructTypeSpec.read_from(reader, type_index)
        if discriminant == EnumTypeSpec.DISCRIMINANT:
            return EnumTypeSpec.read_from(reader, type_index)


@enforce_typing.enforce_types
@dataclasses.dataclass(frozen=True, slots=True)
class StructTypeSpec(NamedTypeSpec):
    '''
    Struct type specification.
    '''

    fields: list[FieldAbi]

    DISCRIMINANT = 0x01

    @staticmethod
    def read_from(reader: BinaryReader, type_index: int):
        '''
        Deserialize this 'TypeSpec' type from 'BinaryReader'.
        '''
        name = reader.readString()
        variants = reader.readList(FieldAbi.read_from)
        return StructTypeSpec(name, type_index, variants)

    def read_element_from(self, reader: BinaryReader, type_env,
                          mode: SerializeMode):
        '''
        Deserialize elements for this 'TypeSpec' from RPC 'BinaryReader'.
        '''
        result = {}
        for field in self.fields:
            result[field.name] = field.type.read_element_from(
                reader, type_env, mode)
        result['__type'] = self.name
        return result


@enforce_typing.enforce_types
@dataclasses.dataclass(frozen=True, slots=True)
class EnumVariant:
    '''
    Enum variant specification.
    '''

    discriminant: int
    definition: NamedTypeRef

    @staticmethod
    def read_from(reader: BinaryReader):
        '''
        Deserialize an 'EnumVariant' from 'BinaryReader'.
        '''
        discriminant = reader.readUInt8()
        type = NamedTypeRef.read_from_inner(reader)
        return EnumVariant(discriminant, type)


@enforce_typing.enforce_types
@dataclasses.dataclass(frozen=True, slots=True)
class EnumTypeSpec(NamedTypeSpec):
    '''
    Enum type specification.
    '''

    variants_by_discriminant: Mapping[int, EnumVariant]

    DISCRIMINANT = 0x02

    @staticmethod
    def read_from(reader: BinaryReader, type_index: int):
        '''
        Deserialize this 'TypeSpec' type from 'BinaryReader'.
        '''
        name = reader.readString()
        variants = reader.readList(EnumVariant.read_from)
        variants_by_discriminant = frozendict(
            {v.discriminant: v for v in variants})
        return EnumTypeSpec(name, type_index, variants_by_discriminant)

    def read_element_from(self, reader: BinaryReader, type_env,
                          mode: SerializeMode):
        '''
        Deserialize elements for this 'TypeSpec' from RPC 'BinaryReader'.
        '''
        discriminant = reader.readUInt8()
        variant = self.variants_by_discriminant[discriminant]
        return variant.definition.read_element_from(reader, type_env, mode)


@enforce_typing.enforce_types
@dataclasses.dataclass(frozen=True, slots=True)
class FnAbi:
    '''
    Function definition.
    '''

    kind: FnKind
    name: Identifier
    shortname: int
    arguments: list[ArgumentAbi]
    secret_argument: ArgumentAbi | None

    @staticmethod
    def read_from(reader: BinaryReader):
        '''
        Deserialize 'FnAbi' from 'BinaryReader'.
        '''
        kind = FnKind.read_from(reader)
        name = reader.readString()
        shortname = reader.readLeb128()
        arguments = reader.readList(ArgumentAbi.read_from)
        secret_argument = None
        if kind == FnKind.ZK_SECRET_INPUT_WITH_EXPLICIT_TYPE:
            secret_argument = ArgumentAbi.read_from(reader)

        return FnAbi(kind, name, shortname, arguments, secret_argument)

@enforce_typing.enforce_types
@dataclasses.dataclass(frozen=True, slots=True)
class ContractAbi:
    '''
    Contract definition.
    '''

    named_types_by_id: Mapping[Identifier, NamedTypeSpec]
    named_types_by_idx: Mapping[int, NamedTypeSpec]
    hooks: Mapping[Identifier, FnAbi]
    state_type: TypeSpec

    @staticmethod
    def read_from(reader: BinaryReader | str):
        '''
        Deserialize 'ContractAbi' from 'BinaryReader'.
        '''

        num_named_types = reader.readUInt32BigEndian()
        named_types = []
        for type_index in range(0, num_named_types):
            named_types.append(NamedTypeSpec.read_from(reader, type_index))

        hooks = reader.readList(FnAbi.read_from)
        state_type = TypeSpec.read_from(reader)
        return ContractAbi(frozendict({t.name: t for t in named_types}),
                           frozendict({t.type_index: t for t in named_types}),
                           frozendict({t.name: t for t in hooks}), state_type)

    def read_type_element_from_rpc(self, type_name: Identifier,
                                   rpc: BinaryReader):
        '''
        Reads element of the given type name from RPC 'BinaryReader'.
        '''
        return self.named_types_by_id[type_name].read_element_from(
            rpc, self.named_types_by_idx, SerializeMode.RPC)

    def read_state(self, reader: BinaryReader):
        if not isinstance(reader, BinaryReader):
            reader = BinaryReader(reader)

        return self.state_type.read_element_from(reader, self.named_types_by_idx,
                                                 SerializeMode.STATE)


@enforce_typing.enforce_types
@dataclasses.dataclass(frozen=True, slots=True)
class FileAbi:
    '''
    File definition.
    '''

    version_binder: Version
    version_client: Version
    contract: ContractAbi

    @staticmethod
    def read_from(reader: BinaryReader):
        '''
        Deserialize 'FileAbi' from 'BinaryReader'.
        '''

        if not isinstance(reader, BinaryReader):
            reader = BinaryReader(reader)

        header = reader.readBytes(6)
        assert header == b'PBCABI'
        version_binder = Version.read_from(reader)
        version_client = Version.read_from(reader)

        assert version_client >= Version(5, 0, 0), version_client

        contract = ContractAbi.read_from(reader)
        return FileAbi(version_binder, version_client, contract)