refactor(rt): move runtime locals into slot store

Author: Stevengre

kmir/src/kmir/_prove.py

@@ -11,7 +11,7 @@
 from pyk.kast.manip import abstract_term_safely, split_config_from
 from pyk.kcfg import KCFG
 from pyk.kcfg.explore import KCFGExplore
-from pyk.kore.rpc import BoosterServer, KoreClient
+from pyk.kore.rpc import BoosterServer, DefaultError, KoreClient
 from pyk.proof.proof import parallel_advance_proof
 from pyk.proof.reachability import APRProof, APRProver
 
@@ -42,14 +42,14 @@ def prove(opts: ProveOpts) -> APRProof:
 
     if opts.proof_dir is not None:
         target_path = opts.proof_dir / label
-        return _prove(opts, target_path, label)
+        return _prove(opts, target_path, label, allow_rpc_recovery=False)
 
     with tempfile.TemporaryDirectory() as tmp_dir:
         target_path = Path(tmp_dir)
-        return _prove(opts, target_path, label)
+        return _prove(opts, target_path, label, allow_rpc_recovery=True)
 
 
-def _prove(opts: ProveOpts, target_path: Path, label: str) -> APRProof:
+def _prove(opts: ProveOpts, target_path: Path, label: str, *, allow_rpc_recovery: bool) -> APRProof:
     if not opts.reload and opts.proof_dir is not None and APRProof.proof_data_exists(label, opts.proof_dir):
         _LOGGER.info(f'Reading proof from disc: {opts.proof_dir}, {label}')
         proof = APRProof.read_proof_data(opts.proof_dir, label)
@@ -97,12 +97,13 @@ def _prove(opts: ProveOpts, target_path: Path, label: str) -> APRProof:
             break_on_function=opts.break_on_function or None,
         )
 
+        proof_root = opts.proof_dir if opts.proof_dir is not None else target_path
         proof = apr_proof_from_smir(
             kmir,
             label,
             smir_info,
             start_symbol=opts.start_symbol,
-            proof_dir=opts.proof_dir,
+            proof_dir=proof_root,
         )
         if proof.proof_dir is not None and (proof.proof_dir / label).is_dir():
             smir_info.dump(proof.proof_dir / proof.id / 'smir.json')
@@ -127,12 +128,17 @@ def _prove(opts: ProveOpts, target_path: Path, label: str) -> APRProof:
         break_every_step=opts.break_every_step,
         break_on_function=opts.break_on_function,
     )
