refactor Edge for merging node

runtimeverification · Jul 23, 2024 · c00d6d7 · c00d6d7
1 parent 5246edf
commit c00d6d7
Showing 1 changed file with 30 additions and 10 deletions.
diff --git a/pyk/src/pyk/kcfg/kcfg.py b/pyk/src/pyk/kcfg/kcfg.py
@@ -207,20 +207,36 @@ def targets(self) -> tuple[KCFG.Node, ...]:
     class Edge(EdgeLike):
         source: KCFG.Node
         target: KCFG.Node
-        depth: int
-        rules: tuple[str, ...]
+        depths: tuple[int, ...]
+        rules_list: tuple[tuple[str, ...], ...]
+        csubsts: tuple[CSubst, ...]
+
+        @property
+        def depth(self) -> int:
+            return max(self.depths)
+
+        @property
+        def rules(self) -> tuple[str, ...]:
+            return single(self.rules_list)
 
         def to_dict(self) -> dict[str, Any]:
             return {
                 'source': self.source.id,
                 'target': self.target.id,
-                'depth': self.depth,
-                'rules': list(self.rules),
+                'depths': list(self.depths),
+                'rules_list': [list(rules) for rules in self.rules_list],
+                'csubsts': [csubst.to_dict() for csubst in self.csubsts],
             }
 
         @staticmethod
         def from_dict(dct: dict[str, Any], nodes: Mapping[int, KCFG.Node]) -> KCFG.Edge:
-            return KCFG.Edge(nodes[dct['source']], nodes[dct['target']], dct['depth'], tuple(dct['rules']))
+            return KCFG.Edge(
+                nodes[dct['source']],
+                nodes[dct['target']],
+                tuple(dct['depths']),
+                tuple(tuple(rules) for rules in dct['rules_list']),
+                tuple(CSubst.from_dict(csubst) for csubst in dct['csubsts']),
+            )
 
         def to_rule(self, label: str, claim: bool = False, priority: int | None = None) -> KRuleLike:
             def is_ceil_condition(kast: KInner) -> bool:
@@ -243,11 +259,11 @@ def _simplify_config(config: KInner) -> KInner:
 
         def replace_source(self, node: KCFG.Node) -> KCFG.Edge:
             assert node.id == self.source.id
-            return KCFG.Edge(node, self.target, self.depth, self.rules)
+            return KCFG.Edge(node, self.target, self.depths, self.rules_list, self.csubsts)
 
         def replace_target(self, node: KCFG.Node) -> KCFG.Edge:
             assert node.id == self.target.id
-            return KCFG.Edge(self.source, node, self.depth, self.rules)
+            return KCFG.Edge(self.source, node, self.depths, self.rules_list, self.csubsts)
 
     @final
     @dataclass(frozen=True)
@@ -375,7 +391,7 @@ def with_single_target(self, target: KCFG.Node) -> KCFG.NDBranch:
 
         @property
         def edges(self) -> tuple[KCFG.Edge, ...]:
-            return tuple(KCFG.Edge(self.source, target, 1, ()) for target in self.targets)
+            return tuple(KCFG.Edge(self.source, target, (1,), (), ()) for target in self.targets)
 
         def replace_source(self, node: KCFG.Node) -> KCFG.NDBranch:
             assert node.id == self.source.id
@@ -798,12 +814,16 @@ def contains_edge(self, edge: Edge) -> bool:
             return edge == other
         return False
 
-    def create_edge(self, source_id: NodeIdLike, target_id: NodeIdLike, depth: int, rules: Iterable[str] = ()) -> Edge:
+    def create_edge(self, source_id: NodeIdLike, target_id: NodeIdLike,
+                    depth: int | Iterable[int],
+                    rules: Iterable[str] | Iterable[Iterable[str]] = (),
+                    csubsts: Iterable[CSubst] = ()
+                    ) -> Edge:
         if depth <= 0:
             raise ValueError(f'Cannot build KCFG Edge with non-positive depth: {depth}')
         source = self.node(source_id)
         target = self.node(target_id)
-        edge = KCFG.Edge(source, target, depth, tuple(rules))
+        edge = KCFG.Edge(source, target, (depth, ), tuple(rules), )
         self.add_successor(edge)
         return edge