sbt-instr-c

#!/usr/bin/env python

import sys
import json
import os
import getopt
from subprocess import call, Popen, PIPE, STDOUT

sbtinstr='./bin/sbt-instr'

# ---------------------
# Functions from Symbiotic
# ---------------------

usage_msg = """
Usage: ./instr OPTS config.json source.c
and OPTS can be following:
    --output=FILE   Specify output file
    --help          Show help message
    --bc            Given file is a bytecode
    --ll            Generate .ll file from output .bc file.
"""


COLORS = {
    'RED'   : '\033[1;31m',
    'GREEN' : '\033[1;32m',
    'BROWN' : '\033[0;33m',
    'RESET' : '\033[0m'
}

def print_stream(msg, stream, prefix = None, print_nl = True, color = None):
    """
    Print message to stderr/stdout
    @ msg      : str    message to print
    @ prefix   : str    prefix for the message
    @ print_nl : bool  print new line after the message
    @ color    : str    color to use when printing, default None
    """

    # Do not print color when the output is redirected
    # to a file
    if not stream.isatty():
        color = None

    if not color is None:
        stream.write(COLORS[color])

    if msg == '':
        return
    if not prefix is None:
        stream.write(prefix)

    stream.write(msg)

    if not color is None:
        stream.write(COLORS['RESET'])

    if print_nl:
        stream.write('\n')

    stream.flush()

def print_stderr(msg, prefix = None, print_nl = True, color = None):
    print_stream(msg, sys.stderr, prefix, print_nl, color)

def print_stdout(msg, prefix = None, print_nl = True, color = None):
    print_stream(msg, sys.stdout, prefix, print_nl, color)

def err(msg, color = 'RED'):
    print_stderr(msg, 'ERROR: ', color=color)
    sys.exit(1)

def dbg(msg, domain = 'all', print_nl = True):
        global debug_opts
        if domain == 'all' or domain in debug_opts:
            print_stderr(msg, 'DBG: ', print_nl)

# --------------------
# Instrumentation part
# --------------------

final_output = "out.bc"
generate_ll = False
source_is_bc = False
configFile = None
cFile = None
fileToLink = None

def parse_command_line():
    try:
        opts, args = getopt.getopt(sys.argv[1:], '',
                                   ['help','bc','ll','output='])
    except getopt.GetoptError as e:
        err('{0}'.format(str(e)))

    for opt, arg in opts:
        if opt == '--help':
            print(usage_msg)
            sys.exit(0)
        elif opt == '--output':
            global final_output
            final_output = arg
            dbg('Output will be stored to {0}'.format(arg))
        elif opt == '--bc':
            global source_is_bc
            source_is_bc = True
            dbg('Given code is bytecode.')
        elif opt == '--ll':
            global generate_ll
            generate_ll = True
            dbg('Will generate .ll file.')

    return args

def parseJson():
    with open(configFile) as data_file:
        data = json.load(data_file)

    return data["file"]

def run():
    # Convert .c file from config.json
    if not fileToLink.endswith(".bc"):
        bcConfigFile = "config.bc"
        ret = call(['clang -g -c -emit-llvm -o {0} {1}'.format(bcConfigFile, fileToLink)], shell=True)
        if ret != 0:
            err('Converting config source file to .bc failed!')
            return
        else:
            dbg('Config source file converted to .bc.')
    else:
        bcConfigFile = fileToLink

    bcSourceFile = "source.bc"

    if not source_is_bc:
        ret = call(['clang -g -c -emit-llvm -o {0} {1}'.format(bcSourceFile, cFile)], shell=True)
        if ret != 0:
            err('Converting source file to .bc failed!')
            return
        else:
            dbg('Source file converted to .bc.')

    ret = call(['{0} {1} {2} {3} {4}'.format(sbtinstr, configFile, bcSourceFile, bcConfigFile, final_output)],
               shell=True)
    if ret != 0:
        err('Instrumentation failed!')
        return
    else:
        dbg('Instrumentation done.')

#     ret = call(['llvm-link {0} {1} -o {2}'.format(final_output, bcConfigFile, final_output)], shell=True)
#     if ret != 0:
#         err('Linking instrumentation functions failed!')
#         return
#     else:
#         dbg('Linked instrumentation functions to the output file.')

    if(generate_ll):
       ret = call(['llvm-dis {0}'.format(final_output)], shell=True)
       if ret != 0:
           err('Generating .ll file failed!')
           return
       else:
           dbg('.ll file generated')

if __name__ == "__main__":

    args = parse_command_line()
    if len(args) < 1:
        print(usage_msg)
        sys.exit(1)

    configFile = args[0]
    cFile = args[1]
    fileToLink = os.path.join(os.path.dirname(configFile), parseJson())
    run()

    sys.stderr.flush()
    sys.stdout.flush()

    sys.exit(0)