notded/dotnet/types.py

import enum
from itertools import zip_longest

from destruct import parse, Type, Struct
from .common import CLRCodedToken, CLRTableType, CLRElementType, MultiEnum



class CLRType:
    pass

class CLRPrimitiveType(CLRType):
    def __init__(self, type: CLRElementType):
        self.type = type

    def __str__(self):
        return self.type.name.lower()

class CLRPointerType(CLRType):
    def __init__(self, child: CLRType):
        self.child = child

    def __str__(self):
        return str(self.child) + '*'

class CLRRefType(CLRType):
    def __init__(self, child: CLRType):
        self.child = child

    def __str__(self):
        return str(self.child) + '&'

class CLRPinnedType(CLRType):
    def __init__(self, child: CLRType):
        self.child = child

    def __str__(self):
        return str(self.child) + ' fixed&'

class CLRUserTypeKind(enum.Enum):
    Class     = 'class'
    ValueType = 'enum'

class CLRUserType(CLRType):
    def __init__(self, token, kind):
        self.token = token
        self.kind = kind

    def __str__(self):
        if self.token.table == CLRTableType.TypeRef:
            prefix = 'ext-'
        else:
            prefix = ''
        return '{}{}#{}'.format(prefix, self.kind.value, self.token.row)

class CLROptionalType(CLRType):
    def __init__(self, child):
        self.child = child

    def __str__(self):
        return str(self.child) + '?'

class CLRRequiredType(CLRType):
    def __init__(self, child):
        self.child = child

    def __str__(self):
        return str(self.child) + '!'

class CLRGenericParamScope(enum.Enum):
    Type   = 'type'
    Method = 'method'

class CLRGenericParamType(CLRType):
    def __init__(self, index, scope):
        self.index = index
        self.scope = scope

    def __str__(self):
        return '#' + str(self.index)

class CLRGenericInstantiationType(CLRType):
    def __init__(self, child, args):
        self.child = child
        self.args = args

    def __str__(self):
        return str(self.child) + '[' + ',  '.join(str(a) for a in self.args) + ']'

class CLRArrayType(CLRType):
    def __init__(self, child, rank, sizes=None, lowers=None):
        self.child = child
        self.rank = rank
        self.sizes = sizes or []
        self.lowers = lowers or []

    def __str__(self):
        return str(self.child) + ''.join(
            '[{}{}]'.format(str(lower) + '...' if lower else '', upper + (lower or 0) if upper else '')
            for _, lower, upper in zip_longest(range(self.rank), self.lowers, self.sizes)
        )

class CLRElement(Type):
    def __init__(self, child):
        self.child = child

    def parse(self, input, context):
        nested_types = {
            CLRElementType.Pointer: CLRPointerType,
            CLRElementType.ByRef:   CLRRefType,
            CLRElementType.Pinned:  CLRPinnedType,
        }
        token_types = {
            CLRElementType.Class:       lambda t: CLRUserType(t, CLRUserTypeKind.Class),
            CLRElementType.ValueType:   lambda t: CLRUserType(t, CLRUserTypeKind.ValueType),
            CLRElementType.OptModifier: CLROptionalType,
            CLRElementType.ReqModifier: CLRRequiredType,
        }
        int_types = {
            CLRElementType.Var:  lambda i: CLRGenericParamType(i, CLRGenericParamScope.Type),
            CLRElementType.MVar: lambda i: CLRGenericParamType(i, CLRGenericParamScope.Method),
        }
        val = CLRElementType(parse(self.child, input, context))
        if val in nested_types:
            with context.enter(val.name, self):
                val = nested_types[val](self.parse(input, context))
        elif val in token_types:
            with context.enter(val.name, self.child):
                raw = parse(self.child, input, context)
                token = parse(CLRCodedToken([CLRTableType.TypeDef, CLRTableType.TypeRef]), raw.to_bytes(4, 'little'), context)
                val = token_types[val](token)
        elif val in int_types:
            with context.enter(val.name, self.child):
                val = int_types[val](parse(self.child, input, context))
        elif val == CLRElementType.Array:
            with context.enter(val.name, self):
                type = self.parse(input, context)
                rank = parse(self.child, input, context)
                nbounds = parse(self.child, input, context)
                bounds = [parse(self.child, input, context) for _ in range(nbounds)]
                nlo = parse(self.child, input, context)
                lo  = [parse(self.child, input, context) for _ in range(nlo)]
            val = CLRArrayType(type, rank, bounds, lo)
        elif val == CLRElementType.GenericInst:
            with context.enter(val.name, self):
                type = self.parse(input, context)
                nargs = parse(self.child, input, context)
                args = []
                for i in range(nargs):
                    with context.enter(i, self):
                        args.append(self.parse(input, context))
            val = CLRGenericInstantiationType(type, args)
        else:
            val = CLRPrimitiveType(val)
        return val


