#!/usr/bin/env python3

# This is a helper script for importing the things we use from LLVM.
#
# NOTE: We have slightly modified copies of some of the files, to reduce the
# dependency surface area, so if you run this to update the sources, please
# examine the diffs to remove things that don't make sense and update 'llvm-patch.diff'.

import errno
import optparse
import os
import re
import shutil
import sys
import subprocess

# ADT types.
ADT_imports = ['EpochTracker', 'iterator', 'iterator_range', 'Hashing', 'None', 'Optional',
               'Statistic', 'StringExtras', 'STLExtras', 'AllocatorList', 'Triple']

# ADT basic data structures.
ADT_imports += ['APFloat', 'APInt', 'APSInt', 'ArrayRef', 'bit', 'PointerIntPair', 'SetVector', 'StringRef', 'StringSwitch',
                'Twine', 'IntrusiveRefCntPtr', 'ilist', 'ilist_base', 'ilist_node', 'ilist_node_base',
                'ilist_node_options', 'ilist_iterator', 'simple_ilist', 'OptionSet', 'PointerUnion']

# ADT Mapping structures.
ADT_imports += ['DenseMap', 'DenseMapInfo', 'DenseSet', 'FoldingSet', 'StringMap', 'StringSet']

# ADT "Small" structures.
ADT_imports += ['SmallPtrSet', 'SmallSet', 'SmallString', 'SmallVector']

# ADT Algorithms.
ADT_imports += ['edit_distance']

# Support types and infrastructure.
Support_imports = [
    'AlignOf', 'Allocator', 'Atomic', 'CBindingWrapping', 'Casting', 'Capacity', 'CommandLine', 'Compiler',
    'Endian', 'Errno', 'ErrorHandling', 'Errc', 'ErrorOr', 'Error', 'Format', 'FormatAdapters',
    'ManagedStatic', 'MathExtras', 'Mutex', 'MutexGuard', 'Memory',
    'MemoryBuffer', 'PointerLikeTypeTraits', 'Recycler', 'SwapByteOrder',
    'Timer', 'TimeValue', 'Threading', 'Unicode', 'UniqueLock', 'Unix', 'WindowsError',
    'Valgrind', 'circular_raw_ostream', 'raw_ostream', 'Signposts', 'type_traits', 'JSON', 'WindowsSupport']

# Stuff we don't want, but have to pull in.
Support_imports += [
    'COFF', 'ConvertUTF', 'ConvertUTFWrapper', 'Debug', 'FileSystem',
    'FileUtilities', 'Host', 'Locale', 'MachO', 'Path', 'Process', 'Program', 'SMLoc',
    'SourceMgr', 'Signals', 'StringSaver', 'ToolOutputFile', 'TrailingObjects', 'Unicode', 'UnicodeCharRanges',
    'MemAlloc', 'Chrono', 'FormatProviders', 'FormatVariadic', 'FormatCommon',
    'FormatVariadicDetails', 'NativeFormatting', 'DJB', 'ReverseIteration', 'MD5',
    'SmallVectorMemoryBuffer', 'WithColor', 'Options', 'PrettyStackTrace', 'Watchdog',
    'TargetParser', 'ARMBuildAttributes', 'ARMTargetParser', 'AArch64TargetParser',
    'ARMTargetParser.def', 'AArch64TargetParser.def', 'X86TargetParser.def', 'LineIterator', 'VersionTuple']

# Dependencies from llvm-c needed by Support.
C_imports = ['Types', 'DataTypes', 'Support', 'Error', 'ErrorHandling']

# Support data structures.
Support_imports += ['YAMLParser', 'YAMLTraits']

# Source files to exclude.
Support_source_excludes = set([])

llvm_srcroot = None
indexstoredb_srcroot = None

def note(msg):
    msg = msg.replace(llvm_srcroot, "<LLVM>")
    msg = msg.replace(indexstoredb_srcroot, "<INDEXSTOREDB>")
    print >>sys.stderr, "note: %s" % (msg,)

def mkdir_p(path):
    try:
        os.makedirs(path)
        note("mkdir -p %r" % (path,))
    except OSError as e:
        if e.errno != errno.EEXIST:
            raise

def copyfile(src, dst):
    note("cp %r %r" % (src, dst))
    shutil.copyfile(src, dst)

def includes_prefix(source, include_dir):
    defs = subprocess.check_output([
        'clang', '-std=c++14', '-dM', '-E',
        '-I', include_dir,
        '-x', 'c++', source])

    return "INDEXSTOREDB_PREFIX_H" in defs

def is_llvm_leaf_source(source):
    # AIXDataTypes is a false negative.
    return subprocess.call("grep '# *include *\"llvm' %s | grep --quiet -v AIXDataTypes" % source, shell=True)

