#!/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 argparse
import re
import subprocess
import sys
import tempfile


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())

                it = re.finditer(r'\blabel %"?([^\s"]+)\b', self.last_line)
                for match in it:
                    self.succs.append(match.group(1))

        return self.succs


def parse_args():
    parser = argparse.ArgumentParser(description="""
    Simple regex based transform of SIL and LLVM-IR into graphviz form.
    """)
    parser.add_argument("output_suffix", nargs='?')
    parser.add_argument("--input_file", type=argparse.FileType('r'),
                        default=sys.stdin)
    parser.add_argument("--renderer", choices=["Graphviz", "dot"],
                        default="Graphviz", help="Default is Graphviz")
    return parser.parse_args()


def main():
    args = parse_args()

    suffix = ""
    if args.output_suffix:
        suffix = args.output_suffix

    block_list = []
    block_map = {}
    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'^(\d+):? *; *preds =(.*)?$')

    # Scan the input file.
    for line in args.input_file.readlines():
        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 block_map:
                break
            cur_block = None

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

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

    block_map = dict(block_map.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 block in block_list:
        for pred_name in block.preds:
            pred_block = block_map[pred_name]
            if block.name not in pred_block.get_succs():
                pred_block.get_succs().append(block.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 block in block_list:
            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 = block_map[succ_name]
                out_file.write(
                    "\tNode" + str(block.index) + " -> Node" +
                    str(succ_block.index) + ";\n")

        out_file.write("}\n")

    # Open the dot file with the appropriate renderer.
    if args.renderer == 'Graphviz':
        subprocess.call(["open", "-a", "Graphviz", file_name])
    elif args.renderer == 'dot':
        pdf = '%s.pdf' % file_name
        subprocess.call(["dot", "-Tpdf", file_name, '-o%s' % pdf])
        subprocess.call(["open", "-a", "Preview", pdf])
    else:
        raise RuntimeError('Unknown renderer?!')


if __name__ == '__main__':
    main()