From bf59d173ade41a5f4630e59751028dfa42fa2a91 Mon Sep 17 00:00:00 2001
From: Stevengre <zhaojianhong96@gmail.com>
Date: Fri, 22 Nov 2024 13:25:58 +0800
Subject: [PATCH] hold the nodes' information in the MergedEdge

---
 pyk/src/pyk/kcfg/kcfg.py        | 25 ++++++++++++++++++++-----
 pyk/src/tests/unit/test_kcfg.py | 12 +++++++++++-
 2 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/pyk/src/pyk/kcfg/kcfg.py b/pyk/src/pyk/kcfg/kcfg.py
index e91db28000..c6a0a6f0dd 100644
--- a/pyk/src/pyk/kcfg/kcfg.py
+++ b/pyk/src/pyk/kcfg/kcfg.py
@@ -264,21 +264,36 @@ class MergedEdge(EdgeLike):
         edges: tuple[KCFG.Edge, ...]
 
         def to_dict(self) -> dict[str, Any]:
+
+            def _merged_edge_to_dict(edge: KCFG.Edge) -> dict[str, Any]:
+                return {
+                    'source': edge.source.to_dict(),
+                    'target': edge.target.to_dict(),
+                    'depth': edge.depth,
+                    'rules': list(edge.rules),
+                }
+
             return {
                 'source': self.source.id,
                 'target': self.target.id,
-                'edges': [edge.to_dict() for edge in self.edges],
+                'edges': [_merged_edge_to_dict(e) for e in self.edges],
             }
 
         @staticmethod
         def from_dict(dct: dict[str, Any], nodes: Mapping[int, KCFG.Node]) -> KCFG.Successor:
+
+            def _merged_edge_from_dict(dct: dict[str, Any]) -> KCFG.Edge:
+                return KCFG.Edge(
+                    KCFG.Node.from_dict(dct['source']),
+                    KCFG.Node.from_dict(dct['target']),
+                    dct['depth'],
+                    tuple(dct['rules']),
+                )
+
             return KCFG.MergedEdge(
                 nodes[dct['source']],
                 nodes[dct['target']],
-                tuple(
-                    KCFG.Edge(edge['source'], edge['target'], edge['depth'], tuple(edge['rules']))
-                    for edge in dct['edges']
-                ),
+                tuple(_merged_edge_from_dict(e) for e in dct['edges']),
             )
 
         def replace_source(self, node: KCFG.Node) -> KCFG.Successor:
diff --git a/pyk/src/tests/unit/test_kcfg.py b/pyk/src/tests/unit/test_kcfg.py
index 254073ba3a..e55f560b14 100644
--- a/pyk/src/tests/unit/test_kcfg.py
+++ b/pyk/src/tests/unit/test_kcfg.py
@@ -113,9 +113,19 @@ def _make_edge_dict(i: int, j: int, depth: int = 1, rules: tuple[str, ...] = ())
 
 
 def merged_edge_dicts(*merged_edges: Iterable) -> list[dict[str, Any]]:
+    def _make_edge_dicts(*edges: Iterable) -> list[dict[str, Any]]:
+        def _make_edge_dict(i: int, j: int, depth: int = 1, rules: tuple[str, ...] = ()) -> dict[str, Any]:
+            return {
+                'source': node_dicts(1, start=i)[0],
+                'target': node_dicts(1, start=j)[0],
+                'depth': depth,
+                'rules': list(rules),
+            }
+
+        return [_make_edge_dict(*edge) for edge in edges]
 
     def _make_merged_edge_dict(s: int, t: int, *edges: Iterable) -> dict[str, Any]:
-        return {'source': s, 'target': t, 'edges': edge_dicts(*edges)}
+        return {'source': s, 'target': t, 'edges': _make_edge_dicts(*edges)}
 
     return [_make_merged_edge_dict(*merged_edge) for merged_edge in merged_edges]