#!/usr/bin/env python

import argparse
import subprocess
import tempfile
import os
import sys

def createAssemblyFile(tmpDir, file, number, args, polly):
  # Check if file exists
  if not os.path.isfile(file):
    print "File '" + file + "' does not exist. Compilation aborted."
    sys.exit(1)

  # Create an LLVM-IR .ll file
  tmpDir = tmpDir +  "/" + str(number)
  subprocess.call(["mkdir", tmpDir])
  llFile = tmpDir + '/output.ll'
  includeFlags = ['-I' + x for x in args.includes]
  preprocessorFlags = ['-D' + x for x in args.preprocessor]

  if args.progress:
    print "Optimizing file: " + file
    print "- Compiling C to LLVM-IR"

  commandLine = ["clang","-emit-llvm", "-S", file, '-o', llFile,
                 '-std=' + args.standard] + includeFlags + preprocessorFlags \
                 + args.unknown_args

  if args.commands:
    print ' '.join(commandLine)

  exit = subprocess.call(commandLine)

  if exit:
    print 'Translation from C to LLVM-IR failed. This command did not work:'
    print ' '.join(commandLine)
    sys.exit(1)

  # Preprocess the programm such that as many optimization opportunities
  # as possible are exposed.
  preoptFile = tmpDir + '/output.preopt.ll'
  preoptPasses = '-basicaa -mem2reg -simplify-libcalls ' \
                 '-simplifycfg -instcombine -tailcallelim -loop-simplify ' \
                 '-lcssa -loop-rotate -lcssa -loop-unswitch -instcombine ' \
                 '-loop-simplify -lcssa -indvars -loop-deletion -instcombine ' \
                 '-polly-prepare -polly-region-simplify -indvars'

  if args.progress:
    print "- Preoptimizing to expose optimization opportunities"

  commandLine = ['opt', '-load', polly] + preoptPasses.split(' ') \
                + ['-S', llFile, '-o', preoptFile]

  if args.commands:
    print ' '.join(commandLine)

  exit = subprocess.call(commandLine)

  if exit:
    print 'Preoptimizing failed: This command did not work'
    print ' '.join(commandLine)
    sys.exit(1)

  # Run the Polly optimizers
  pollyFile = tmpDir + '/output.polly.ll'

  if args.view or args.viewonly:
    commandLine = ['opt', '-load', polly, '-basicaa'] \
                  + [preoptFile, '-disable-output']

    if args.view:
      commandLine.append('-view-scops')
    else:
      commandLine.append('-view-scops-only')

    if args.commands:
      print ' '.join(commandLine)

    exit = subprocess.call(commandLine)

    if exit:
      print 'Polly optimizations failed: This command did not work'
      print ' '.join(commandLine)
      sys.exit(1)

  if args.fpolly:
    commandLine = ['opt', '-load', polly, '-basicaa'] \
                  + [preoptFile, '-disable-polly-legality']

    if args.pollyimport:
      commandLine.append('-polly-import-jscop')

    if args.fpluto:
      commandLine.append('-polly-optimize')

    if args.foptimize:
      commandLine.append('-polly-optimize-isl')

    if args.fatLeastOnce:
      commandLine.append('-enable-polly-atLeastOnce')

    if args.faligned:
      commandLine.append('-enable-polly-aligned')

    if args.ftile:
      commandLine.append('-enable-pluto-tile')

    if args.pluto_fuse:
      commandLine.append('-pluto-fuse')
      commandLine.append(args.pluto_fuse)

    if args.fparallel:
      commandLine.append('-enable-polly-openmp')

    if args.fvector:
      commandLine.append('-enable-polly-vector')
      commandLine.append('-enable-schedule-prevector')
      if (args.ftile or args.fpluto):
	commandLine.append('-enable-pluto-prevector')

    if args.pollyexport:
      commandLine.append('-polly-export-jscop')

    if args.debug:
      debugCommandLine = commandLine + ['-polly-cloog', '-analyze', '-q']
      subprocess.call(debugCommandLine)

    if args.progress:
      print "- Running polly",

      if args.fpluto:
        print "with pluto",

      if args.ftile:
        print "and tiling",

      if args.fparallel:
        print "and parallel",

      print ""

    commandLine = commandLine + ['-polly-codegen', '-S', '-o', pollyFile]

    if args.commands:
      print ' '.join(commandLine)

    exit = subprocess.call(commandLine)

    if exit:
      print 'Polly optimizations failed: This command did not work'
      print ' '.join(commandLine)
      sys.exit(1)


  # Run optimizations to clean up the redundant calculations Polly introduced
  # and to take advantage of non polyhedral optimizations.
  optFile = tmpDir + '/output.opt.ll'
  commandLine = ['opt', '-S', '-o', optFile + "tmp"]
  commandLine2 = ['opt', '-S', '-o', optFile]

  if args.optlevel != '0':
    commandLine.append('-O' + args.optlevel)
    commandLine2.append('-O' + args.optlevel)

  if args.fpolly:
    commandLine.append(pollyFile)
  else:
    commandLine.append(preoptFile)

  commandLine2.append(optFile + "tmp");

  if args.progress:
    print "- Running normal LLVM -O3 passes"

  exit = subprocess.call(commandLine)

  if exit:
    print 'LLVM -O3 optimizations failed: This command did not work'
    print ' '.join(commandLine)
    sys.exit(1)

  exit = subprocess.call(commandLine2)

  if exit:
    print 'LLVM -O3 optimizations failed: This command did not work'
    print ' '.join(commandLine)
    sys.exit(1)


  if args.emitllvm:
    return optFile


  # Create assembly file
  assemblyFile = tmpDir + '/output.opt.s'
  commandLine = ['llc', optFile, '-o', assemblyFile]

  if args.progress:
    print "- Translating LLVM-IR to assembly"

  subprocess.call(commandLine)

  if exit:
    print 'Translating LLVM-IR to assembly failed: This command did not work'
    print ' '.join(commandLine)
    sys.exit(1)

  return assemblyFile

