More enum decoding (#740)

* Decode single-variant enums
* Decode multi-variant enums with direct tag encoding

Author: tothtamas28

kmir/src/kmir/decoding.py

@@ -7,7 +7,23 @@
 from pyk.kast.prelude.string import stringToken
 
 from .alloc import Allocation, AllocInfo, Memory, ProvenanceEntry, ProvenanceMap
-from .ty import ArrayT, Bool, EnumT, Int, IntTy, PtrT, RefT, Str, Uint
+from .ty import (
+    ArbitraryFields,
+    ArrayT,
+    BoolT,
+    Direct,
+    EnumT,
+    Initialized,
+    IntT,
+    IntTy,
+    Multiple,
+    PrimitiveInt,
+    PtrT,
+    RefT,
+    Single,
+    StrT,
+    UintT,
+)
 from .value import (
     NO_METADATA,
     AggregateValue,
@@ -26,7 +42,7 @@
 
     from pyk.kast import KInner
 
-    from .ty import Ty, TypeMetadata, UintTy
+    from .ty import FieldsShape, LayoutShape, MachineSize, Scalar, TagEncoding, Ty, TypeMetadata, UintTy
     from .value import Metadata
 
 
@@ -126,16 +142,26 @@ def decode_value_or_unable(data: bytes, type_info: TypeMetadata, types: Mapping[
 
 def decode_value(data: bytes, type_info: TypeMetadata, types: Mapping[Ty, TypeMetadata]) -> Value:
     match type_info:
-        case Bool():
+        case BoolT():
             return _decode_bool(data)
-        case Str():
+        case StrT():
             return _decode_str(data)
-        case Uint(int_ty) | Int(int_ty):
+        case UintT(int_ty) | IntT(int_ty):
             return _decode_int(data, int_ty)
         case ArrayT(elem_ty, length):
             return _decode_array(data, elem_ty, length, types)
-        case EnumT(discriminants=discriminants, fields=fields):
-            return _decode_enum(data, discriminants, fields)
+        case EnumT(
+            discriminants=discriminants,
+            fields=fields,
+            layout=layout,
+        ):
+            return _decode_enum(
+                data=data,
+                discriminants=discriminants,
+                fields=fields,
+                layout=layout,
+                types=types,
+            )
         case _:
             raise ValueError(f'Unsupported type: {type_info}')
 
@@ -195,18 +221,145 @@ def _decode_array(
 
 
 def _decode_enum(
+    *,
+    data: bytes,
+    discriminants: list[int],
+    fields: list[list[Ty]],
+    layout: LayoutShape | None,
+    types: Mapping[Ty, TypeMetadata],
+) -> Value:
+    if not layout:
+        raise ValueError('Enum layout not provided')
+
+    offsets = _extract_offsets(layout.fields)
+
+    match layout.variants:
+        case Single(index):
+            return _decode_enum_single(
+                data=data,
+                discriminants=discriminants,
+                fields=fields,
+                offsets=offsets,
+                # ---
+                tag_index=index,
+                # ---
+                types=types,
+            )
+        case Multiple(
+            tag=tag,
+            tag_encoding=tag_encoding,
+            tag_field=tag_field,
+            variants=variants,
+        ):
+            return _decode_enum_multiple(
+                data=data,
+                discriminants=discriminants,
+                fields=fields,
+                offsets=offsets,
+                # ---
+                tag=tag,
+                tag_encoding=tag_encoding,
+                tag_field=tag_field,
+                variant_layouts=variants,
+                # ---
+                types=types,
+            )
+        case _:
+            raise AssertionError('Undhandled case')
+
+
+def _extract_offsets(fields_shape: FieldsShape) -> list[MachineSize]:
+    match fields_shape:
+        case ArbitraryFields(offsets=offsets):
+            return offsets
+        case _:
+            raise ValueError(f'Unsupported fields shape: {fields_shape}')
+
+
+def _decode_enum_single(
+    *,
     data: bytes,
     discriminants: list[int],
     fields: list[list[Ty]],
+    offsets: list[MachineSize],
+    tag_index: int,
+    types: Mapping[Ty, TypeMetadata],
 ) -> Value:
-    # The only supported case for now is when there are no fields
-    if any(tys for tys in fields):
-        raise ValueError('TODO - implement this case')
+    assert len(fields) == 1, 'Expected a single list of field types for single-variant enum'
+    tys = fields[0]
+
+    assert len(discriminants) == 1, 'Expected a single discriminant for single-variant enum'
+    discriminant = discriminants[0]
+    assert tag_index == discriminant, 'Assumed tag_index to be the same as the discriminant'
+
+    field_values = _decode_fields(data=data, tys=tys, offsets=offsets, types=types)
+    return AggregateValue(0, field_values)
+
+
+def _decode_enum_multiple(
+    *,
+    data: bytes,
+    discriminants: list[int],
+    fields: list[list[Ty]],
+    offsets: list[MachineSize],
+    # ---
+    tag: Scalar,
+    tag_encoding: TagEncoding,
+    tag_field: int,
+    variant_layouts: list[LayoutShape],
+    # ---
+    types: Mapping[Ty, TypeMetadata],
+) -> Value:
+    if not isinstance(tag_encoding, Direct):
+        raise ValueError(f'Unsupported encoding: {tag_encoding}')
+
+    assert tag_field == 0, 'Assumed tag field to be zero'
+    assert len(offsets) == 1, 'Assumed offsets to only contain the tag offset'
+    tag_offset = offsets[0]
+    tag_value = _extract_tag_value(data=data, tag_offset=tag_offset, tag=tag)
 
-    tag = int.from_bytes(data, byteorder='little', signed=False)
     try:
-        variant_idx = discriminants.index(tag)
+        variant_idx = discriminants.index(tag_value)
     except ValueError as err:
-        raise ValueError(f'Tag not found: {tag}') from err
+        raise ValueError(f'Tag not found: {tag_value}') from err
+
+    tys = fields[variant_idx]
+
+    variant_layout = variant_layouts[variant_idx]
+    field_offsets = _extract_offsets(variant_layout.fields)
+    assert isinstance(variant_layout.variants, Single)
 
-    return AggregateValue(variant_idx, ())
+    field_values = _decode_fields(data=data, tys=tys, offsets=field_offsets, types=types)
+    return AggregateValue(variant_idx, field_values)
+
+
+def _decode_fields(
+    *,
+    data: bytes,
+    tys: list[Ty],
+    offsets: list[MachineSize],
+    types: Mapping[Ty, TypeMetadata],
+) -> list[Value]:
+    res: list[Value] = []
+    for ty, offset in zip(tys, offsets, strict=True):
+        type_info = types[ty]
+        size_in_bytes = type_info.nbytes(types)
+        field_data = data[offset.in_bytes : offset.in_bytes + size_in_bytes]
+        value = decode_value(field_data, type_info, types)
+        res.append(value)
+    return res
+
+
+def _extract_tag_value(*, data: bytes, tag_offset: MachineSize, tag: Scalar) -> int:
+    match tag:
+        case Initialized(
+            value=PrimitiveInt(
+                length=length,
+                signed=signed,
+            ),
+            valid_range=_,
+        ):
+            tag_data = data[tag_offset.in_bytes : tag_offset.in_bytes + length.value]
+            return int.from_bytes(tag_data, byteorder='little', signed=signed)
+        case _:
+            raise ValueError('Unsupported tag: {tag}')

kmir/src/kmir/kast.py

@@ -9,7 +9,7 @@
 from pyk.kast.prelude.ml import mlEqualsTrue
 from pyk.kast.prelude.utils import token
 
-from .ty import ArrayT, Bool, EnumT, Int, PtrT, RefT, StructT, TupleT, Uint, UnionT
+from .ty import ArrayT, BoolT, EnumT, IntT, PtrT, RefT, StructT, TupleT, UintT, UnionT
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
@@ -146,15 +146,15 @@ def _fresh_var(self, prefix: str) -> KVariable:
     def _symbolic_value(self, ty: Ty, mutable: bool) -> tuple[KInner, Iterable[KInner], KInner | None]:
         # returns: symbolic value of given type, related constraints, related pointer metadata
         match self.smir_info.types.get(ty):
-            case Int(info):
+            case IntT(info):
                 val, constraints = int_var(self._fresh_var('ARG_INT'), info.value, True)
                 return val, constraints, None
 
-            case Uint(info):
+            case UintT(info):
                 val, constraints = int_var(self._fresh_var('ARG_UINT'), info.value, False)
                 return val, constraints, None
 
-            case Bool():
+            case BoolT():
                 val, constraints = bool_var(self._fresh_var('ARG_BOOL'))
                 return val, constraints, None