Reusing backend information to construct KCFGs more efficiently

pyk/src/pyk/kcfg/explore.py

@@ -218,19 +218,20 @@ def extend_cterm(
         cut_point_rules: Iterable[str] = (),
         terminal_rules: Iterable[str] = (),
         module_name: str | None = None,
-    ) -> KCFGExtendResult:
+    ) -> list[KCFGExtendResult]:
+
         def log(message: str, *, warning: bool = False) -> None:
             _LOGGER.log(logging.WARNING if warning else logging.INFO, f'Extend result for {self.id}: {message}')
 
         custom_step_result = self.kcfg_semantics.custom_step(_cterm)
         if custom_step_result is not None:
             log(f'custom step node: {node_id}')
-            return custom_step_result
+            return [custom_step_result]
 
         abstract_cterm = self.kcfg_semantics.abstract_node(_cterm)
         if _cterm != abstract_cterm:
             log(f'abstraction node: {node_id}')
-            return Abstract(abstract_cterm)
+            return [Abstract(abstract_cterm)]
 
         cterm, next_states, depth, vacuous, next_node_logs = self.cterm_symbolic.execute(
             _cterm,
@@ -240,33 +241,35 @@ def log(message: str, *, warning: bool = False) -> None:
             module_name=module_name,
         )
 
+        extend_results: list[KCFGExtendResult] = []
+
         # Basic block
         if depth > 0:
             log(f'basic block at depth {depth}: {node_id}')
-            return Step(cterm, depth, next_node_logs, self._extract_rule_labels(next_node_logs))
+            extend_results.append(Step(cterm, depth, next_node_logs, self._extract_rule_labels(next_node_logs)))
 
         # Stuck or vacuous
         if not next_states:
             if vacuous:
                 log(f'vacuous node: {node_id}', warning=True)
-                return Vacuous()
-            log(f'stuck node: {node_id}')
-            return Stuck()
-
+                extend_results.append(Vacuous())
+            else:
+                log(f'stuck node: {node_id}')
+                extend_results.append(Stuck())
         # Cut rule
-        if len(next_states) == 1:
+        elif len(next_states) == 1:
             log(f'cut-rule basic block at depth {depth}: {node_id}')
-            return Step(
-                next_states[0].state,
-                1,
-                next_node_logs,
-                self._extract_rule_labels(next_node_logs),
-                cut=True,
+            extend_results.append(
+                Step(
+                    next_states[0].state,
+                    1,
+                    next_node_logs,
+                    self._extract_rule_labels(next_node_logs),
+                    cut=True,
+                )
             )
-
         # Branch
-        assert len(next_states) > 1
-        if all(branch_constraint for _, branch_constraint in next_states):
+        elif all(branch_constraint for _, branch_constraint in next_states):
             branch_preds = [flatten_label('#And', not_none(rule_predicate)) for _, rule_predicate in next_states]
             common_preds = list(
                 unique(
@@ -283,9 +286,11 @@ def log(message: str, *, warning: bool = False) -> None:
                 )
             constraint_strs = [self.pretty_print(ml_pred_to_bool(bc)) for bc in branches]
             log(f'{len(branches)} branches: {node_id} -> {constraint_strs}')
-            return Branch(branches)
+            extend_results.append(Branch(branches))
         else:
             # NDBranch
             log(f'{len(next_states)} non-deterministic branches: {node_id}')
             next_cterms = [cterm for cterm, _ in next_states]
-            return NDBranch(next_cterms, next_node_logs, self._extract_rule_labels(next_node_logs))
+            extend_results.append(NDBranch(next_cterms, next_node_logs, self._extract_rule_labels(next_node_logs)))
+
+        return extend_results

pyk/src/pyk/proof/reachability.py

@@ -674,6 +674,7 @@ class APRProver(Prover[APRProof, APRProofStep, APRProofResult]):
     direct_subproof_rules: bool
     assume_defined: bool
     kcfg_explore: KCFGExplore
+    next_steps: dict[NodeIdLike, KCFGExtendResult]
 
     def __init__(
         self,
@@ -698,6 +699,7 @@ def __init__(
         self.fast_check_subsumption = fast_check_subsumption
         self.direct_subproof_rules = direct_subproof_rules
         self.assume_defined = assume_defined
+        self.next_steps = {}
 
     def close(self) -> None:
         self.kcfg_explore.cterm_symbolic._kore_client.close()
@@ -793,14 +795,24 @@ def step_proof(self, step: APRProofStep) -> list[APRProofResult]:
         if step.circularity and not step.nonzero_depth and (execute_depth is None or execute_depth > 1):
             execute_depth = 1
 
-        extend_result = self.kcfg_explore.extend_cterm(
-            step.node.cterm,
-            execute_depth=execute_depth,
-            cut_point_rules=cut_rules,
-            terminal_rules=self.terminal_rules,
-            module_name=step.module_name,
-            node_id=step.node.id,
-        )
+        if step.node.id in self.next_steps:
+            extend_results = [self.next_steps.pop(step.node.id)]
+        else:
+            extend_results = self.kcfg_explore.extend_cterm(
+                step.node.cterm,
+                execute_depth=execute_depth,
+                cut_point_rules=cut_rules,
+                terminal_rules=self.terminal_rules,
+                module_name=step.module_name,
+                node_id=step.node.id,
+            )
+
+        assert len(extend_results) == 1 or len(extend_results) == 2
+        extend_result = extend_results[0]
+        if len(extend_results) == 2:
+            assert step.node.id not in self.next_steps
+            self.next_steps[step.node.id] = extend_results[1]
+
         return [
             APRProofExtendResult(
                 node_id=step.node.id, extend_result=extend_result, prior_loops_cache_update=prior_loops