def parseArguments():
  description = 'pollycc is a simple replacement for compiler drivers like ' \
                'gcc, clang or icc. It uses clang to compile C code and can ' \
                'optimize the code with Polly. It will either produce an ' \
                'optimized binary or an optimized \'.o\' file.'

  parser = argparse.ArgumentParser(description=description)
  parser.add_argument('files', nargs='+')
  parser.add_argument('-o', dest='output', help='the name of the output file')
  parser.add_argument('-I', action='append', dest='includes', default=[],
                      help='include path to pass to clang')
  parser.add_argument('-D', action='append', dest='preprocessor', default=[],
                      help='preprocessor directives to pass to clang')
  parser.add_argument('-l', action='append', dest='libraries', default=[],
                      help='library flags to pass to the linker')
  parser.add_argument('-L', action='append', dest='librarypath', default=[],
                      help='library paths to pass to the linker')
  parser.add_argument('-O', dest='optlevel', choices='0123s', default='3',
                      help='optimization level.')
  parser.add_argument('-S', action='store_true',
                      help='compile only; do not link or assemble')
  parser.add_argument('-emit-llvm', action='store_true', dest='emitllvm',
                      help='output LLVM-IR instead of assembly if -S is set')
  parser.add_argument('-std', dest='standard', help='The C standard to use',
                      default='gnu89')
  parser.add_argument('-p', '--progress', action='store_true',
                      help='Show the compilation progress')
  parser.add_argument('-c', action='store_true',
                      help='compile and assemble, but do not link')
  parser.add_argument('-fpolly', help='enable polly', action='store_true')
  parser.add_argument('-fpluto', help='enable pluto', action='store_true')
  parser.add_argument('-fpluto-fuse', dest='pluto_fuse', help='enable pluto')
  parser.add_argument('-foptimize',
                      help='Optimize schedule with isl pluto like algorithm',
                      action='store_true')
  parser.add_argument('-faligned', help='Assume aligned vector accesses',
                      action='store_true')
  parser.add_argument('-fatLeastOnce',
                      help='Assume all loops are executed at least once',
                      action='store_true')
  parser.add_argument('-fview-scops', dest='view',
                      help='Show the scops with graphviz',
                      action='store_true')
  parser.add_argument('-fview-scops-only', dest='viewonly',
                      help='Show the scops with graphviz (Only Basic Blocks)',
                      action='store_true')
  parser.add_argument('-ftile', '-fpluto-tile',
                      help='enable pluto tiling', action='store_true')
  parser.add_argument('-fparallel',
                      help='enable openmp code generation (in development)',
		      action='store_true')
  parser.add_argument('-fvector',
                      help='enable SIMD code generation (in development)',
		      action='store_true')
  parser.add_argument('-fpolly-export', dest='pollyexport',
                      help='Export Polly jscop', action='store_true')
  parser.add_argument('-fpolly-import', dest='pollyimport',
                      help='Import Polly jscop', action='store_true')
  parser.add_argument('-commands', help='print command lines executed',
                      action='store_true')
  parser.add_argument('-d', '--debug', help='print debugging output',
                      action='store_true')
  parser.add_argument('-v', '--version', dest='version', action='store_true',
                      help='print version info')

  arguments, argv = parser.parse_known_args()

  if argv:
    arguments.unknown_args = filter(lambda x: x[0] == '-', argv)
    arguments.files = arguments.files + filter(lambda x: x[0] != '-', argv)

    if arguments.unknown_args:
      print 'warning: The following arguments are unknown and will \n' \
            '         be passed directly to clang: ' + \
            ' '.join(arguments.unknown_args)
  else:
    arguments.unknown_args = []

  # Post process arguments
  if arguments.ftile:
    arguments.fpluto = True

  if arguments.pollyimport:
    arguments.fpolly = True

  if arguments.view or arguments.viewonly:
    arguments.fpolly = True

  if arguments.fpluto:
    arguments.fpolly = True

  if arguments.foptimize:
    arguments.fpolly = True

  return arguments

