Execute drop glue for Drop terminators

kmir/src/kmir/kdist/mir-semantics/kmir.md

@@ -219,6 +219,20 @@ NB that a stack height of `0` cannot occur here, because the compiler prevents l
 If the local `_0` does not have a value (i.e., it remained uninitialised), the function returns unit and writing the value is skipped.
 
 ```k
+  rule [termReturnIgnored]: <k> #execTerminator(terminator(terminatorKindReturn, _SPAN)) ~> _
+         =>
+           #execBlockIdx(TARGET)
+       </k>
+       <currentFunc> _ => CALLER </currentFunc>
+       <currentBody> _ => #getBlocks(CALLER) </currentBody>
+       <caller> CALLER => NEWCALLER </caller>
+       <dest> place(local(-1), .ProjectionElems) => NEWDEST </dest>
+       <target> someBasicBlockIdx(TARGET) => NEWTARGET </target>
+       <unwind> _ => UNWIND </unwind>
+       <locals> _ => NEWLOCALS </locals>
+       <stack> ListItem(StackFrame(NEWCALLER, NEWDEST, NEWTARGET, UNWIND, NEWLOCALS)) STACK => STACK </stack>
+    [priority(40)]
+
   rule [termReturnSome]: <k> #execTerminator(terminator(terminatorKindReturn, _SPAN)) ~> _
          =>
            #setLocalValue(DEST, #decrementRef(VAL)) ~> #execBlockIdx(TARGET)
@@ -474,6 +488,8 @@ An operand may be a `Reference` (the only way a function could access another fu
       =>
        ListItem(newLocal(TY, MUT)) #reserveFor(REST)
 
+  syntax Operand ::= operandValue ( Value )
+
   syntax KItem ::= #setArgsFromStack ( Int, Operands)
                  | #setArgFromStack ( Int, Operand)
                  | #execIntrinsic ( MonoItemKind, Operands, Place, Span )
@@ -493,6 +509,12 @@ An operand may be a `Reference` (the only way a function could access another fu
         ...
        </k>
 
+  rule <k> #setArgFromStack(IDX, operandValue(VAL))
+        =>
+           #setLocalValue(place(local(IDX), .ProjectionElems), #incrementRef(VAL))
+        ...
+       </k>
+
   rule <k> #setArgFromStack(IDX, operandCopy(place(local(I), .ProjectionElems)))
         =>
            #setLocalValue(place(local(IDX), .ProjectionElems), #incrementRef(getValue(CALLERLOCALS, I)))
@@ -669,12 +691,52 @@ Other terminators that matter at the MIR level "Runtime" are `Drop` and `Unreach
 Drops are elaborated to Noops but still define the continuing control flow. Unreachable terminators lead to a program error.
 
 ```k
-  rule [termDrop]: <k> #execTerminator(terminator(terminatorKindDrop(_PLACE, TARGET, _UNWIND), _SPAN))
+  syntax MaybeTy ::= #dropPlaceTy(Place, List) [function, total]
+
+  rule #dropPlaceTy(place(local(I), PROJS), LOCALS)
+    => getTyOf(tyOfLocal({LOCALS[I]}:>TypedLocal), PROJS)
+    requires 0 <=Int I andBool I <Int size(LOCALS)
+     andBool isTypedLocal(LOCALS[I])
+    [preserves-definedness]
+
+  rule #dropPlaceTy(_, _) => TyUnknown [owise]
+
+  syntax KItem ::= #execDropCall ( MaybeTy, Place, BasicBlockIdx, UnwindAction, Span )
+                 | #callDropGlue ( Ty, BasicBlockIdx, UnwindAction, Span )
+
+  rule [termDropGlue]: <k> #execTerminator(terminator(terminatorKindDrop(PLACE, TARGET, UNWIND), SPAN))
+         =>
+           #execDropCall(#lookupDropFunctionTy(#dropPlaceTy(PLACE, LOCALS)), PLACE, TARGET, UNWIND, SPAN)
+        ...
+       </k>
+       <locals> LOCALS </locals>
+
+  rule <k> #execDropCall(TyUnknown, _PLACE, TARGET, _UNWIND, _SPAN)
          =>
            #execBlockIdx(TARGET)
         ...
        </k>
 
+  rule <k> #execDropCall(FTY:Ty, PLACE, TARGET, UNWIND, SPAN)
+         =>
+           rvalueAddressOf(mutabilityMut, PLACE) ~> #callDropGlue(FTY, TARGET, UNWIND, SPAN)
+        ...
+       </k>
+
+  rule <k> PTR:Value ~> #callDropGlue(FTY, TARGET, UNWIND, SPAN)
+         =>
+           #execTerminatorCall(
+             FTY,
+             lookupFunction(FTY),
+             operandValue(PTR) .Operands,
+             place(local(-1), .ProjectionElems),
+             someBasicBlockIdx(TARGET),
+             UNWIND,
+             SPAN
+           )
+        ...
+       </k>
+
   syntax MIRError ::= "ReachedUnreachable"
 
   rule [termUnreachable]: <k> #execTerminator(terminator(terminatorKindUnreachable, _SPAN))

kmir/src/kmir/kdist/mir-semantics/rt/types.md

@@ -183,6 +183,14 @@ To make this function total, an optional `MaybeTy` is used.
   syntax MaybeTy ::= Ty
                    | "TyUnknown"
 
+  syntax MaybeTy ::= lookupDropFunctionTy ( Ty ) [function, total, symbol(lookupDropFunctionTy)]
+                   | #lookupDropFunctionTy ( MaybeTy ) [function, total]
+
+  rule #lookupDropFunctionTy(TY:Ty) => lookupDropFunctionTy(TY)
+  rule #lookupDropFunctionTy(TyUnknown) => TyUnknown
+
+  rule lookupDropFunctionTy(_) => TyUnknown [owise]
+
   syntax MaybeTy ::= #transparentFieldTy ( TypeInfo ) [function, total]
 
   rule #transparentFieldTy(typeInfoStructType(_, _, FIELD .Tys, LAYOUT)) => FIELD

kmir/src/kmir/kompile.py

@@ -487,6 +487,20 @@ def get_int_arg(app: KInner) -> int:
 
     type_equations = _make_stratified_rules(kmir, 'lookupTy', 'Ty', 'TypeInfo', 'ty', type_assocs, invalid_type)
 
+    drop_function_assocs: list[tuple[int, KInner]]
+    drop_function_assocs = [
+        (int(ty), KApply('ty', (intToken(int(drop_ty)),))) for ty, drop_ty in smir_info.drop_function_tys.items()
+    ]
+    drop_function_equations = _make_stratified_rules(
+        kmir,
+        'lookupDropFunctionTy',
+        'Ty',
+        'MaybeTy',
+        'ty',
+        drop_function_assocs,
+        KApply('TyUnknown_RT-TYPES_MaybeTy', ()),
+    )
+
     invalid_alloc_n = KApply(
         'InvalidAlloc(_)_RT-VALUE-SYNTAX_Evaluation_AllocId', (KApply('allocId', (KVariable('N'),)),)
     )
@@ -501,7 +515,7 @@ def get_int_arg(app: KInner) -> int:
     if break_on_function:
         break_on_rules.append(_mk_break_on_functions_rule(kmir, break_on_function))
 
-    return [*equations, *type_equations, *alloc_equations, *break_on_rules]
+    return [*equations, *type_equations, *drop_function_equations, *alloc_equations, *break_on_rules]
 
 
 def _functions(kmir: KMIR, smir_info: SMIRInfo) -> dict[int, KInner]:

kmir/src/kmir/smir.py

@@ -7,7 +7,7 @@
 from typing import TYPE_CHECKING, NewType
 
 from .alloc import AllocInfo
-from .ty import EnumT, RefT, StructT, Ty, TypeMetadata, UnionT
+from .ty import ArrayT, EnumT, PtrT, RefT, StructT, TupleT, Ty, TypeMetadata, UnionT
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -172,6 +172,37 @@ def function_tys(self) -> dict[str, int]:
             res[name] = fun_syms[sym][0]
         return res
 
+    @cached_property
+    def drop_function_tys(self) -> dict[Ty, Ty]:
+        res: dict[Ty, Ty] = {}
+
+        for sym, item in self.items.items():
+            if sym not in self.function_symbols_reverse:
+                continue
+
+            mono_item = item['mono_item_kind'].get('MonoItemFn')
+            if mono_item is None:
+                continue
+
+            if not mono_item['name'].startswith('std::ptr::drop_in_place::<'):
+                continue
+
+            body = mono_item.get('body')
+            if body is None or body['arg_count'] < 1:
+                continue
+
+            arg_ty = Ty(body['locals'][1]['ty'])
+            arg_type = self.types.get(arg_ty)
+            if not isinstance(arg_type, (PtrT, RefT)):
+                _LOGGER.debug(f'Skipping drop glue {sym}: unexpected argument type {arg_type}')
+                continue
+
+            pointee_ty = Ty(arg_type.pointee_type)
+            fn_ty = Ty(self.function_symbols_reverse[sym][0])
+            res[pointee_ty] = fn_ty
+
+        return res
+
     @cached_property
     def spans(self) -> dict[int, tuple[Path, int, int, int, int]]:
         return {id: (p, sr, sc, er, ec) for id, [p, sr, sc, er, ec] in self._smir['spans']}
@@ -225,29 +256,101 @@ def call_edges(self) -> dict[Ty, set[Ty]]:
             else:
                 called_tys = set()
                 for block in body['blocks']:
-                    if 'Call' not in block['terminator']['kind']:
-                        continue
-                    call = block['terminator']['kind']['Call']
-
-                    # 1. Direct call: the function being called
-                    func = call['func']
-                    if 'Constant' in func:
-                        called_tys.add(Ty(func['Constant']['const_']['ty']))
-
-                    # 2. Indirect call: function pointers passed as arguments
-                    #    These are ZeroSized constants whose ty is in the function table
-                    for arg in call.get('args', []):
-                        if 'Constant' in arg:
-                            const_ = arg['Constant'].get('const_', {})
-                            if const_.get('kind') == 'ZeroSized':
-                                ty = const_.get('ty')
-                                if isinstance(ty, int) and ty in function_tys:
-                                    called_tys.add(Ty(ty))
+                    terminator = block['terminator']['kind']
+
+                    if 'Call' in terminator:
+                        call = terminator['Call']
+
+                        # 1. Direct call: the function being called
+                        func = call['func']
+                        if 'Constant' in func:
+                            called_tys.add(Ty(func['Constant']['const_']['ty']))
+
+                        # 2. Indirect call: function pointers passed as arguments
+                        #    These are ZeroSized constants whose ty is in the function table
+                        for arg in call.get('args', []):
+                            if 'Constant' in arg:
+                                const_ = arg['Constant'].get('const_', {})
+                                if const_.get('kind') == 'ZeroSized':
+                                    ty = const_.get('ty')
+                                    if isinstance(ty, int) and ty in function_tys:
+                                        called_tys.add(Ty(ty))
+
+                    if 'Drop' in terminator:
+                        drop = terminator['Drop']
+                        drop_ty = self._place_ty(body, drop['place'])
+                        if drop_ty is None:
+                            continue
+                        drop_fn_ty = self.drop_function_tys.get(drop_ty)
+                        if drop_fn_ty is not None:
+                            called_tys.add(drop_fn_ty)
 
             for ty in self.function_symbols_reverse[sym]:
                 result[Ty(ty)] = called_tys
         return result
 
+    def _place_ty(self, body: dict, place: dict) -> Ty | None:
+        local = place.get('local')
+        if not isinstance(local, int):
+            return None
+        locals_ = body.get('locals', [])
+        if not (0 <= local < len(locals_)):
+            return None
+
+        current_ty: Ty | None = Ty(locals_[local]['ty'])
+        for projection in place.get('projection', []):
+            assert current_ty is not None
+            current_ty = self._projected_ty(current_ty, projection)
+            if current_ty is None:
+                return None
+
+        return current_ty
+
+    def _projected_ty(self, ty: Ty, projection: object) -> Ty | None:
+        type_info = self.types.get(ty)
+        if type_info is None:
+            return None
+
+        if projection == 'Deref':
+            if isinstance(type_info, (PtrT, RefT)):
+                return Ty(type_info.pointee_type)
+            return None
+
+        if not isinstance(projection, dict):
+            return None
+
+        if 'Field' in projection:
+            index, _ = projection['Field']
+            if isinstance(type_info, (StructT, UnionT)):
+                fields = type_info.fields
+            elif isinstance(type_info, TupleT):
+                fields = type_info.components
+            elif isinstance(type_info, EnumT):
+                # Field projection into enums is only well-defined after a downcast.
+                return None
+            else:
+                return None
+
+            if 0 <= index < len(fields):
+                return Ty(fields[index])
+            return None
+
+        if 'ConstantIndex' in projection and isinstance(type_info, ArrayT):
+            return Ty(type_info.element_type)
+
+        if 'Subslice' in projection and isinstance(type_info, ArrayT):
+            return ty
+
+        if 'OpaqueCast' in projection or 'Subtype' in projection:
+            key = 'OpaqueCast' if 'OpaqueCast' in projection else 'Subtype'
+            cast_ty = projection[key]
+            return Ty(cast_ty) if isinstance(cast_ty, int) else None
+
+        if 'Downcast' in projection:
+            return ty
+
+        return None
+
 
 def compute_closure(start: Ty, edges: dict[Ty, set[Ty]]) -> set[Ty]:
     work = deque([start])

kmir/src/tests/integration/data/prove-rs/interior-mut-fail.rs → kmir/src/tests/integration/data/prove-rs/interior-mut.rs

@@ -20,12 +20,10 @@ impl Counter {
 
 fn main() {
     let counter = Counter::new(0);
-    // println!("Before: {}", counter.get());
 
-    // We only have &counter, but can still mutate inside
+    // We only have &counter, but can still mutate inside.
     counter.increment();
     counter.increment();
 
     assert!(2 == counter.get());
-    // println!("After: {}", counter.get());
 }

kmir/src/tests/integration/test_integration.py

@@ -43,7 +43,6 @@
 }
 PROVE_SHOW_SPECS = [
     'local-raw-fail',
-    'interior-mut-fail',
     'interior-mut3-fail',
     'iter_next_3',
     'assert_eq_exp',