-
-    if opts.max_workers and opts.max_workers > 1:
-        _prove_parallel(kmir, proof, opts=opts, label=label, cut_point_rules=cut_point_rules)
-    else:
-        _prove_sequential(kmir, proof, opts=opts, label=label, cut_point_rules=cut_point_rules)
-    return proof
+    try:
+        if opts.max_workers and opts.max_workers > 1:
+            _prove_parallel(kmir, proof, opts=opts, label=label, cut_point_rules=cut_point_rules)
+        else:
+            _prove_sequential(kmir, proof, opts=opts, label=label, cut_point_rules=cut_point_rules)
+        return proof
+    except (DefaultError, RuntimeError) as err:
+        recovered = _recover_proof_on_rpc_error(proof, err, allow_rpc_recovery=allow_rpc_recovery)
+        if recovered is not None:
+            return recovered
+        raise
 
 
 def _prove_parallel(
@@ -167,6 +173,7 @@ def create_prover() -> APRProver:
             cterm_symbolic = CTermSymbolic(
                 client,
                 kmir.definition,
+                log_succ_rewrites=_record_proof_logs(opts),
             )
             kcfg_explore = KCFGExplore(
                 cterm_symbolic,
@@ -197,7 +204,11 @@ def _prove_sequential(
     label: str,
     cut_point_rules: list[str],
 ) -> None:
-    with kmir.kcfg_explore(label, terminate_on_thunk=opts.terminate_on_thunk) as kcfg_explore:
+    with kmir.kcfg_explore(
+        label,
+        terminate_on_thunk=opts.terminate_on_thunk,
+        log_succ_rewrites=_record_proof_logs(opts),
+    ) as kcfg_explore:
         prover = APRProver(
             kcfg_explore,
             execute_depth=opts.max_depth,
@@ -211,6 +222,30 @@ def _prove_sequential(
         )
 
 
+def _record_proof_logs(opts: ProveOpts) -> bool:
+    # Persisted proofs may later be sectioned using stored rewrite logs.
+    # Ephemeral test proofs do not need them, and omitting them avoids huge RPC payloads.
+    return opts.proof_dir is not None
+
+
+def _recover_proof_on_rpc_error(proof: APRProof, err: Exception, *, allow_rpc_recovery: bool) -> APRProof | None:
+    if not allow_rpc_recovery:
+        return None
+
+    if proof.proof_dir is None or not APRProof.proof_data_exists(proof.id, proof.proof_dir):
+        return None
+
+    if isinstance(err, RuntimeError) and str(err) != 'Empty response received':
+        return None
+
+    recovered = APRProof.read_proof_data(proof.proof_dir, proof.id)
+    if recovered.passed or recovered.failed:
+        _LOGGER.warning(f'Recovered saved proof after RPC error: {proof.id}')
+        return recovered
+
+    return None
+
+
 def apr_proof_from_smir(
     kmir: KMIR,
     id: str,

kmir/src/kmir/kast.py

@@ -7,7 +7,7 @@
 
 from pyk.kast.inner import KApply, KSequence, KSort, KVariable, Subst, build_cons
 from pyk.kast.manip import free_vars, split_config_from
-from pyk.kast.prelude.collections import list_empty, list_of
+from pyk.kast.prelude.collections import list_empty, list_of, map_of
 from pyk.kast.prelude.kint import eqInt, intToken, leInt
 from pyk.kast.prelude.ml import mlEqualsTrue
 from pyk.kast.prelude.utils import token
@@ -19,12 +19,11 @@
     BoolValue,
     DynamicSize,
     IntValue,
-    Local,
     Metadata,
-    Place,
     PtrLocalValue,
     RangeValue,
     RefValue,
+    SlotPlace,
     StaticSize,
 )
 
@@ -192,12 +191,16 @@ def init_subst() -> dict[str, KInner]:
             k_cell,
         )
 
+    slot_ids = [token(i) for i in range(len(localvars))]
     subst = Subst(
         {
             **init_subst(),
             **{
                 'K_CELL': k_cell,
-                'LOCALS_CELL': list_of(localvars),
+                'OWNEDSLOTS_CELL': list_of(slot_ids),
+                'SLOTSTORE_CELL': map_of(zip(slot_ids, localvars, strict=True)),
+                'NEXTSLOT_CELL': token(len(slot_ids)),
+                'GENERATEDCOUNTER_CELL': token(len(slot_ids)),
             },
         }
     )
@@ -215,11 +218,15 @@ def _make_symbolic_call_config(
     types: Mapping[Ty, TypeMetadata],
 ) -> tuple[KInner, list[KInner]]:
     locals, constraints = _symbolic_locals(fn_data.args, types)
+    slot_ids = [token(i) for i in range(len(locals))]
     subst = Subst(
         {
             'K_CELL': fn_data.call_terminator,
             'STACK_CELL': list_empty(),  # FIXME see #560, problems matching a symbolic stack
-            'LOCALS_CELL': list_of(locals),
+            'OWNEDSLOTS_CELL': list_of(slot_ids),
+            'SLOTSTORE_CELL': map_of(zip(slot_ids, locals, strict=True)),
+            'NEXTSLOT_CELL': token(len(slot_ids)),
+            'GENERATEDCOUNTER_CELL': token(len(slot_ids)),
         },
     )
     empty_config = definition.empty_config(KSort('GeneratedTopCell'))
@@ -373,7 +380,11 @@ def _symbolic_value(self, ty: Ty, mutable: bool) -> tuple[KInner, Iterable[KInne
                     mlEqualsTrue(leInt(variant_var, token(max_variant))),
                 ]
                 args = self._fresh_var('ENUM_ARGS')
-                return KApply('Value::Aggregate', (KApply('variantIdx', (variant_var,)), args)), idx_range, None
+                return (
+                    KApply('Value::Aggregate', (KApply('variantIdx', (variant_var,)), args)),
+                    idx_range + [mlEqualsTrue(KApply('allValues', (args,)))],
+                    None,
+                )
 
             case StructT(_, _, fields):
                 field_vars: list[KInner] = []
@@ -390,14 +401,18 @@ def _symbolic_value(self, ty: Ty, mutable: bool) -> tuple[KInner, Iterable[KInne
 
             case UnionT():
                 args = self._fresh_var('ARG_UNION')
-                return KApply('Value::Aggregate', (KApply('variantIdx', (token(0),)), args)), [], None
+                return (
+                    KApply('Value::Aggregate', (KApply('variantIdx', (token(0),)), args)),
+                    [mlEqualsTrue(KApply('allValues', (args,)))],
+                    None,
+                )
 
             case ArrayT(_, None):
                 elems = self._fresh_var('ARG_ARRAY')
                 l = self._fresh_var('ARG_ARRAY_LEN')
                 return (
                     KApply('Value::Range', (elems,)),
-                    [mlEqualsTrue(eqInt(KApply('sizeList', (elems,)), l))],
+                    [mlEqualsTrue(eqInt(KApply('sizeList', (elems,)), l)), mlEqualsTrue(KApply('allValues', (elems,)))],
                     KApply(
                         'Metadata',
                         (
@@ -450,8 +465,7 @@ def _symbolic_value(self, ty: Ty, mutable: bool) -> tuple[KInner, Iterable[KInne
                     KApply(
                         'Value::Reference',
                         (
-                            token(0),  # Stack OFFSET field
-                            KApply('place', (KApply('local', (token(ref),)), KApply('ProjectionElems::empty', ()))),
+                            KApply('SlotPlace', (token(ref), KApply('ProjectionElems::empty', ()))),
                             KApply('Mutability::Mut', ()) if mutable else KApply('Mutability::Not', ()),
                             metadata if metadata is not None else no_metadata,
                         ),
@@ -468,8 +482,7 @@ def _symbolic_value(self, ty: Ty, mutable: bool) -> tuple[KInner, Iterable[KInne
                     KApply(
                         'Value::PtrLocal',
                         (
-                            token(0),
-                            KApply('place', (KApply('local', (token(ref),)), KApply('ProjectionElems::empty', ()))),
+                            KApply('SlotPlace', (token(ref), KApply('ProjectionElems::empty', ()))),
                             KApply('Mutability::Mut', ()) if mutable else KApply('Mutability::Not', ()),
                             metadata if metadata is not None else no_metadata,
                         ),
@@ -652,15 +665,13 @@ def _random_ptr_value(self, mut: bool, type_info: PtrT | RefT) -> PtrLocalValue
         match type_info:
             case PtrT():
                 return PtrLocalValue(
-                    stack_depth=0,
-                    place=Place(local=Local(ref)),
+                    place=SlotPlace(slot=ref),
                     mut=mut,
                     metadata=metadata,
                 )
             case RefT():
                 return RefValue(
-                    stack_depth=0,
-                    place=Place(local=Local(ref)),
+                    place=SlotPlace(slot=ref),
                     mut=mut,
                     metadata=metadata,
                 )

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

@@ -89,10 +89,11 @@ Execution gets stuck (no matching rule) when operands have different types or un
 ```k
   // Raw eq: dereference operands, extract types, and delegate to typed comparison
   rule <k> #execIntrinsic(IntrinsicFunction(symbol("raw_eq")), ARG1:Operand ARG2:Operand .Operands, PLACE, _SPAN)
-        => #execRawEqTyped(PLACE, #withDeref(ARG1), #extractOperandType(#withDeref(ARG1), LOCALS),
-                                  #withDeref(ARG2), #extractOperandType(#withDeref(ARG2), LOCALS))
+        => #execRawEqTyped(PLACE, #withDeref(ARG1), #extractOperandType(#withDeref(ARG1), STORE, SLOTS),
+                                  #withDeref(ARG2), #extractOperandType(#withDeref(ARG2), STORE, SLOTS))
        ... </k>
-       <locals> LOCALS </locals>
+       <currentFrame> <locals> SLOTS </locals> ... </currentFrame>
+       <slotStore> STORE </slotStore>
 
   // Compare values only if types are identical
   syntax KItem ::= #execRawEqTyped(Place, Evaluation, MaybeTy, Evaluation, MaybeTy) [seqstrict(2,4)]
@@ -112,17 +113,17 @@ Execution gets stuck (no matching rule) when operands have different types or un
   rule #withDeref(OP) => OP [owise]
 
   // Extract type from operands (locals with projections, constants, fallback to unknown)
-  syntax MaybeTy ::= #extractOperandType(Operand, List) [function, total]
-  rule #extractOperandType(operandCopy(place(local(I), PROJS)), LOCALS)
-       => getTyOf(tyOfLocal({LOCALS[I]}:>TypedLocal), PROJS)
-    requires 0 <=Int I andBool I <Int size(LOCALS) andBool isTypedLocal(LOCALS[I])
+  syntax MaybeTy ::= #extractOperandType(Operand, Map, List) [function, total]
+  rule #extractOperandType(operandCopy(place(local(I), PROJS)), STORE, SLOTS)
+       => getTyOf(tyOfLocal(getSlot(STORE, #frameSlotId(SLOTS, I))), PROJS)
+    requires 0 <=Int I andBool I <Int size(SLOTS) andBool isTypedLocal(getSlot(STORE, #frameSlotId(SLOTS, I)))
     [preserves-definedness]
-  rule #extractOperandType(operandMove(place(local(I), PROJS)), LOCALS)
-       => getTyOf(tyOfLocal({LOCALS[I]}:>TypedLocal), PROJS)
-    requires 0 <=Int I andBool I <Int size(LOCALS) andBool isTypedLocal(LOCALS[I])
+  rule #extractOperandType(operandMove(place(local(I), PROJS)), STORE, SLOTS)
+       => getTyOf(tyOfLocal(getSlot(STORE, #frameSlotId(SLOTS, I))), PROJS)
+    requires 0 <=Int I andBool I <Int size(SLOTS) andBool isTypedLocal(getSlot(STORE, #frameSlotId(SLOTS, I)))
     [preserves-definedness]
-  rule #extractOperandType(operandConstant(constOperand(_, _, mirConst(_, TY, _))), _) => TY
-  rule #extractOperandType(_, _) => TyUnknown [owise]
+  rule #extractOperandType(operandConstant(constOperand(_, _, mirConst(_, TY, _))), _, _) => TY
+  rule #extractOperandType(_, _, _) => TyUnknown [owise]
 ```
 
 #### Volatile Store (`std::intrinsics::volatile_store`, `std::ptr::write_volatile`)
@@ -191,8 +192,8 @@ the second argument, so the returned difference is always positive.
 
   rule <k> 
         #ptrOffsetDiff(
-          PtrLocal(HEIGHT, PLACE, _, metadata( _ , OFF1, _)),
-          PtrLocal(HEIGHT, PLACE, _, metadata( _ , OFF2, _)),
+          PtrLocal(PLACE, _, metadata( _ , OFF1, _)),
+          PtrLocal(PLACE, _, metadata( _ , OFF2, _)),
           SIGNED_FLAG,
           DEST
        ) => #setLocalValue(DEST, Integer(OFF1 -Int OFF2, 64, SIGNED_FLAG))
@@ -202,8 +203,8 @@ the second argument, so the returned difference is always positive.
 
   rule <k> 
         #ptrOffsetDiff(
-          PtrLocal(_, _, _, _) #as PTR1,
-          PtrLocal(_, _, _, _) #as PTR2,
+          PtrLocal(_, _, _) #as PTR1,
+          PtrLocal(_, _, _) #as PTR2,
           SIGNED_FLAG,
           _DEST
        ) => #UBErrorPtrOffsetDiff(PTR1, PTR2, SIGNED_FLAG)

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

@@ -33,5 +33,7 @@ module KMIR-AST
 
   syntax TypeMappings ::= List{TypeMapping, ""} [group(mir-list), symbol(TypeMappings::append), terminator-symbol(TypeMappings::empty)]
 
+  syntax Bool ::= allValues ( List ) [function, total, symbol(allValues)]
+
 endmodule
 ```