def createAssemblyFiles(files, args, pollyLib):
  assemblyFiles = []

  tmpDir = tempfile.mkdtemp(prefix='pollycc-');

  number = 0
  for file in args.files:
    assemblyFiles.append(createAssemblyFile(tmpDir, file, number, args, pollyLib))
    if args.progress:
      print ""

    number = number + 1

  return assemblyFiles

def createOutputFiles(assemblyFiles, args):
  if args.output:
    outputFile = args.output
  elif len(args.files) == 1 and args.c:
    outputFile = args.files[0].replace('.c', '.o')
  elif len(args.files) == 1 and args.S:
    outputFile = args.files[0].replace('.c', '.s')
  else:
    outputFile = 'a.out'

  linkerFlags = ['-l' + x for x in args.libraries]
  libraryPath = ['-L' + x for x in args.librarypath]

  commandLine = ['gcc', '-o', outputFile] + assemblyFiles

  if args.S:
    subprocess.call(['mv', assemblyFiles[0], outputFile])
    return

  if args.c:
    commandLine.append('-c')
  else:
    commandLine = commandLine + linkerFlags + libraryPath

  if args.fparallel:
    commandLine.append('-lgomp')

  if args.progress:
    print "Final linking"

  subprocess.call(commandLine)

def checkExecutables(pollyLib, pluto):
  commandLine = ['opt', '-load', pollyLib, '-help']
  try:
    proc = subprocess.Popen(commandLine, stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE)
    stdout_value = proc.communicate()[0]

    if not stdout_value.count('polly-prepare'):
      sys.exit('Polly support not available in opt')

    if pluto and not stdout_value.count('Optimize the scop using pocc'):
      sys.exit('Polly compiled without POCC/Pluto support. -ftile and -fpluto'
	       ' will not work.')
  except OSError:
    print 'error: opt cannot be executed: '
    print 'failing command: \n' + " ".join(commandLine)
    sys.exit(1)

  commandLine = ['clang', '-v']
  try:
    subprocess.call(commandLine, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  except OSError:
    print 'error: clang cannot be executed: '
    print 'failing command: \n' + " ".join(commandLine)
    sys.exit(1)

  commandLine = ['lli', '-help']
  try:
    subprocess.call(commandLine, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  except OSError:
    print 'error: lli cannot be executed: '
    print 'failing command: \n' + " ".join(commandLine)
    sys.exit(1)

def main():
  args = parseArguments()

  if not 'LIBPOLLY' in os.environ:
    sys.exit('Polly library not provided. Please set LIBPOLLY environment ' + \
             'variable to the LLVMPolly.so file')

  pollyLib = os.environ['LIBPOLLY']

  checkExecutables(pollyLib, args.fpluto)

  assemblyFiles = createAssemblyFiles(args.files, args, pollyLib)
  createOutputFiles(assemblyFiles, args)

if __name__ == '__main__':
  main()