diff --git a/vyper/venom/context.py b/vyper/venom/context.py index 2e35967dfe..222886b518 100644 --- a/vyper/venom/context.py +++ b/vyper/venom/context.py @@ -53,6 +53,14 @@ def append_data(self, opcode: str, args: list[IROperand]) -> None: """ self.data_segment.append(IRInstruction(opcode, args)) # type: ignore + def as_graph(self) -> str: + s = ["digraph G {"] + for fn in self.functions.values(): + s.append(fn.as_graph(True)) + s.append("\n") + s.append("}") + return "\n".join(s) + def __repr__(self) -> str: s = ["IRContext:"] for fn in self.functions.values(): diff --git a/vyper/venom/function.py b/vyper/venom/function.py index eace17af0d..fb0dabc99a 100644 --- a/vyper/venom/function.py +++ b/vyper/venom/function.py @@ -214,7 +214,11 @@ def copy(self): new.last_variable = self.last_variable return new - def as_graph(self) -> str: + def as_graph(self, only_subgraph=False) -> str: + """ + Return the function as a graphviz dot string. If only_subgraph is True, only return the + subgraph, not the full digraph -for embedding in a larger graph- + """ import html def _make_label(bb): @@ -227,18 +231,25 @@ def _make_label(bb): return ret # return f"{bb.label.value}:\n" + "\n".join([f" {inst}" for inst in bb.instructions]) - ret = "digraph G {\n" + ret = [] + + if not only_subgraph: + ret.append("digraph G {{") + ret.append(f'subgraph "{self.name}" {{') for bb in self.get_basic_blocks(): for out_bb in bb.cfg_out: - ret += f' "{bb.label.value}" -> "{out_bb.label.value}"\n' + ret.append(f' "{bb.label.value}" -> "{out_bb.label.value}"') for bb in self.get_basic_blocks(): - ret += f' "{bb.label.value}" [shape=plaintext, ' - ret += f'label={_make_label(bb)}, fontname="Courier" fontsize="8"]\n' + ret.append(f' "{bb.label.value}" [shape=plaintext, ') + ret.append(f'label={_make_label(bb)}, fontname="Courier" fontsize="8"]') + + ret.append("}\n") + if not only_subgraph: + ret.append("}\n") - ret += "}\n" - return ret + return "\n".join(ret) def __repr__(self) -> str: str = f"IRFunction: {self.name}\n"