cfg_visualizer.py

from graphviz import Digraph


class CFGVisualizer:
    def __init__(self):
        self.graph = Digraph()
        self.node_id = 0

    def new_node(self, label):
        nid = f"N{self.node_id}"
        self.node_id += 1

        self.graph.node(nid, label)
        return nid

    def build(self, program):
        nodes = {}
        edges = []

        prev = None

        for instr in program:
            nid = self.new_node(instr)
            nodes[instr] = nid

            if prev:
                edges.append((prev, nid))

            prev = nid

            # detect jumps (simple CFG logic)
            if instr.startswith("JMPZ"):
                parts = instr.split()
                target = parts[-1]

                jump_node = self.new_node(f"→ {target}")
                edges.append((nid, jump_node))

        for a, b in edges:
            self.graph.edge(a, b)

        return self.graph

    def render(self, program, filename="cfg"):
        g = self.build(program)
        g.render(filename, format="png", cleanup=True)
        print(f"CFG saved as {filename}.png")