#!/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

def checkbool(name, default=False):
    v = os.getenv(name)
    if v:
        try:
            return bool(int(v))
        except:
            pass
    return default

CCC_LOG = checkenv('CCC_LOG')
CCC_ECHO = checkbool('CCC_ECHO')
CCC_NATIVE = checkbool('CCC_NATIVE','1')
CCC_FALLBACK = checkbool('CCC_FALLBACK')
CCC_LANGUAGES = checkenv('CCC_LANGUAGES','c,c++,c-cpp-output,objective-c,objective-c++,objective-c-cpp-output,assembler-with-cpp')
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 syntaxonly(args):
    command = [CLANG,'-fsyntax-only']
    run(command + args)
    
def compile_fallback(args):
    command = [CC,'-c']
    run(command + args)
    
def compile(args, native, save_temps=False, asm_opts=[]):
    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] + asm_opts)
        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, asm_opts):
    if CCC_LANGUAGES and language and language not in CCC_LANGUAGES:
        log('fallback', args)
        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, asm_opts)
        except:
            log('fallback-on-fail', args)
            print >>sys.stderr, 'WARNING: ccc: Using fallback compiler for: %s'%(' '.join(map(quote, args)),)
            compile_fallback(args)
    else:
        compile(args, native, save_temps, asm_opts)
    
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 ""

def log(name, item):
    if CCC_LOG:
        f = open(CCC_LOG,'a')
        print >>f, (name, item)
        f.close()

def inferaction(args):
    if '-E' in args:
        return 'preprocess'
    if '-fsyntax-only' in args:
        return 'syntax-only'
    if '-c' in args:
        return 'compile'
    for arg in args:
        if arg.startswith('-print-prog-name'):
            return 'pring-prog-name'
    return 'link'

def main(args):
    log('invoke', args)

    action = inferaction(args)
    output = ''
    asm_opts = []
    compile_opts = []
    link_opts = []
    files = []
    save_temps = 0
    language = ''
    native = CCC_NATIVE

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

        if '=' in arg:
            argkey,argvalue = arg.split('=',1)
        else:
            argkey,argvalue = arg,None

        # Modes ccc supports
        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 argkey in ('-std', '-mmacosx-version-min'):
            compile_opts.append(arg)

        # Special case debug options to only pass -g to clang. This is
        # wrong.
        if arg in ('-g', '-gdwarf-2'):
            compile_opts.append('-g')

        # 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 no arguments that should pass through
        if (arg in ('-dynamiclib', '-bundle', '-headerpad_max_install_names',
                    '-nostdlib', '-static', '-dynamic', '-r') or
            arg.startswith('-Wl,')):
            link_opts.append(arg)

        # Options with one argument that should pass through
        if arg in ('-framework', '-multiply_defined', '-bundle_loader',
                   '-weak_framework',
                   '-e', '-install_name',
                   '-unexported_symbols_list', '-exported_symbols_list', 
                   '-compatibility_version', '-current_version', '-init',
                   '-seg1addr', '-dylib_file', '-Xlinker', '-undefined'):
            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])
            asm_opts.append(arg)
            asm_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]
            compile_opts.append(arg)
            compile_opts.append(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
            args = []
            if language:
                args.extend(['-x', language])
            if poutput:
                args += ['-o', poutput, file] + compile_opts
            else:
                args += [file] + compile_opts
            preprocess(args)
            # Discard the explicit language after used once
            language = ''

    if action == 'syntax-only':
        for i, file in enumerate(files):
            if not language:
                language = inferlanguage(extension(file))
            args = []
            if language:
                args.extend(['-x', language])
            args += [file] + compile_opts
            syntaxonly(args)
            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 = []
            if language:
                args.extend(['-x', language])
            args += ['-o', coutput, file] + compile_opts
            checked_compile(args, native, language, save_temps, asm_opts)
            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 = []
                if language:
                    args.extend(['-x', language])
                args = ['-o', out, file] + compile_opts
                checked_compile(args, native, language, save_temps, asm_opts)
                language = ''
                files[i] = out
        if not output:
            output = 'a.out'
        args = ['-o', output] + files + link_opts
        link(args, native)

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