<commit_before>import html
import os
from typing import List, Union, Dict, Optional, Tuple

from hdlConvertorAst.to.hdlUtils import iter_with_last
from hwtHls.hlsStreamProc.debugCodeSerializer import CopyBasicBlockLabelsToCode
from hwtHls.hlsStreamProc.ssa.analysis.liveness import EdgeLivenessDict
from hwtHls.hlsStreamProc.ssa.basicBlock import SsaBasicBlock
from hwtHls.hlsStreamProc.ssa.phi import SsaPhi
from hwtHls.hlsStreamProc.statements import HlsStreamProcCodeBlock
from hwtHls.netlist.toGraphwiz import GraphwizNode, GraphwizLink, \
    HwtHlsNetlistToGraphwiz


class SsaToGraphwiz():
    """
    Convert SSA to graphwiz for visualization.
    """

    def __init__(self, name: str):
        self.name = name
        self.nodes: List[GraphwizNode] = []
        self.links: List[GraphwizLink] = []
        self.obj_to_node: Dict[Union[SsaBasicBlock, Tuple[SsaBasicBlock, SsaBasicBlock]], GraphwizNode] = {}

    def construct(self, begin: SsaBasicBlock,
                  code: Optional[HlsStreamProcCodeBlock],
                  pipelines: Optional[List[List[SsaBasicBlock]]],
                  edge_var_live: Optional[EdgeLivenessDict]):
        self._node_from_SsaBasicBlock(begin, True, edge_var_live)
        if pipelines is not None:
            bb_to_pipe_i = {}
            for i, pipe in enumerate(pipelines):
                for bb in pipe:
                    bb_to_pipe_i[bb] = i

            for bb, n in self.obj_to_node.items():
                if not isinstance(bb, SsaBasicBlock):
                    continue
                pipe_i = bb_to_pipe_i[bb]
                n.body = n.body.replace("<begin> ", f"<begin> p{pipe_i} ")

        if code is not None:
            CopyBasicBlockLabelsToCode().visit(begin)
            code_body = '[shape=plaintext,fontname=monospace,label="' + self._escape(repr(code)).replace("\n", "\\l\\\n") + '\l"]'
            self.nodes.append(GraphwizNode("code", code_body))

    def _node_from_SsaBasicBlock(self, bb: SsaBasicBlock, is_start: bool, edge_var_live: Optional[EdgeLivenessDict]):
        try:
            return self.obj_to_node[bb]
        except KeyError:
            pass

        # node needs to be constructed before connecting because graph may contain loops
        node = GraphwizNode(f"bb{len(self.nodes)}", None)
        self.obj_to_node[bb] = node
        self.nodes.append(node)

        # construct new node
        top_str = '\<start\> ' if is_start else ''
        body_rows = [f"<begin> {top_str:s}{bb.label:s}:"]
        for phi in bb.phis:
            phi: SsaPhi
            ops = ", ".join(
                f"[{self._escape(repr(o))}, {b.label:s}]"
                for (o, b) in phi.operands
            )
            body_rows.append(f"{self._escape(repr(phi.dst))} = phi {self._escape(repr(phi.dst._dtype))} {ops:s}\\l")

        for stm in bb.body:
            body_rows.append(self._escape(repr(stm)) + "\\l")

        for i, (cond, dst_bb) in enumerate(bb.successors.targets):
            branch_label = f"br{i:d}"
            cond_str = "" if cond is None else self._escape(repr(cond))
            body_rows.append(f"{{\\<{branch_label:s}\\> | <{branch_label:s}> {cond_str:s} }}")
            dst_node = self._node_from_SsaBasicBlock(dst_bb, False, edge_var_live)
            _src = f"{node.label:s}:{branch_label:s}"
            _dst = f"{dst_node.label}:begin"
            if edge_var_live is None:
                self.links.append(GraphwizLink(_src, _dst))
            else:
                var_rows = []
                for var in edge_var_live.get(bb, {}).get(dst_bb, ()):
                    n = html.escape(var._name)
                    var_rows.append(f"<tr><td>{n:s}</td></tr>")
                if var_rows:
                    link_var_node = GraphwizNode(f"bb{len(self.nodes)}", None)
                    link_var_node.body = '[shape=plaintext,fontname=monospace,label=<\n    <table color="gray">%s</table>\n>]' % "\n".join(var_rows)
                    self.obj_to_node[(bb, dst_bb)] = link_var_node
                    self.nodes.append(link_var_node)
                    self.links.append(GraphwizLink(_src, f"{link_var_node.label:s}"))
                    self.links.append(GraphwizLink(f"{link_var_node.label:s}", _dst))
                else:
                    self.links.append(GraphwizLink(_src, _dst))

        buff = []
        buff.append('[shape=record,fontname=monospace,label="{')
        for last, row in iter_with_last(body_rows):
            buff.append(row)
            if not last:
                buff.append("|\n")

        buff.append('}"]')
        node.body = "".join(buff)
        return node

    @staticmethod
    def _escape(s: str) -> str:
        return HwtHlsNetlistToGraphwiz._escape(s)

    def dumps(self):
        buff = ["digraph ", self.name, " {\n", " graph [splines=ortho]\n"]
        for n in self.nodes:
            n.dumps(buff)
        for link in self.links:
            link.dumps(buff)
        buff.append("}\n")
        return "".join(buff)


