python ccfg2.py && dot -Tpng cfg2.dot -o cfg2.png; https://github.com/Iainmon/CS-261

c_cfg.py

from pycparser import c_parser, c_generator


# -----------------------------
# CFG data structures
# -----------------------------
class BasicBlock:
    def __init__(self, block_id):
        self.id = block_id
        self.stmts = []

    def __repr__(self):
        return f"BasicBlock({self.id})"


class Edge:
    def __init__(self, src, dst, label=None):
        self.src = src
        self.dst = dst
        self.label = label

    def __repr__(self):
        return f"Edge({self.src} -> {self.dst}, label={self.label})"


class CFG:
    def __init__(self):
        self.blocks = []
        self.edges = []
        self.entry = self.new_block()
        self.exit = self.new_block()

    def new_block(self):
        block = BasicBlock(len(self.blocks))
        self.blocks.append(block)
        return block

    def add_edge(self, src, dst, label=None):
        self.edges.append(Edge(src.id, dst.id, label))


# -----------------------------
# CFG builder
# -----------------------------
class CFGBuilder:
    def __init__(self):
        self.cfg = CFG()
        self.break_targets = []
        self.continue_targets = []

    def build_function(self, func_def):
        cur = self.cfg.entry
        end = self.visit_stmt(func_def.body, cur)

        if end is not None:
            self.cfg.add_edge(end, self.cfg.exit)

        return self.cfg

    def visit_stmt(self, node, cur):
        kind = node.__class__.__name__

        if kind == "Compound":
            return self.visit_compound(node, cur)

        if kind == "If":
            return self.visit_if(node, cur)

        if kind in ("While", "For"):
            return self.visit_loop(node, cur)

        if kind == "Return":
            cur.stmts.append(node)
            self.cfg.add_edge(cur, self.cfg.exit, "return")
            return None

        if kind == "Break":
            cur.stmts.append(node)
            self.cfg.add_edge(cur, self.break_targets[-1], "break")
            return None

        if kind == "Continue":
            cur.stmts.append(node)
            self.cfg.add_edge(cur, self.continue_targets[-1], "continue")
            return None

        # Default: assignment, declaration, expression, etc.
        cur.stmts.append(node)
        return cur

    def visit_compound(self, node, cur):
        for stmt in node.block_items or []:
            if cur is None:
                cur = self.cfg.new_block()
            cur = self.visit_stmt(stmt, cur)
        return cur

    def visit_if(self, node, cur):
        # Put the condition inside the current block
        cur.stmts.append(node.cond)

        then_block = self.cfg.new_block()
        else_block = self.cfg.new_block()
        merge_block = self.cfg.new_block()

        self.cfg.add_edge(cur, then_block, "true")
        self.cfg.add_edge(cur, else_block, "false")

        then_end = self.visit_stmt(node.iftrue, then_block)
        if then_end is not None:
            self.cfg.add_edge(then_end, merge_block)

        if node.iffalse is not None:
            else_end = self.visit_stmt(node.iffalse, else_block)
        else:
            else_end = else_block

        if else_end is not None:
            self.cfg.add_edge(else_end, merge_block)

        return merge_block

    def visit_loop(self, node, cur):
        # Handle "for(init; cond; next)"
        if node.__class__.__name__ == "For" and node.init is not None:
            cur.stmts.append(node.init)

        cond_block = self.cfg.new_block()
        body_block = self.cfg.new_block()
        after_block = self.cfg.new_block()

        self.cfg.add_edge(cur, cond_block)

        if getattr(node, "cond", None) is not None:
            cond_block.stmts.append(node.cond)
            self.cfg.add_edge(cond_block, body_block, "true")
            self.cfg.add_edge(cond_block, after_block, "false")
        else:
            # Infinite loop form: for(;;)
            self.cfg.add_edge(cond_block, body_block, "loop")

        self.break_targets.append(after_block)
        self.continue_targets.append(cond_block)

        body_end = self.visit_stmt(node.stmt, body_block)

        self.continue_targets.pop()
        self.break_targets.pop()

        if node.__class__.__name__ == "For" and node.next is not None:
            step_block = self.cfg.new_block()
            if body_end is not None:
                self.cfg.add_edge(body_end, step_block)
            step_block.stmts.append(node.next)
            self.cfg.add_edge(step_block, cond_block)
        else:
            if body_end is not None:
                self.cfg.add_edge(body_end, cond_block)

        return after_block


