1#!/usr/bin/env python3
2# -- coding: utf-8 --
3# Copyright (c) 2024 Shenzhen Kaihong Digital Industry Development Co., Ltd.
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16# Copyright (c) 2024 Huawei Device Co., Ltd.
17# Licensed under the Apache License, Version 2.0 (the "License");
18# you may not use this file except in compliance with the License.
19# You may obtain a copy of the License at
20#
21# http://www.apache.org/licenses/LICENSE-2.0
22#
23# Unless required by applicable law or agreed to in writing, software
24# distributed under the License is distributed on an "AS IS" BASIS,
25# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
26# See the License for the specific language governing permissions and
27# limitations under the License.
28
29from typing import List
30import graphviz as gv
31import sys
32import os
33import argparse
34
35
36class BasicBlock:
37    def __init__(self, ident: int) -> None:
38        self.id = ident
39        self.insts = []
40        self.preds = []
41        self.succs = []
42        self.props = []
43
44    def text(self, no_insts=False) -> str:
45        s = 'BB ' + str(self.id) + '\n'
46        s += 'props: '
47        for i, prop in enumerate(self.props):
48            if i > 0:
49                s += ', '
50            s += prop
51        s += '\n'
52
53        if not no_insts:
54            for inst in self.insts:
55                s += inst + '\n'
56        return s
57
58
59class Function:
60    def __init__(self, method: str, blocks: List[BasicBlock]) -> None:
61        self.method = method
62        self.blocks = blocks
63
64
65class GraphDumpParser:
66    def __init__(self) -> None:
67        self.method = None
68        self.block = None
69        self.blocks = []
70        self.functions = []
71
72    def finish_block(self):
73        if self.block:
74            self.blocks.append(self.block)
75            self.block = None
76
77    def finish_function(self):
78        self.finish_block()
79        if self.method:
80            self.functions.append(Function(self.method, self.blocks))
81            self.blocks = []
82
83    def begin_block(self, line: str):
84        self.finish_block()
85
86        preds_start = line.find('preds:')
87        block_id = int(line[len('BB'):preds_start].strip())
88        self.block = BasicBlock(block_id)
89        if preds_start != -1:
90            self.parse_block_preds(line, preds_start)
91
92    def begin_function(self, line: str):
93        self.finish_function()
94        self.method = line[len('method:'):].strip()
95
96    def parse_block_preds(self, line: str, preds_start: int):
97        preds = line[preds_start+len('preds:'):].strip().strip('[]')
98        for pred in preds.split(','):
99            pred = pred.strip()
100            self.block.preds.append(int(pred[len('bb'):].strip()))
101
102    def parse_block_props(self, line: str):
103        for s in line[len('prop:'):].strip().split(','):
104            s = s.strip()
105            if not s.startswith('bc'):
106                self.block.props.append(s)
107
108    def parse_block_succs(self, line: str):
109        succs = line[len('succs:'):].strip().strip('[]')
110        for succ in succs.split(','):
111            succ = succ.strip()
112            self.block.succs.append(int(succ[len('bb'):].strip()))
113        self.finish_block()
114
115    def parse(self, lines: List[str]) -> List[Function]:
116        for line in lines:
117            if line.startswith('Method:'):
118                self.begin_function(line)
119            elif line.startswith('BB'):
120                self.begin_block(line)
121            elif self.block:
122                if line.startswith('prop:'):
123                    self.parse_block_props(line)
124                elif line.startswith('succs:'):
125                    self.parse_block_succs(line)
126                else:
127                    self.block.insts.append(line)
128        self.finish_function()
129        return self.functions
130
131
132def draw_function(function: Function, out_dir=None, no_insts=False):
133    dot = gv.Digraph(format='png')
134    for block in function.blocks:
135        dot.node(str(block.id), block.text(no_insts), shape='box')
136    for block in function.blocks:
137        for succ in block.succs:
138            dot.edge(str(block.id), str(succ))
139    basename = 'cfg_' + function.method
140    dotfile_path = os.path.join(out_dir, basename)
141    dot.render(basename, out_dir, format="png")
142    os.rename(dotfile_path, dotfile_path + '.dot')
143
144
145def main():
146    parser = argparse.ArgumentParser(description="A tool for drawing CFGs by reading ir dump from stdin")
147    parser.add_argument("--no-insts", action="store_true", help="drawing without ir instructions")
148    parser.add_argument("--out", type=str, default="./out", help="output directory, default to './out'")
149    args = parser.parse_args()
150
151    functions = GraphDumpParser().parse(sys.stdin.readlines())
152    for function in functions:
153        draw_function(function, args.out, no_insts=args.no_insts)
154
155
156if __name__ == "__main__":
157    main()
158