#!/usr/bin/env python
# viewcfg - A script for viewing the CFG of SIL and LLVM IR -*- python -*-
#
# This source file is part of the Swift.org open source project
#
# Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
# Licensed under Apache License v2.0 with Runtime Library Exception
#
# See https://swift.org/LICENSE.txt for license information
# See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
#
# ----------------------------------------------------------------------------
#
# For vim users: use the following lines in .vimrc...
#
#   com! -nargs=? Funccfg silent ?{$?,/^}/w !viewcfg <args>
#   com! -range -nargs=? Viewcfg silent <line1>,<line2>w !viewcfg <args>
#
# ...to add these commands:
#
#   :Funccfg        displays the CFG of the current SIL/LLVM function.
#   :<range>Viewcfg displays the sub-CFG of the selected range.
#
# Note: viewcfg should be in the $PATH and .dot files should be associated
# with the Graphviz app.
#
# ----------------------------------------------------------------------------

from __future__ import print_function

import re
import subprocess
import sys
import tempfile


def help():
    print("""\
Usage:

viewcfg [output-suffix] < file

By default all CFGs are opened in the same window.
Use the a unique output-suffix to open a CFG in a new window.
""")


class Block(object):

    current_index = 0

    def __init__(self, name, preds):
        self.name = name
        self.content = None
        self.preds = []
        self.succs = None
        self.last_line = None
        self.index = Block.current_index
        Block.current_index += 1
        if preds is not None:
            for pred in re.split("[, %]", preds):
                can_pred = pred.strip()
                if can_pred:
                    self.preds.append(can_pred)

    def add_line(self, text):
        if self.content is None:
            self.content = ""
        escaped_text = re.sub(r'([\\<>{}"|])', r'\\\1', text[0:80]).rstrip()
        self.content += escaped_text + '\\l'
        self.last_line = text

    def get_succs(self):
        if self.succs is None:
            self.succs = []
            if self.last_line is not None:
                for match in re.finditer(r'\bbb[0-9]+\b', self.last_line):
                    self.succs.append(match.group())

                for match in re.finditer(r'\blabel %(\S+)\b', self.last_line):
                    self.succs.append(match.group(1))

        return self.succs


def main():
    suffix = ""
    if len(sys.argv) >= 2:
        if sys.argv[1].startswith('-'):
            help()
            return
        suffix = sys.argv[1]

    blocks = {}
    cur_block = None
    sil_block_pattern = re.compile(r'^(\S+)(\(.*\))?: *(\/\/ *Preds:(.*))?$')
    llvm_block_pattern1 = re.compile(r'^(\S+): *; *preds =(.*)?$')
    llvm_block_pattern2 = re.compile(r'^; <label>:(\d+) *; *preds =(.*)?$')

    # Scan the input file.

    for line in sys.stdin:
        sil_block_match = sil_block_pattern.match(line)
        llvm_block_match1 = llvm_block_pattern1.match(line)
        llvm_block_match2 = llvm_block_pattern2.match(line)
        block_name = None
        preds = None
        if sil_block_match:
            block_name = sil_block_match.group(1)
            preds = sil_block_match.group(4)
        elif llvm_block_match1:
            block_name = llvm_block_match1.group(1)
            preds = llvm_block_match1.group(2)
        elif llvm_block_match2:
            block_name = llvm_block_match2.group(1)
            preds = llvm_block_match2.group(2)
        elif line.startswith(' '):
            if cur_block is not None:
                cur_block.add_line(line)
        elif not line[:1].isspace():
            if line.startswith('}') and blocks:
                break
            cur_block = None

        if block_name is not None:
            cur_block = Block(block_name, preds)
            cur_block.add_line(line)
            blocks[block_name] = cur_block

    # Add empty blocks which we didn't see, but which are referenced.
    new_blocks = {}
    for name, block in blocks.iteritems():
        for adj_name in (block.preds + block.get_succs()):
            if adj_name not in blocks:
                new_blocks[adj_name] = Block(adj_name, None)

    blocks = dict(blocks.items() + new_blocks.items())

    # Add missing edges if we didn't see a successor in the terminator
    # but the block is mentioned in the pred list of the successor.

    for name, block in blocks.iteritems():
        for pred_name in block.preds:
            pred_block = blocks[pred_name]
            if name not in pred_block.get_succs():
                pred_block.get_succs().append(name)

    # Write the output dot file.

    file_name = tempfile.gettempdir() + "/viewcfg" + suffix + ".dot"
    with open(file_name, 'w') as out_file:
        out_file.write('digraph "CFG" {\n')
        for name, block in blocks.iteritems():
            if block.content is not None:
                out_file.write(
                    "\tNode" + str(block.index) +
                    " [shape=record,label=\"{" + block.content + "}\"];\n")
            else:
                out_file.write(
                    "\tNode" + str(block.index) +
                    " [shape=record,color=gray,fontcolor=gray,label=\"{" +
                    block.name + "}\"];\n")

            for succ_name in block.get_succs():
                succ_block = blocks[succ_name]
                out_file.write(
                    "\tNode" + str(block.index) + " -> Node" +
                    str(succ_block.index) + ";\n")

        out_file.write("}\n")

    # Open the dot file (should be done with Graphviz).

    subprocess.call(["open", file_name])

if __name__ == '__main__':
    main()