# -----------------------------
# DOT visualization
# -----------------------------
class CFGDotRenderer:
    def __init__(self):
        self.cgen = c_generator.CGenerator()

    def _escape(self, text):
        return (
            text.replace("\\", "\\\\")
                .replace('"', '\\"')
                .replace("{", "\\{")
                .replace("}", "\\}")
                .replace("\n", " ")
                .replace("<", "\\<")
                .replace(">", "\\>")
        )

    def _stmt_to_text(self, stmt):
        try:
            text = self.cgen.visit(stmt)
        except Exception:
            text = stmt.__class__.__name__
        return self._escape(text)

    def block_label(self, block, cfg, show_block_labels=False):
        if block is cfg.entry:
            title = f"B{block.id} (ENTRY)" # " (ENTRY)"
            # title = f"B{block.id}" # " (ENTRY)"
        elif block is cfg.exit:
            title = f"B{block.id} (EXIT)"
        else:
            title = f"B{block.id}"

        if block.stmts:
            body = "".join(f"{self._stmt_to_text(stmt)}\\l" for stmt in block.stmts) # works

            # body = "".join(f"{self._stmt_to_text(stmt)}\\l" for stmt in block.stmts)
            # body = "".join(f"{self._stmt_to_text(stmt)}" for stmt in block.stmts)
        else:
            body = "\\<empty\\>\\l"
            # body = "<empty>"

        if show_block_labels:
            return f'{{{title}|{body}}}'
        else:
            return f'{{{body}}}'


    def render(self, cfg, show_block_labels=False):
        lines = [
            "digraph CFG {",
            '  rankdir=TB;',
            '  graph [fontname="Courier"];',
            '  node [shape=record, fontname="Courier"];',
            '  edge [fontname="Courier"];',
            ""
        ]

        # Emit nodes
        for block in cfg.blocks:
            label = self.block_label(block, cfg, show_block_labels=show_block_labels)

            attrs = [f'label="{label}"']

            if block is cfg.entry:
                attrs.append('style="filled"')
                attrs.append('fillcolor="lightgreen"')
            elif block is cfg.exit:
                attrs.append('style="filled"')
                attrs.append('fillcolor="lightcoral"')
            else:
                attrs.append('style="filled"')
                attrs.append('fillcolor="lightblue"')


            lines.append(f'  B{block.id} [{", ".join(attrs)}];')

        lines.append("")

        # Emit edges
        for edge in cfg.edges:
            if edge.label:
                label = self._escape(edge.label)
                lines.append(f'  B{edge.src} -> B{edge.dst} [label="{label}"];')
            else:
                lines.append(f'  B{edge.src} -> B{edge.dst};')

        lines.append("}")
        return "\n".join(lines)
    
code = r"""
int f(int x) {
    if (x > 0) {
        return 1;
    } else {
        x = x + 1;
    }

    while (x < 10) {
        x++;
    }

    return x;
}
"""


# code = r"""
# int f(int x) {
#     int y = 0;

#     if (x > 0) {
#         y = 1;
#     } else {
#         y = 2;
#     }

#     while (y < 5) {
#         y = y + 69;
#     }

#     return y;
# }
# """

# void print(int condition);
# void assert(int condition);


code = r"""
int f() {
    int x = 0;
    x = x + 1;
    while (x < 10) {
        if (5 < x) {
            print(x);
        } else {
            assert(x < 6);
        }
        x = x + 1;
    }
    return x;
}
"""

code = r"""
int f(int n) {
    int i = 0;
    int sum = 0;
    int product = 1;

    while (i < n) {
        if (i % 2 == 0) {
            sum = sum + i;
        } else {
            product = product * i;
        }

        if (sum > 50) {
            break;
        }

        i = i + 1;
    }

    if (product > sum) {
        return product;
    } else {
        return sum;
    }
}
"""

parser = c_parser.CParser()
ast = parser.parse(code)

func = ast.ext[0]

builder = CFGBuilder()
cfg = builder.build_function(func)

dot = CFGDotRenderer().render(cfg,show_block_labels=False)

with open("cfg2.dot", "w") as f:
    f.write(dot)

print(dot)