#!@PYTHON_EXECUTABLE@
#
#===- streamexecutor-config - Build config script for SE -----*- python -*--===#
#
#                     The LLVM Compiler Infrastructure
#
# This file is distributed under the University of Illinois Open Source
# License. See LICENSE.TXT for details.
#
#===------------------------------------------------------------------------===#

r"""
Get configuration info needed to compile programs which use StreamExecutor.

Runs llvm-config and adds StreamExecutor-specific flags to the output. Supports
only the subset of llvm-config flags that are relevant for applications
compiling against StreamExecutor.

This utility will typically be used to construct a compile command line for an
application which depends on the StreamExecutor library.

For example:
  c++ example.cpp -o example \
      $(streamexecutor-config \
          --cppflags --cxxflags --ldflags --libs --system-libs)
"""

import argparse
import errno
import os
import shlex
import subprocess
import sys

# The following functions are configured by cmake. They use raw triple-quoted
# strings to surround values that are substituted by cmake at configure time.
# This kind of quoting should allow for paths that contain spaces.

def get_llvm_config_dir():
  """Gets the path to the llvm-config executable."""
  return r"""@LLVM_BINARY_DIR@/bin"""

def get_cmake_install_prefix():
  """Gets the value of the cmake variable CMAKE_INSTALL_PREFIX."""
  return r"""@CMAKE_INSTALL_PREFIX@"""

def convert_library_name(library_name):
  """Converts a library name ending in '.framework' into a '-framework' flag.

  This is used to support OS X.

  >>> convert_library_name('')
  ''

  >>> convert_library_name('/usr/local/lib64/libcuda.so')
  '/usr/local/lib64/libcuda.so'

  >>> convert_library_name('/Library/Frameworks/cuda.framework')
  '-framework cuda'
  """
  framework_suffix = '.framework'
  if library_name.endswith(framework_suffix):
    framework_name = os.path.basename(library_name)[:-len(framework_suffix)]
    library_name = '-framework ' + framework_name
  return library_name

def get_cuda_driver_library():
  """Gets the value of the cmake variable CUDA_DRIVER_LIBRARY."""
  return convert_library_name(r"""@CUDA_DRIVER_LIBRARY@""")

def cuddle_flag(flag, tokens):
  """If flag appears by itself in tokens, combines it with the next token.

  >>> tokens = ['-I', '/usr/include']
  >>> cuddle_flag('-I', tokens)
  >>> tokens
  ['-I/usr/include']

  >>> tokens = ['-L', '/usr/lib']
  >>> cuddle_flag('-L', tokens)
  >>> tokens
  ['-L/usr/lib']

  >>> tokens = ['-I']
  >>> cuddle_flag('-I', tokens)
  >>> tokens
  ['-I']

  >>> tokens = ['-I', '/usr/include', '-I', '/usr/local/include']
  >>> cuddle_flag('-I', tokens)
  >>> tokens
  ['-I/usr/include', '-I/usr/local/include']
  """
  start = 0
  while True:
    try:
      index = tokens.index(flag, start)
    except ValueError:
      return
    if index + 1 < len(tokens):
      follower = tokens.pop(index + 1)
      tokens[index] = flag + follower
    start = index + 1

def get_llvm_config_output_for_dir(llvm_config_dir, flags_string):
  """Calls llvm-config at the given path and returns the output with -I and -L
  flags cuddled."""
  output = subprocess.check_output(
      ['%s/llvm-config' % llvm_config_dir] + flags_string.split()).strip()
  tokens = shlex.split(output)
  cuddle_flag('-I', tokens)
  cuddle_flag('-L', tokens)
  return ' '.join(tokens)

def has_token(token, string):
  """Checks if the given token appears in the string.

  The token argument must be a single shell token.

  >>> string = '-I/usr/include -L"/usr/lib"'
  >>> has_token('-I/usr/include', string)
  True
  >>> has_token('-I/usr/local/include', string)
  False
  >>> has_token('-I"/usr/include"', string)
  True
  >>> has_token('-L"/usr/lib"', string)
  True
  >>> has_token('-L/usr/lib', string)
  True
  """
  split_token = shlex.split(token)
  if len(split_token) > 1:
    raise ValueError('has_token called with a multi-token token: ' + token)
  escaped_token = split_token[0]
  return escaped_token in shlex.split(string)

def main():
  parser = argparse.ArgumentParser(
      prog='streamexecutor-config',
      formatter_class=argparse.RawDescriptionHelpFormatter,
      description=__doc__)

  parser.add_argument(
      '--cppflags',
      action='store_true',
      help=
        'C preprocessor flags for files that include StreamExecutor headers.')

  parser.add_argument(
      '--cxxflags',
      action='store_true',
      help='C++ compiler flags for files that include StreamExecutor headers.')

  parser.add_argument(
      '--ldflags',
      action='store_true',
      help='Print linker flags.')

  parser.add_argument(
      '--libs',
      action='store_true',
      help='Libraries needed to link against StreamExecutor.')

  parser.add_argument(
      '--system-libs',
      action='store_true',
      help='System libraries needed to link against StreamExecutor.')

  parser.add_argument(
      '--llvm-config-dir',
      default=get_llvm_config_dir(),
      help='Directory containing the llvm-config executable. '\
          'If not specified, defaults to the cmake-configured location')

  args = parser.parse_args()

  # Print the help message if the user did not pass any flag arguments.
  if not any(
      getattr(args, flag)
        for flag in ('cppflags', 'cxxflags', 'ldflags', 'libs', 'system_libs')):
    parser.print_help()
    sys.exit(1)

  # Check for the presence of the llvm-config executable.
  if not os.path.isfile('%s/llvm-config' % args.llvm_config_dir):
    sys.exit('llvm-config not found in: ' + args.llvm_config_dir)
  if not os.access('%s/llvm-config' % args.llvm_config_dir, os.X_OK):
    sys.exit('llvm-config not executable in: ' + args.llvm_config_dir)

  # We will always use args.llvm_config_dir as the second argument to
  # get_llvm_config_output_for_path.
  get_llvm_config_output = lambda flags : get_llvm_config_output_for_dir(
      args.llvm_config_dir, flags)

  all_flags = []

  if args.cppflags:
    llvm_flags = get_llvm_config_output('--cppflags')
    all_flags.append(llvm_flags)
    se_flag = "-I%s/include" % get_cmake_install_prefix()
    if not has_token(token=se_flag, string=llvm_flags):
      all_flags.append(se_flag)

  if args.cxxflags:
    all_flags.append(get_llvm_config_output('--cxxflags'))

  if args.ldflags:
    llvm_flags = get_llvm_config_output('--ldflags')
    all_flags.append(llvm_flags)
    se_flag = "-L%s/lib" % get_cmake_install_prefix()
    if not has_token(token=se_flag, string=llvm_flags):
      all_flags.append(se_flag)

  if args.libs:
    llvm_flags = get_llvm_config_output('--libs support symbolize')
    se_flag = '-lstreamexecutor'
    if not has_token(token=se_flag, string=llvm_flags):
      all_flags.append(se_flag)
    cuda_driver_library = get_cuda_driver_library()
    if cuda_driver_library:
      all_flags.append(cuda_driver_library)
    all_flags.append(llvm_flags)

  if args.system_libs:
    all_flags.append(get_llvm_config_output('--system-libs'))

  print(' '.join(all_flags))

if __name__ == '__main__':
  main()