swift-mirror/utils/viewcfg

#!/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 - 2016 Apple Inc. and the Swift project authors
# Licensed under Apache License v2.0 with Runtime Library Exception
#
# See http://swift.org/LICENSE.txt for license information
# See http://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 sys
import tempfile
import subprocess

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:

    currentIndex = 0

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

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

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

                for match in re.finditer(r'\blabel %(\S+)\b', self.lastLine):
                    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 = {}
    curBlock = None
    silBlockPattern = re.compile(r'^(\S+)(\(.*\))?: *(\/\/ *Preds:(.*))?$')
    llvmBlockPattern1 = re.compile(r'^(\S+): *; *preds =(.*)?$')
    llvmBlockPattern2 = re.compile(r'^; <label>:(\d+) *; *preds =(.*)?$')

    # Scan the input file.

    for line in sys.stdin:
        silBlockMatch = silBlockPattern.match(line)
        llvmBlockMatch1 = llvmBlockPattern1.match(line)
        llvmBlockMatch2 = llvmBlockPattern2.match(line)
        blockName = None
        preds = None
        if silBlockMatch:
            blockName = silBlockMatch.group(1)
            preds = silBlockMatch.group(4)
        elif llvmBlockMatch1:
            blockName = llvmBlockMatch1.group(1)
            preds = llvmBlockMatch1.group(2)
        elif llvmBlockMatch2:
            blockName = llvmBlockMatch2.group(1)
            preds = llvmBlockMatch2.group(2)
        elif line.startswith(' '):
            if curBlock is not None:
                curBlock.addLine(line)
        elif not line[:1].isspace():
            if line.startswith('}') and blocks:
                break
            curBlock = None

        if blockName is not None:
            curBlock = Block(blockName, preds)
            curBlock.addLine(line)
            blocks[blockName] = curBlock

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

    blocks = dict(blocks.items() + newBlocks.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 predName in block.preds:
            predBlock = blocks[predName]
            if name not in predBlock.getSuccs():
                predBlock.getSuccs().append(name)

    # Write the output dot file.

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

            for succName in block.getSuccs():
                succBlock = blocks[succName]
                outFile.write("\tNode" + str(block.index) + " -> Node" +
                    str(succBlock.index) + ";\n")

        outFile.write("}\n")

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

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

if __name__ == '__main__':
    main()