class CLRCompressedInt(Type):
    def __init__(self, signed=True):
        self.signed = signed

    def parse(self, input, context):
        val = input.read(1)[0]
        if (val >> 5) == 0b110:
            nbytes = 4
            b = bytearray([val & 0b11111])
        elif (val >> 6) == 0b10:
            nbytes = 2
            b = bytearray([val & 0b111111])
        else:
            nbytes = 1
            b = bytearray([val])
        b.extend(input.read(nbytes - 1))
        return int.from_bytes(b, byteorder='big', signed=self.signed)

class CLRCompressedUInt(Type):
    def __new__(self):
        return CLRCompressedInt(signed=False)

class CLRSignatureType(enum.Enum):
    Default  = 0x0
    C        = 0x1
    StdCall  = 0x2
    ThisCall = 0x3
    FastCall = 0x4
    VarArg   = 0x5
    Field    = 0x6
    LocalVar = 0x7
    Property = 0x8

class CLRSignatureFlags(enum.Flag):
    Generic      = 0x10
    HasThis      = 0x20
    ExplicitThis = 0x40

class CLRSignatureAttributes(MultiEnum):
    type  = (CLRSignatureType,  0x0F)
    flags = (CLRSignatureFlags, 0xF0)

class CLRMethodSignature(Struct):
    param_count = CLRCompressedUInt()
    ret_type    = CLRElement(CLRCompressedUInt())
    params      = Arr(CLRElement(CLRCompressedUInt()))

    def on_param_count(self, spec, context):
        spec.params.count = self.param_count

    def __str__(self):
        return '(' + ', '.join(str(p) for p in self.params) + ') -> ' + str(self.ret_type)

class CLRFieldSignature(Struct):
    type = CLRElement(CLRCompressedUInt())

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

class CLRPropertySignature(Struct):
    param_count = CLRCompressedUInt()
    type        = CLRElement(CLRCompressedUInt())
    params      = Arr(CLRElement(CLRCompressedUInt()))

    def on_param_count(self, spec, context):
        spec.params.count = self.param_count

    def __str__(self):
        return '(' + ', '.join(str(p) for p in self.params) + ') -> ' + str(self.ret_type)

class CLRLocalVarSignature(Struct):
    count = CLRCompressedUInt()
    vars  = Arr(CLRElement(CLRCompressedUInt()))

    def on_count(self, spec, context):
        spec.vars.count = self.count

class CLRSignature(Struct):
    attribs   = Enum(CLRSignatureAttributes, UInt(8))
    signature = Switch(options={
        CLRSignatureType.Default:  CLRMethodSignature,
        CLRSignatureType.C:        CLRMethodSignature,
        CLRSignatureType.StdCall:  CLRMethodSignature,
        CLRSignatureType.ThisCall: CLRMethodSignature,
        CLRSignatureType.FastCall: CLRMethodSignature,
        CLRSignatureType.VarArg:   CLRMethodSignature,
        CLRSignatureType.Field:    CLRFieldSignature,
        CLRSignatureType.Property: CLRPropertySignature,
        CLRSignatureType.LocalVar: CLRLocalVarSignature,
    })

    def on_attribs(self, spec, context):
        spec.signature.selector = self.attribs.type

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