swift-mirror/utils/viewcfg

#!/usr/bin/env python

# A script for viewing the CFG of SIL and llvm IR.

# 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.

import re
import sys
import tempfile
import subprocess
import os

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 not adjName 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 not name in predBlock.getSuccs():
                predBlock.getSuccs().append(name)

    # Write the output dot file.

    fileName = tempfile.gettempdir() + "/viewcfg" + suffix + ".dot"
    outFile = open(fileName, "w")

    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")
    outFile.flush()
    os.fsync(outFile.fileno())
    outFile.close

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

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

main()