#!/usr/bin/env python
#
#                     The LLVM Compiler Infrastructure
#
# This file is distributed under the University of Illinois Open Source
# License. See LICENSE.TXT for details.
#
##===----------------------------------------------------------------------===##
#
# This script attempts to be a drop-in replacement for gcc.
#
##===----------------------------------------------------------------------===##

import os
import sys
import subprocess

def checkenv(name, alternate=None):
    """checkenv(var, alternate=None) - Return the given environment var,
    or alternate if it is undefined or empty."""
    v = os.getenv(name)
    if v and v.strip():
        return v.strip()
    return alternate

CCC_ECHO = checkenv('CCC_ECHO','1')
CCC_NATIVE = checkenv('CCC_NATIVE')
CCC_FALLBACK = checkenv('CCC_FALLBACK')
CCC_LANGUAGES = checkenv('CCC_LANGUAGES','c,c++,c-cpp-output,objective-c,objective-c++,objective-c-cpp-output')
if CCC_LANGUAGES:
    CCC_LANGUAGES = set([s.strip() for s in CCC_LANGUAGES.split(',')])

# We want to support use as CC or LD, so we need different defines.
CLANG = checkenv('CLANG', 'clang')
LLC = checkenv('LLC', 'llc')
AS = checkenv('AS', 'as')
CC = checkenv('CCC_CC', 'cc')
LD = checkenv('CCC_LD', 'c++')

def error(message):
    print >> sys.stderr, 'ccc: ' + message
    sys.exit(1)

def quote(arg):
    if '"' in arg or ' ' in arg:
        return repr(arg)
    return arg

def stripoutput(args):
    """stripoutput(args) -> (output_name, newargs)
    
    Remove the -o argument from the arg list and return the output
    filename and a new argument list. Assumes there will be at most
    one -o option. If no output argument is found the result is (None,
    args)."""
    for i,a in enumerate(args):
        if a.startswith('-o'):
            if a=='-o':
                if i+1<len(args):
                    return args[i+1],args[:i]+args[i+2:]
            elif a.startswith('-o='):
                opt,arg = a.split('=',1)
                return arg,args[:i]+args[i+1:]
    return None,args

def run(args):
    if CCC_ECHO:
        print ' '.join(map(quote, args))
        sys.stdout.flush()
    code = subprocess.call(args)
    if code > 255:
        code = 1
    if code:
        sys.exit(code)

def remove(path):
    """remove(path) -> bool - Attempt to remove the file at path (if any).

    The result indicates if the remove was successful. A warning is
    printed if there is an error removing the file."""
    if os.path.exists(path):
        try:
            os.remove(path)
        except:
            print >>sys.stderr, 'WARNING: Unable to remove temp "%s"'%(path,)
            return False
    return True

def preprocess(args):
    command = [CLANG,'-E']
    run(command + args)

def compile_fallback(args):
    command = [CC,'-c']
    run(command + args)
    
def compile(args, native, save_temps=False):
    if native:
        output,args = stripoutput(args)
        if not output:
            raise ValueError,'Expected to always have explicit -o in compile()'

        # I prefer suffixing these to changing the extension, which is
        # more likely to overwrite other things. We could of course
        # use temp files.
        bc_output = output + '.bc'
        s_output = output + '.s'
        command = [CLANG,'-emit-llvm-bc']
        try:
            run(command + args + ['-o', bc_output])
            # FIXME: What controls relocation model?
            run([LLC, '-relocation-model=pic', '-f', '-o', s_output, bc_output])
            run([AS, '-o', output, s_output])
        finally:
            if not save_temps:
                remove(bc_output)
                remove(s_output)
    else:
        command = [CLANG,'-emit-llvm-bc']
        run(command + args)

def checked_compile(args, native, language, save_temps):
    if CCC_LANGUAGES and language and language not in CCC_LANGUAGES:
        print >>sys.stderr, 'NOTE: ccc: Using fallback compiler for: %s'%(' '.join(map(quote, args)),)
        compile_fallback(args)
    elif CCC_FALLBACK:
        try:
            compile(args, native, save_temps)
        except:
            print >>sys.stderr, 'WARNING: ccc: Using fallback compiler for: %s'%(' '.join(map(quote, args)),)
            compile_fallback(args)
    else:
        compile(args, native, save_temps)
    
def link(args, native):
    if native:
        run([LD] + args)
    else:
        command = ['llvm-ld', '-native', '-disable-internalize']
        run(command + args)

def extension(path):
    return path.split(".")[-1]

def changeextension(path, newext):
    i = path.rfind('.')
    if i < 0:
        return path
    j = path.rfind('/', 0, i)
    if j < 0:
        return path[:i] + "." + newext
    return path[j+1:i] + "." + newext

