Implement decoding for structs (#762)

Closes #761

Author: tothtamas28

kmir/src/kmir/decoding.py

@@ -21,6 +21,7 @@
     RefT,
     Single,
     StrT,
+    StructT,
     UintT,
 )
 from .value import (
@@ -149,6 +150,8 @@ def decode_value(data: bytes, type_info: TypeMetadata, types: Mapping[Ty, TypeMe
             return _decode_int(data, int_ty)
         case ArrayT(elem_ty, length):
             return _decode_array(data, elem_ty, length, types)
+        case StructT(fields=fields, layout=layout):
+            return _decode_struct(data=data, fields=fields, layout=layout, types=types)
         case EnumT(
             discriminants=discriminants,
             fields=fields,
@@ -219,6 +222,28 @@ def _decode_array(
     return RangeValue(elems)
 
 
+def _decode_struct(
+    *,
+    data: bytes,
+    fields: list[Ty],
+    layout: LayoutShape | None,
+    types: Mapping[Ty, TypeMetadata],
+) -> Value:
+    if not layout:
+        raise ValueError('Struct layout not provided')
+
+    offsets = _extract_offsets(layout.fields)
+
+    match layout.variants:
+        case Single(index=0):
+            pass
+        case _:
+            raise ValueError(f'Unexpected layout variants in struct: {layout.variants}')
+
+    field_values = _decode_fields(data=data, tys=fields, offsets=offsets, types=types)
+    return AggregateValue(0, field_values)
+
+
 def _decode_enum(
     *,
     data: bytes,

kmir/src/kmir/ty.py

@@ -146,7 +146,7 @@ def from_raw(data: Any) -> EnumT:
                     adt_def=adt_def,
                     discriminants=list(discriminants),
                     fields=[list(tys) for tys in fields],
-                    layout=LayoutShape.from_raw(layout),
+                    layout=LayoutShape.from_raw(layout) if layout is not None else None,
                 )
             case _:
                 raise _cannot_parse_as('EnumT', data)
@@ -467,6 +467,7 @@ class StructT(TypeMetadata):
     name: str
     adt_def: int
     fields: list[Ty]
+    layout: LayoutShape | None
 
     @staticmethod
     def from_raw(data: Any) -> StructT:
@@ -476,16 +477,25 @@ def from_raw(data: Any) -> StructT:
                     'name': name,
                     'adt_def': adt_def,
                     'fields': fields,
+                    'layout': layout,
                 }
             }:
                 return StructT(
                     name=name,
                     adt_def=adt_def,
                     fields=list(fields),
+                    layout=LayoutShape.from_raw(layout) if layout is not None else None,
                 )
             case _:
                 raise _cannot_parse_as('StructT', data)
 
+    def nbytes(self, types: Mapping[Ty, TypeMetadata]) -> int:
+        match self.layout:
+            case None:
+                raise ValueError(f'Cannot determine size, layout is missing for: {self}')
+            case LayoutShape(size=size):
+                return size.in_bytes
+
 
 @dataclass
 class UnionT(TypeMetadata):

kmir/src/tests/integration/data/decode-value/struct-simple-permuted-fields.expected

@@ -0,0 +1,3 @@
+Aggregate ( variantIdx ( 0 ) , ListItem ( Integer ( 15 , 8 , false ) )
+ListItem ( Integer ( 31 , 8 , false ) )
+ListItem ( BoolVal ( true ) ) )

kmir/src/tests/integration/data/decode-value/struct-simple-permuted-fields.json

@@ -0,0 +1,65 @@
+{
+  "bytes": [
+    1,
+    15,
+    31
+  ],
+  "types": [
+    [
+      0,
+      {
+        "PrimitiveType": {
+          "Uint": "U8"
+        }
+      }
+    ],
+    [
+      1,
+      {
+        "PrimitiveType": "Bool"
+      }
+    ]
+  ],
+  "typeInfo": {
+    "StructType": {
+      "name": "core::ops::RangeInclusive<u8>",
+      "adt_def": 100,
+      "fields": [
+        0,
+        0,
+        1
+      ],
+      "layout": {
+        "fields": {
+          "Arbitrary": {
+            "offsets": [
+              {
+                "num_bits": 8
+              },
+              {
+                "num_bits": 16
+              },
+              {
+                "num_bits": 0
+              }
+            ]
+          }
+        },
+        "variants": {
+          "Single": {
+            "index": 0
+          }
+        },
+        "abi": {
+          "Aggregate": {
+            "sized": true
+          }
+        },
+        "abi_align": 1,
+        "size": {
+          "num_bits": 24
+        }
+      }
+    }
+  }
+}

kmir/src/tests/integration/test_decode_value.py

@@ -52,14 +52,16 @@ def definition(definition_dir: Path) -> KDefinition:
     from pyk.kast.outer import read_kast_definition
 
     res = read_kast_definition(definition_dir / 'compiled.json')
+    _patch_definition(res)
+    return res
+
 
+def _patch_definition(definition: KDefinition) -> None:
     # Monkey patch __repr__ on the fixture to avoid flooding the output on test failure
-    cls = res.__class__
+    cls = definition.__class__
     new_repr = lambda self: repr('KMIR LLVM definition')
     new_cls = type(f'{cls.__name__}WithCustomRepr', (cls,), {'__repr__': new_repr})
-    object.__setattr__(res, '__class__', new_cls)
-
-    return res
+    object.__setattr__(definition, '__class__', new_cls)
 
 
 def dedent(s: str) -> str:
@@ -181,6 +183,7 @@ def load_test_types():
     'enum-option-nonzero-none',
     'enum-option-nonzero-some',
     'str',
+    'struct-simple-permuted-fields',
 )
 
 
@@ -232,15 +235,31 @@ def test_decode_value(
     assert test_data.expected == actual
 
 
+@pytest.fixture(scope='module')
+def kmir_definition_dir() -> Path:
+    from kmir.build import LLVM_DEF_DIR
+
+    return LLVM_DEF_DIR
+
+
+@pytest.fixture(scope='module')
+def kmir_definition(kmir_definition_dir: Path) -> KDefinition:
+    from pyk.kast.outer import read_kast_definition
+
+    res = read_kast_definition(kmir_definition_dir / 'compiled.json')
+    _patch_definition(res)
+    return res
+
+
 @pytest.mark.parametrize(
     'test_data',
     TEST_DATA,
     ids=[test_id for test_id, *_ in TEST_DATA],
 )
 def test_python_decode_value(
     test_data: _TestData,
-    definition_dir: Path,
-    definition: KDefinition,
+    kmir_definition_dir: Path,
+    kmir_definition: KDefinition,
     tmp_path: Path,
 ) -> None:
     from pyk.kast.inner import KSort
@@ -260,9 +279,9 @@ def test_python_decode_value(
         types=types,
     )
     kast = value.to_kast()
-    kore = kast_to_kore(definition, kast, KSort('Value'))
+    kore = kast_to_kore(kmir_definition, kast, KSort('Value'))
     actual = kore_print(
-        definition_dir=definition_dir,
+        definition_dir=kmir_definition_dir,
         pattern=kore,
         output='pretty',
     )