Skip to content
34 changes: 28 additions & 6 deletions kmir/src/kmir/decoding.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@

from pyk.kast import KInner

from .alloc import DefId
from .ty import FieldsShape, LayoutShape, MachineSize, Scalar, TagEncoding, Ty, TypeMetadata, UintTy
from .value import MetadataSize

Expand All @@ -62,7 +63,12 @@ def to_kast(self) -> KInner:
)


def decode_alloc_or_unable(alloc_info: AllocInfo, types: Mapping[Ty, TypeMetadata]) -> Value:
def decode_alloc_or_unable(
alloc_info: AllocInfo,
types: Mapping[Ty, TypeMetadata],
statics: Mapping[DefId, Allocation] | None = None,
) -> Value:
statics = statics or {}
match alloc_info:
case AllocInfo(
ty=ty,
Expand All @@ -78,12 +84,28 @@ def decode_alloc_or_unable(alloc_info: AllocInfo, types: Mapping[Ty, TypeMetadat
data = bytes(n or 0 for n in bytez)
return _decode_memory_alloc_or_unable(data=data, ptrs=ptrs, ty=ty, types=types)
case AllocInfo(
ty=_,
# `Static` currently only carries `def_id`; we ignore it here.
global_alloc=Static(),
ty=ty,
global_alloc=Static(def_id=def_id),
):
# Static global alloc does not carry raw bytes here; leave as unable-to-decode placeholder
return UnableToDecodeValue('Static global allocation not decoded')
# The static's bytes live in its `MonoItemStatic` item, not in the alloc entry.
# The alloc's `ty` is a reference to the static; decode its contents against the pointee type.
allocation = statics.get(def_id)
if allocation is None:
return UnableToDecodeValue(f'Static allocation not found for def_id: {def_id}')

try:
type_info = types[ty]
except KeyError:
return UnableToDecodeValue(f'Decoding static allocation with unknown type: {ty}')

pointee_ty = _pointee_ty(type_info)
if pointee_ty is None:
return UnableToDecodeValue(f'Static allocation type is not a reference or a pointer: {type_info}')

data = bytes(n or 0 for n in allocation.bytez)
return _decode_memory_alloc_or_unable(
data=data, ptrs=allocation.provenance.ptrs, ty=pointee_ty, types=types
)
case AllocInfo(
ty=_,
global_alloc=Function(
Expand Down
2 changes: 1 addition & 1 deletion kmir/src/kmir/kompile.py
Original file line number Diff line number Diff line change
Expand Up @@ -662,7 +662,7 @@ def _decode_alloc(smir_info: SMIRInfo, raw_alloc: Any) -> tuple[KInner, KInner]:

alloc_id = raw_alloc['alloc_id']
alloc_info = smir_info.allocs[alloc_id]
value = decode_alloc_or_unable(alloc_info=alloc_info, types=smir_info.types)
value = decode_alloc_or_unable(alloc_info=alloc_info, types=smir_info.types, statics=smir_info.statics)

match value:
case UnableToDecodeValue(msg):
Expand Down
29 changes: 24 additions & 5 deletions kmir/src/kmir/smir.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
from functools import cached_property
from typing import TYPE_CHECKING, NewType

from .alloc import AllocInfo
from .alloc import Allocation, AllocInfo, DefId
from .ty import EnumT, RefT, StructT, Ty, TypeMetadata, UnionT

if TYPE_CHECKING:
Expand Down Expand Up @@ -54,6 +54,16 @@ def allocs(self) -> dict[AllocId, AllocInfo]:
alloc_info.alloc_id: alloc_info for alloc_info in (AllocInfo.from_dict(dct) for dct in self._smir['allocs'])
}

@cached_property
def statics(self) -> dict[DefId, Allocation]:
res: dict[DefId, Allocation] = {}
for item in self._smir['items']:
kind = item['mono_item_kind']
if 'MonoItemStatic' in kind:
mono_item_static = kind['MonoItemStatic']
res[DefId(mono_item_static['id'])] = Allocation.from_dict(mono_item_static['allocation'])
return res

@cached_property
def types(self) -> dict[Ty, TypeMetadata]:
return {Ty(id): TypeMetadata.from_raw(type) for id, type in self._smir['types']}
Expand Down Expand Up @@ -131,7 +141,11 @@ def function_symbols(self) -> dict[int, dict]:
fnc_symbols[-1] = {'NormalSym': self.main_symbol}

# function items not present in the SMIR lookup table are added with negative Ty ID
missing = [name for name in self.items.keys() if {'NormalSym': name} not in fnc_symbols.values()]
missing = [
name
for name, item in self.items.items()
if SMIRInfo._is_func(item) and {'NormalSym': name} not in fnc_symbols.values()
]

fake_ty = -2
for name in missing:
Expand Down Expand Up @@ -181,6 +195,10 @@ def spans(self) -> dict[int, tuple[Path, int, int, int, int]]:
def _is_func(item: dict[str, dict]) -> bool:
return 'MonoItemFn' in item['mono_item_kind']

@staticmethod
def _is_static(item: dict[str, dict]) -> bool:
return 'MonoItemStatic' in item['mono_item_kind']

def reduce_to(self, start_symbols: str | Sequence[str]) -> SMIRInfo:
# returns a new SMIRInfo with all _items_ removed that are not reachable from the named function(s)
match start_symbols:
Expand All @@ -199,10 +217,11 @@ def reduce_to(self, start_symbols: str | Sequence[str]) -> SMIRInfo:

new_smir = self._smir.copy() # shallow copy, but we can overwrite the `items`

# filter the new symbols to avoid key errors
# filter the new function symbols to avoid key errors
new_syms = [self.function_symbols[ty] for ty in reachable]
new_syms_ = [sym['NormalSym'] for sym in new_syms if 'NormalSym' in sym]
new_smir['items'] = [self.items[sym] for sym in new_syms_ if sym in self.items]
new_syms_ = {sym['NormalSym'] for sym in new_syms if 'NormalSym' in sym}
# Also keep the statics
new_smir['items'] = [item for sym, item in self.items.items() if SMIRInfo._is_static(item) or sym in new_syms_]

return SMIRInfo(new_smir)

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
static NUM: u8 = 55;

fn main() {
let num_ref = #
assert!(*num_ref == 55);
}
Loading
Loading