class HlsNetlistPassToDot():

    def __init__(self, file_name:str):
        self.file_name = file_name

    def apply(self, to_hw: "SsaSegmentToHwPipeline"):
        to_graphwiz = SsaToGraphwiz(os.path.basename(self.file_name))
        with open(self.file_name, "w") as f:
            to_graphwiz.construct(to_hw.start, to_hw.original_code, [to_hw.pipeline, ],<commit_msg>SsaToGraphwiz: fix graph name spec<commit_after>import html
from typing import List, Union, Dict, Optional, Tuple
from pathlib import Path
from hdlConvertorAst.to.hdlUtils import iter_with_last
from hwtHls.hlsStreamProc.debugCodeSerializer import CopyBasicBlockLabelsToCode
from hwtHls.hlsStreamProc.ssa.analysis.liveness import EdgeLivenessDict
from hwtHls.hlsStreamProc.ssa.basicBlock import SsaBasicBlock
from hwtHls.hlsStreamProc.ssa.phi import SsaPhi
from hwtHls.hlsStreamProc.statements import HlsStreamProcCodeBlock
from hwtHls.netlist.toGraphwiz import GraphwizNode, GraphwizLink, \
    HwtHlsNetlistToGraphwiz


class SsaToGraphwiz():
    """
    Convert SSA to graphwiz for visualization.
    """

    def __init__(self, name: str):
        self.name = name
        self.nodes: List[GraphwizNode] = []
        self.links: List[GraphwizLink] = []
        self.obj_to_node: Dict[Union[SsaBasicBlock, Tuple[SsaBasicBlock, SsaBasicBlock]], GraphwizNode] = {}

    def construct(self, begin: SsaBasicBlock,
                  code: Optional[HlsStreamProcCodeBlock],
                  pipelines: Optional[List[List[SsaBasicBlock]]],
                  edge_var_live: Optional[EdgeLivenessDict]):
        self._node_from_SsaBasicBlock(begin, True, edge_var_live)
        if pipelines is not None:
            bb_to_pipe_i = {}
            for i, pipe in enumerate(pipelines):
                for bb in pipe:
                    bb_to_pipe_i[bb] = i

            for bb, n in self.obj_to_node.items():
                if not isinstance(bb, SsaBasicBlock):
                    continue
                pipe_i = bb_to_pipe_i[bb]
                n.body = n.body.replace("<begin> ", f"<begin> p{pipe_i} ")

        if code is not None:
            CopyBasicBlockLabelsToCode().visit(begin)
            code_body = '[shape=plaintext,fontname=monospace,label="' + self._escape(repr(code)).replace("\n", "\\l\\\n") + '\l"]'
            self.nodes.append(GraphwizNode("code", code_body))

    def _node_from_SsaBasicBlock(self, bb: SsaBasicBlock, is_start: bool, edge_var_live: Optional[EdgeLivenessDict]):
        try:
            return self.obj_to_node[bb]
        except KeyError:
            pass

        # node needs to be constructed before connecting because graph may contain loops
        node = GraphwizNode(f"bb{len(self.nodes)}", None)
        self.obj_to_node[bb] = node
        self.nodes.append(node)

        # construct new node
        top_str = '\<start\> ' if is_start else ''
        body_rows = [f"<begin> {top_str:s}{bb.label:s}:"]
        for phi in bb.phis:
            phi: SsaPhi
            ops = ", ".join(
                f"[{self._escape(repr(o))}, {b.label:s}]"
                for (o, b) in phi.operands
            )
            body_rows.append(f"{self._escape(repr(phi.dst))} = phi {self._escape(repr(phi.dst._dtype))} {ops:s}\\l")

        for stm in bb.body:
            body_rows.append(self._escape(repr(stm)) + "\\l")

        for i, (cond, dst_bb) in enumerate(bb.successors.targets):
            branch_label = f"br{i:d}"
            cond_str = "" if cond is None else self._escape(repr(cond))
            body_rows.append(f"{{\\<{branch_label:s}\\> | <{branch_label:s}> {cond_str:s} }}")
            dst_node = self._node_from_SsaBasicBlock(dst_bb, False, edge_var_live)
            _src = f"{node.label:s}:{branch_label:s}"
            _dst = f"{dst_node.label}:begin"
            if edge_var_live is None:
                self.links.append(GraphwizLink(_src, _dst))
            else:
                var_rows = []
                for var in edge_var_live.get(bb, {}).get(dst_bb, ()):
                    n = html.escape(var._name)
                    var_rows.append(f"<tr><td>{n:s}</td></tr>")
                if var_rows:
                    link_var_node = GraphwizNode(f"bb{len(self.nodes)}", None)
                    link_var_node.body = '[shape=plaintext,fontname=monospace,label=<\n    <table color="gray">%s</table>\n>]' % "\n".join(var_rows)
                    self.obj_to_node[(bb, dst_bb)] = link_var_node
                    self.nodes.append(link_var_node)
                    self.links.append(GraphwizLink(_src, f"{link_var_node.label:s}"))
                    self.links.append(GraphwizLink(f"{link_var_node.label:s}", _dst))
                else:
                    self.links.append(GraphwizLink(_src, _dst))

        buff = []
        buff.append('[shape=record,fontname=monospace,label="{')
        for last, row in iter_with_last(body_rows):
            buff.append(row)
            if not last:
                buff.append("|\n")

        buff.append('}"]')
        node.body = "".join(buff)
        return node

    @staticmethod
    def _escape(s: str) -> str:
        return HwtHlsNetlistToGraphwiz._escape(s)

    def dumps(self):
        buff = ["digraph ", self.name, " {\n", " graph [splines=ortho]\n"]
        for n in self.nodes:
            n.dumps(buff)
        for link in self.links:
            link.dumps(buff)
        buff.append("}\n")
        return "".join(buff)


class HlsNetlistPassToDot():

    def __init__(self, file_name:str):
        self.file_name = file_name

    def apply(self, to_hw: "SsaSegmentToHwPipeline"):
        to_graphwiz = SsaToGraphwiz(Path(self.file_name).stem)
        with open(self.file_name, "w") as f:
            to_graphwiz.construct(to_hw.start, to_hw.original_code, [to_hw.pipeline, ],
{'fraction_non_alphanumeric': 0.1259974021154203, 'fraction_numerical': 0.0, 'mean_word_length': 2.291297709923664, 'pattern_counts': {'":': 0, '<': 29, '<?xml version=': 0, '>': 31, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1685749', 'n_tokens_mistral': 3633, 'n_tokens_neox': 3491, 'n_words': 648}