szc

#!/usr/bin/env python

import os
import sys
import random
import argparse
from distutils import util

parser = argparse.ArgumentParser(description="Stabilizer Compiler Driver")

# Which randomizations should be run
parser.add_argument('-R', action='append', choices=['code', 'heap', 'stack', 'link'], default=[])

# Driver control arguments
parser.add_argument('-v', action='store_true')
parser.add_argument('-lang', choices=['c', 'c++', 'fortran'])
parser.add_argument('-platform', choices=['auto', 'linux', 'osx'], default='auto')
parser.add_argument('-frontend', choices=['gcc', 'clang'], default='gcc')

# Compiler pass-through arguments
parser.add_argument('-c', action='store_true')
parser.add_argument('-o')
parser.add_argument('-O', type=int, default=2)
parser.add_argument('-g', action='store_true')
parser.add_argument('-f', action='append', default=[])
parser.add_argument('-D', action='append', default=[])
parser.add_argument('-L', action='append', default=[])
parser.add_argument('-I', action='append', default=[])
parser.add_argument('-l', action='append', default=[])
parser.add_argument('input', nargs='+')

# Do the parse
args = parser.parse_args()

def getPlatform():
	if util.get_platform().startswith('macosx'):
		return 'osx'
	elif util.get_platform().startswith('linux'):
		return 'linux'
	else:
		print 'Unsupported platform'
		exit(2)

def arg(flag, values):
	if not isinstance(values, list):
		values = [values]
	
	cmd = ''
	for v in values:
		if v == True:
			cmd += ' -'+flag
		elif v == False:
			pass
		else:
			cmd += ' -'+flag+v
	return cmd

if args.platform == 'auto':
	args.platform = getPlatform()

STABILIZER_HOME = os.path.dirname(__file__)

if args.platform == 'osx':
	LIBSUFFIX = 'dylib'
	args.frontend = 'clang'
else:
	LIBSUFFIX = 'so'

opts = []

args.l.append('stdc++')
#args.v = True

if 'code' in args.R:
	opts.append('lower-intrinsics')
	opts.append('lowerswitch')
	opts.append('lowerinvoke')
	opts.append('stabilize-code')

if 'stack' in args.R:
	opts.append('stabilize-stack')

if 'heap' in args.R:
	opts.append('stabilize-heap')

if 'code' in args.R or 'heap' in args.R or 'stack' in args.R:
	args.L.append(STABILIZER_HOME)
	args.l.append('stabilizer')
	opts.append('stabilize')

def compile(input):
	if input.endswith('.o'):
		return input
	
	needsAssembly = False
	
	if args.lang == 'fortran':
		cmd = 'gfortran -O0 -fplugin=dragonegg -S -fplugin-arg-dragonegg-emit-ir'
		cmd += arg('o ', args.o+'.s')
		needsAssembly = True
		
	elif args.frontend == 'gcc':
		cmd = 'gcc -O0 -fplugin=dragonegg -S -fplugin-arg-dragonegg-emit-ir'
		cmd += arg('o ', args.o+'.s')
		needsAssembly = True
	
	else:
		cmd = 'clang -O0 -c -emit-llvm'
		cmd += arg('o ', args.o)
	
	cmd += arg('O', 0)
	cmd += arg('g', args.g)
	cmd += arg('I', args.I)
	cmd += arg('f', args.f)
	cmd += arg('D', args.D)
	
	cmd += ' '+input
	
	if args.v:
		print cmd
	os.system(cmd)
	
	if needsAssembly:
		cmd = 'llvm-as -o ' + args.o + ' ' + args.o + '.s'
		if args.v:
			print cmd
		os.system(cmd)
	
	return args.o

def link(inputs):
	cmd = 'llvm-link -o ' + args.o + '.bc '

	if 'link' in args.R:
		random.shuffle(inputs)

	cmd += ' '.join(inputs)

	if args.v:
		print cmd
	os.system(cmd)
	return args.o + '.bc'

def transform(input):
	cmd = 'opt -o ' + args.o + '.opt.bc'
	cmd += ' ' + input
	
	if args.O > 0:
		cmd += ' -O'+str(args.O)
	
	cmd += ' -load='+STABILIZER_HOME+'/LLVMStabilizer.'+LIBSUFFIX
	
	cmd += arg('', opts)
	
	if args.v:
		print cmd
	os.system(cmd)
	return args.o + '.opt.bc'

def codegen(input):
	cmd = 'llc -O0 -relocation-model=pic -disable-fp-elim'
	cmd += ' -o ' + args.o + '.s'
	cmd += ' ' + input
	
	if args.v:
		print cmd
	os.system(cmd)
	
	if args.lang == 'fortran':
		cmd = 'gfortran'
	else:
		cmd = args.frontend
	
	cmd += ' ' + args.o + '.s'
	
	cmd += arg('o ', args.o)
	cmd += arg('f', args.f)
	cmd += arg('L', args.L)
	cmd += arg('l', args.l)
	
	if args.v:
		print cmd
	os.system(cmd)
	
	return args.o

# Build up program arguments
object_files = map(compile, args.input)

if not args.c:
	linked = link(object_files)
	transformed = transform(linked)
	codegen(transformed)