def add_prefix(str):
    snippet = '#include "llvm/Config/indexstoredb-prefix.h"\n'

    m = re.search(r'#ifndef *([A-Za-z_]*)\n#define *\1 *\n', str)
    if m:
        return str[:m.end()] + '\n' + snippet + str[m.end():]

    i = str.find(r'#include')
    if i != -1:
        return str[:i] + snippet + str[i:]

    # If we ever hit this, we could find the end of the copyright header comment.
    print >>sys.stderr, 'error: could not find existing #include or header guards'
    sys.exit(1)

def maybe_add_prefix(source, include_dir):
    ext = os.path.splitext(source)[1]
    if not ext in ['.h', '.cpp', '.c']:
        return
    if not is_llvm_leaf_source(source):
        # Skip files that include other llvm headers; we will add the include to
        # the leaf header.
        return
    if includes_prefix(source, include_dir):
        # Skip files that already include the prefix header.
        return

    note("adding prefix header to %r" % source)
    with open(source, 'r+') as file:
        new = add_prefix(file.read())
        file.seek(0)
        file.write(new)
        file.truncate()

def main():
    parser = optparse.OptionParser("usage: %prog <llvm-source-path>")
    (opts, args) = parser.parse_args()

    if len(args) != 1:
        parser.error("unexpected number of arguments")

    global llvm_srcroot, indexstoredb_srcroot
    llvm_srcroot, = args
    indexstoredb_srcroot = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

    def import_header(dir, name):
        src = os.path.join(llvm_srcroot, 'include', 'llvm', dir, name)
        if os.path.exists(src):
            dst = os.path.join(indexstoredb_srcroot, 'include', 'llvm', dir, name)
            mkdir_p(os.path.dirname(dst))
            copyfile(src, dst)

    def import_c_header(name):
        src = os.path.join(llvm_srcroot, 'include', 'llvm-c', name)
        if os.path.exists(src):
            dst = os.path.join(indexstoredb_srcroot, 'include', 'llvm-c', name)
            mkdir_p(os.path.dirname(dst))
            copyfile(src, dst)

    def import_source(dir, name):
        src = os.path.join(llvm_srcroot, 'lib', dir, name)
        if os.path.exists(src):
            dst = os.path.join(indexstoredb_srcroot, 'lib', 'LLVMSupport', dir, name)
            mkdir_p(os.path.dirname(dst))
            copyfile(src, dst)

    print "note: importing from %r to %r" % (llvm_srcroot, indexstoredb_srcroot)

    for name in ADT_imports:
        import_header('ADT', name+'.h')
        if name not in Support_source_excludes:
            import_source('Support', name+'.c')
            import_source('Support', name+'.cpp')
    for name in Support_imports:
        if name.endswith('.def'):
            import_header('Support', name)
            continue
        import_header('Support', name+'.h')
        if name not in Support_source_excludes:
            import_source('Support', name+'.c')
            import_source('Support', name+'.cpp')
            import_source('Support', os.path.join('Unix', name+'.h'))
            import_source('Support', os.path.join('Unix', name+'.inc'))
            import_source('Support', os.path.join('Windows', name+'.h'))
            import_source('Support', os.path.join('Windows', name+'.inc'))

    for name in C_imports:
        import_c_header(name + '.h')

    print ""
    print "Adding prefix header includes"

    base_dirs = [
        os.path.join(indexstoredb_srcroot, 'include', 'llvm'),
        os.path.join(indexstoredb_srcroot, 'include', 'llvm-c'),
        os.path.join(indexstoredb_srcroot, 'lib', 'LLVMSupport'),
    ]

    include_dir = os.path.join(indexstoredb_srcroot, 'include')

    for base in base_dirs:
        for root, dirs, files in os.walk(base):
            for file in files:
                maybe_add_prefix(os.path.join(root, file), include_dir)

    print("")
    print("Copying llvm config files")
    # NOTE: If you're updating to a newer llvm, you may have to manually merge
    # these with the newer version produced by the llvm build
    indexstore_llvm_config_source = os.path.join(indexstoredb_srcroot, 'Utilities', 'import-llvm.d', 'include', 'llvm', 'Config')
    indexstore_llvm_config_target = os.path.join(include_dir, 'llvm', 'Config')
    shutil.copyfile(os.path.join(indexstore_llvm_config_source, 'config.h'),
                    os.path.join(indexstore_llvm_config_target, 'config.h'))
    shutil.copyfile(os.path.join(indexstore_llvm_config_source, 'llvm-config.h'),
                    os.path.join(indexstore_llvm_config_target, 'llvm-config.h'))

    print ""
    print "Applying patch"

    patch_path = os.path.join(indexstoredb_srcroot, 'Utilities', 'llvm-patch.diff'),
    subprocess.call("patch -p1 < '%s'" % patch_path, shell=True)

if __name__ == '__main__':
    main()