def inferlanguage(extension):
    if extension == "c":
        return "c"
    elif extension in ["cpp", "cc"]:
        return "c++"
    elif extension == "i":
        return "c-cpp-output"
    elif extension == "m":
        return "objective-c"
    elif extension == "mm":
        return "objective-c++"
    elif extension == "mi":
        return "objective-c-cpp-output"
    elif extension == "s":
        return "assembler"
    elif extension == "S":
        return "assembler-with-cpp"
    else:
        return "unknown"

def main(args):
    action = 'link'
    output = ''
    compile_opts = []
    link_opts = []
    files = []
    save_temps = 0
    language = ''
    native = CCC_NATIVE

    i = 0
    while i < len(args):
        arg = args[i]

        # Modes ccc supports
        if arg == '-E':
            action = 'preprocess'
        if arg == '-c':
            action = 'compile'
        if arg.startswith('-print-prog-name'):
            action = 'print-prog-name'
        if arg == '-save-temps':
            save_temps = 1
        if arg == '-emit-llvm' or arg == '--emit-llvm':
            native = False

        # Options with no arguments that should pass through
        if arg in ['-v', '-fobjc-gc', '-fobjc-gc-only', '-fnext-runtime',
                   '-fgnu-runtime']:
            compile_opts.append(arg)
            link_opts.append(arg)
        
        # Options with one argument that should be ignored
        if arg in ['--param', '-u']:
            i += 1

        # Preprocessor options with one argument that should be ignored
        if arg in ['-MT', '-MF']:
            i += 1

        # Prefix matches for the compile mode
        if arg[:2] in ['-D', '-I', '-U', '-F']:
            if not arg[2:]:
                arg += args[i+1]
                i += 1
            compile_opts.append(arg)
        if arg[:5] in ['-std=']:
            compile_opts.append(arg)

        # Options with one argument that should pass through to compiler
        if arg in [ '-include', '-idirafter', '-iprefix',
                    '-iquote', '-isystem', '-iwithprefix',
                    '-iwithprefixbefore']:
            compile_opts.append(arg)
            compile_opts.append(args[i+1])
            i += 1

        # Options with one argument that should pass through
        if arg in ('-framework', '-multiply_defined', '-bundle_loader',
                   '-e'):
            link_opts.append(arg)
            link_opts.append(args[i+1])
            i += 1

        # Options with one argument that should pass through to both
        if arg in ['-isysroot', '-arch']:
            compile_opts.append(arg)
            compile_opts.append(args[i+1])
            link_opts.append(arg)
            link_opts.append(args[i+1])
            i += 1
        
        # Options with three arguments that should pass through
        if arg in ('-sectorder',):
            link_opts.extend(args[i:i+4])
            i += 3

        # Prefix matches for the link mode
        if arg[:2] in ['-l', '-L', '-F', '-R']:
            link_opts.append(arg)

        # Enable threads
        if arg == '-pthread':
          link_opts.append('-lpthread')

        # Input files
        if arg == '-filelist':
            f = open(args[i+1])
            for line in f:
                files.append(line.strip())
            f.close()
            i += 1
        if arg == '-x':
            language = args[i+1]
            i += 1
        if arg[0] != '-':
            files.append(arg)

        # Output file
        if arg == '-o':
            output = args[i+1]
            i += 1

        i += 1

    if action == 'print-prog-name':
        # assume we can handle everything
        print sys.argv[0]
        return

    if not files:
        error('no input files')

    if action == 'preprocess' or save_temps:
        for i, file in enumerate(files):
            if not language:
                language = inferlanguage(extension(file))
            if save_temps and action != 'preprocess':
                # Need a temporary output file
                if language == 'c':
                    poutput = changeextension(file, "i");
                elif language == 'objective-c':
                    poutput = changeextension(file, "mi");
                else:
                    poutput = changeextension(file, "tmp." + extension(file))
                files[i] = poutput
            else:
                poutput = output
            if poutput:
                args = ['-x', language, '-o', poutput, file] + compile_opts
            else:
                args = ['-x', language, file] + compile_opts
            preprocess(args)
            # Discard the explicit language after used once
            language = ''

    if action == 'compile' or save_temps:
        for i, file in enumerate(files):
            if not language:
                language = inferlanguage(extension(file))
            if save_temps and action != "compile":
                # Need a temporary output file
                coutput = changeextension(file, "o");
                files[i] = coutput
            elif not output:
                coutput = changeextension(file, "o")
            else:
                coutput = output
            args = ['-x', language, '-o', coutput, file] + compile_opts
            checked_compile(args, native, language, save_temps)
            language = ''

    if action == 'link':
        for i, file in enumerate(files):
            if not language:
                language = inferlanguage(extension(file))
            ext = extension(file)
            if ext != "o" and ext != "a" and ext != "so":
                out = changeextension(file, "o")
                args = ['-x', language, '-o', out, file] + compile_opts
                checked_compile(args, native, language, save_temps)
                language = ''
                files[i] = out
        if not output:
            output = 'a.out'
        args = ['-o', output] + link_opts + files
        link(args, native)

if __name__ == '__main__':
    main(sys.argv[1:])