From 90b655442eb36e2bcdb28eeedddf94a3b6f1d3e8 Mon Sep 17 00:00:00 2001
From: caandt <caandt@twoha.cc>
Date: Wed, 6 Aug 2025 12:21:37 -0500
Subject: [PATCH] fix cdg

---
 pylingual/control_flow_reconstruction/cfg.py          |  3 +++
 .../control_flow_reconstruction/templates/CDG.py      | 11 +++++------
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/pylingual/control_flow_reconstruction/cfg.py b/pylingual/control_flow_reconstruction/cfg.py
index e5759be..6b406cc 100644
--- a/pylingual/control_flow_reconstruction/cfg.py
+++ b/pylingual/control_flow_reconstruction/cfg.py
@@ -187,6 +187,9 @@ class CFG(DiGraph_CFT):
             self.iteration_graphs[-1].append(dot.to_string())
 
     def cdg(self) -> CFG:
+        self.add_nodes_from([self.start, self.end])
+        self.add_edges_from((self.start, n, EdgeKind.Meta.prop()) for n in self.nodes if self.in_degree(n) == 0)
+        self.add_edges_from((n, self.end, EdgeKind.Meta.prop()) for n in self.nodes if all(e["kind"] is EdgeKind.Exception for _, _, e in self.out_edges(n, data=True)))
         pdt = nx.create_empty_copy(self)
         pdt.add_edges_from((B, A) for A, B in nx.immediate_dominators(self.reverse(), self.end).items())
         pdt.remove_edge(self.end, self.end)
diff --git a/pylingual/control_flow_reconstruction/templates/CDG.py b/pylingual/control_flow_reconstruction/templates/CDG.py
index 5cd96f9..455353a 100644
--- a/pylingual/control_flow_reconstruction/templates/CDG.py
+++ b/pylingual/control_flow_reconstruction/templates/CDG.py
@@ -33,6 +33,9 @@ class CDG(ControlFlowTemplate):
 
         cfg.clear()
         cfg.add_node(cdg)
+        for p, n in nx.dfs_edges(cdg.cdg, cfg.start):
+            cdg.cdg.nodes[n]["indent"] = cdg.cdg.nodes[p].get("indent", -1) + 1
+        cdg.cdg.remove_node(cfg.start)
         return cdg
 
     @override
@@ -41,11 +44,7 @@ class CDG(ControlFlowTemplate):
 
     @override
     def to_indented_source(self, source):
-        cdg = self.cdg
-        for p, n in nx.dfs_edges(cdg, self.start):
-            cdg.nodes[n]["indent"] = cdg.nodes[p].get("indent", -1) + 1
-        cdg.remove_node(self.start)
         src = []
-        for n in sorted(cdg.nodes, key=lambda x: x.offset):
-            src.extend(source[n, cdg.nodes[n].get("indent", 0)])
+        for n in sorted(self.cdg.nodes, key=lambda x: x.offset):
+            src.extend(source[n, self.cdg.nodes[n].get("indent", 0)])
         return src