som

#!/usr/bin/env python3
import argparse
import sys
import os
import shlex
import subprocess

BASE_DIR    = os.path.dirname(os.path.realpath(__file__))
TRUFFLE_DIR = BASE_DIR + '/../graal'
JAVA_HOME   = os.getenv('JAVA_HOME',  None)
JVMCI_BIN   = os.getenv('JVMCI_BIN',  None)
GRAAL_HOME  = os.getenv('GRAAL_HOME',  None)
GRAAL_FLAGS = os.getenv('GRAAL_FLAGS', None)

##
## Defining Argument Parsing
##
parser = argparse.ArgumentParser(
    description='TruffleSOM launcher and tool integration')

parser.add_argument('-d', '--debug', help='wait for debugger to attach',
                    dest='debug', action='store_true', default=False)
parser.add_argument('-t', '--num-threads', help='number of threads to be used',
                    dest='threads', default=None)
parser.add_argument('-dnu', '--stack-trace-on-dnu', help='Print a stack trace on #doesNotUnderstand:',
                    dest='som_dnu', action='store_true', default=False)
parser.add_argument('-di', '--dump-ir', help='Dump the IR, i.e., the AST or bytecode of a method',
                    dest='dump_ir', action='store_true', default=False)


explore = parser.add_argument_group('Explore and Investigate Execution')
explore.add_argument('-i', '--igv', help='dump compilation details to IGV',
                    dest='igv', action='store_true', default=False)
explore.add_argument('-if', '--igv-to-file', help='dump compilation details to file to be loaded by IGV',
                    dest='igv_to_file', action='store_true', default=False)
explore.add_argument('-io', '--igv-only', help='only dump named method, use of * allowed. Uses Invokable.toString()',
                    dest='only_igv', default=None)
explore.add_argument('-l', '--low-level', help='enable low-level optimization output',
                    dest='low_level', action='store_true', default=False)
explore.add_argument('-ti', '--trace-invalidation', help='trace assumption invalidation and transfers to interpreter',
                    dest='trace_invalidation', action='store_true', default=False)
explore.add_argument('-ts', '--trace-splitting', help='trace splitting decisions',
                    dest='trace_splitting', action='store_true', default=False)
explore.add_argument('-w', '--perf-warnings', help='enable performance warnings',
                    dest='perf_warnings', action='store_true', default=False)
explore.add_argument('-o', '--only', help='only compile give methods, comma separated list',
                    dest='only_compile', default=None)

profile = parser.add_argument_group('Profile Execution')
profile.add_argument('-gp', '--graal-profile', help='enable Graal-level profiling after warmup',
                    dest='graal_profile', action='store_true', default=False)
profile.add_argument('-ga', '--graal-profile-allocations', help='enable Graal-level profiling after warmup, and profile allocations',
                    dest='graal_profile_allocations', action='store_true', default=False)
profile.add_argument('-gi', '--graal-profile-intervals', help='enable Graal-level profiling after certain time intervals',
                    dest='graal_profile_timed', action='store_true', default=False)
profile.add_argument('-gb', '--graal-branch-profile', help='enable Graal-level branch profiling',
                    dest='graal_branch_profile', action='store_true', default=False)
profile.add_argument('-tp', '--truffle-profile', help='enable Graal-level profiling after warmup',
                    dest='truffle_profile', action='store_true', default=False)
explore.add_argument('-v', '--visual-vm', help='connect to VisualVM for profiling',
                    dest='visual_vm', action='store_true', default=False)

tools = parser.add_argument_group('Tools for various purposes')
tools.add_argument('-n', '--node-stats', help='collect details about AST nodes. Optionally define output file name. Default: node-stats.yml',
                   dest='nodestats', action='store', nargs='?',
                   const='node-stats.yml', default=False)
tools.add_argument('-cov', '--coverage', help='collect coverage statistics. Optionally define output file. Default is standard out.',
                   dest='coverage', action='store', nargs='?',
                   const='', default=False)
tools.add_argument('-ct', '--coverage-type', help="human readable 'histogram' (per file coverage summary) or 'detailed' (per line coverage summary), machine readable 'json', tool compliant 'lcov'. (default: histogram)",
                   dest='coverage_type', action='store', default='histogram')
tools.add_argument('-cs', '--cpu-sampler', help='Truffle CPU Sampler. [CPUSAMPLE] can be one of histogram, calltree, json',
                   dest='cpusampler', action='store', nargs='?',
                   const='histogram', default=False)
tools.add_argument('-fr', '--flight-recorder', help='profile with Java Flight Recorder',
                   dest='flight_recorder', action='store_true', default=False)


parser.add_argument('-A', '--no-assert', help='execute with assertions disabled',
                    dest='assert_', action='store_false', default=True)
parser.add_argument('-B', '--no-background', help='disable background compilation',
                    dest='background_compilation', action='store_false', default=True)
parser.add_argument('-C', '--no-compilation', help='disable Truffle compilation',
                    dest='no_compilation', action='store_true', default=False)
parser.add_argument('-G', '--interpreter', help='run without Graal',
                    dest='interpreter', action='store_true', default=False)
parser.add_argument('-EG', '--no-embedded-graal', help='run without the embedded Graal. When set, USE_EMBEDDED_GRAAL, JVMCI_BIN, or GRAAL_HOME are used.',
                    dest='use_embedded_graal', action='store', nargs='?',
                    const=False, default=True)
parser.add_argument('-LG', '--no-libgraal', help='run without using the embedded libgraal, which is a precompiled Graal',
                    dest='use_libgraal', action='store_false', default=True)
parser.add_argument('-X', '--java-interpreter', help='run without Graal, and only the Java interpreter',
                    dest='java_interpreter', action='store_true', default=False)
parser.add_argument('--no-graph-pe', help='disable Graph PE',
                    dest='graph_pe', action='store_false', default=True)



parser.add_argument('-vv', '--verbose', action='store_true', default=False,
                    dest='verbose', help="print command-line before executing")
parser.add_argument('--print-graal-options', action='store_true', default=False,
                    dest='print_graal_options', help="print all Graal options")

parser.add_argument('-J', help="Java VM Argument prefix",
                    dest="java_args", action='append')
parser.add_argument('-D', help="define a Java property",
                    dest="java_properties", action='append')


parser.add_argument('args', nargs=argparse.REMAINDER,
                    help='arguments passed to TruffleSOM')

setup = parser.add_argument_group('Setting up TruffleSOM')
setup.add_argument('--setup', choices=['labsjdk', 'mx', 'latest-mx', 'truffle-commit-id'],
                   dest='setup',
                   help='sets up the chosen dependency. LabsJDK is needed for running with Graal and building native images, and mx is the used build tool.')

args, unknown_args = parser.parse_known_args()


def get_mx_suite():
  import importlib.util
  import sys
  spec = importlib.util.spec_from_file_location('suite', BASE_DIR + '/mx.trufflesom/suite.py')
  module = importlib.util.module_from_spec(spec)
  spec.loader.exec_module(module)
  return getattr(module, 'suite')

def get_labs_jdk_id():
  return get_mx_suite()['libraries']['LABS_JDK']['id']

def get_truffle_commit_id():
  return get_mx_suite()['imports']['suites'][0]['version']

def find_mx(exit_when_missing=True):
  possible_mx = [BASE_DIR + '/../mx/mx', 'mx']
  for m in possible_mx:
    try:
      p = subprocess.run([m, '--version'], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
      if p.returncode == 0:
        return m
    except:
      pass

  if exit_when_missing:
    print("The mx build tool is not on the PATH or in the " + BASE_DIR + "/../mx folder. Please run `./som --setup mx`")
    exit(1)

  return None

def get_compiled_graalvm_java_bin(use_libgraal):
  graal_bin = None

  mx = find_mx()
  output = check_output(
    [mx, '--env', 'libgraal' if use_libgraal else 'graal', 'graalvm-home'],
    stderr=STDOUT, cwd=BASE_DIR
  ).decode()

  lines = output.strip().split('\n')
  if not lines:
    return None

  libgraal_jdk_home = lines[-1].strip()
  return libgraal_jdk_home + '/bin/java'


if args.setup == 'mx':
  mx = find_mx(False)
  if mx is None:
    p = subprocess.run(['git', 'clone', '--depth', '1', 'https://github.com/graalvm/mx.git', BASE_DIR + '/../mx'])
    exit(p.returncode)
  exit(0)

if args.setup == 'latest-mx':
  if not os.path.exists(BASE_DIR + '/../mx'):
    p = subprocess.run(['git', 'clone', '--depth', '1', 'https://github.com/graalvm/mx.git', BASE_DIR + '/../mx'])
    exit(p.returncode)
  else:
    subprocess.run(['git', '-C', BASE_DIR + '/../mx', 'fetch', '--all'])
    p = subprocess.run(['git', '-C', BASE_DIR + '/../mx', 'reset', '--hard', 'origin/HEAD'])
    exit(p.returncode)

if args.setup == 'labsjdk':
  if not os.path.exists(BASE_DIR + "/libs"):
    os.mkdir(BASE_DIR + "/libs")
  elif os.path.exists(BASE_DIR + "/libs/jvmci"):
    exit(0)

  mx = find_mx()
  # move the env file out of the way, because of latest mx changes
  os.rename(BASE_DIR + '/mx.trufflesom/env', BASE_DIR + '/mx.trufflesom/env-load-jdk')

  if os.path.exists(TRUFFLE_DIR + "/common.json"):
    graal_common_json = ["--configuration", TRUFFLE_DIR + "/common.json"]
  else:
    graal_common_json = []

  try:
    p = subprocess.run([
      mx, "--quiet", "fetch-jdk"] + graal_common_json + [
        "--strip-contents-home",
        "--jdk-id", get_labs_jdk_id(), "--alias", BASE_DIR + "/libs/jvmci"])
  finally:
    os.rename(BASE_DIR + '/mx.trufflesom/env-load-jdk', BASE_DIR + '/mx.trufflesom/env')
  exit(p.returncode)

if args.setup == 'truffle-commit-id':
  print(get_truffle_commit_id())
  exit(0)

if args.java_interpreter:
    args.interpreter = True

if args.flight_recorder:
    args.interpreter = True

if args.interpreter:
    args.use_embedded_graal = False

# Handle executable names
if sys.argv[0].endswith('fast'):
    args.assert_ = False
if sys.argv[0].endswith('debug'):
    args.perf_warnings = True
    args.trace_splitting = True
    args.background_compilation = False
    args.use_libgraal = False

# Determine JVM to be used
java_bin = None

# try local JVMCI, which ant already needed
local_jvmci_bin = BASE_DIR + '/libs/jvmci'
local_jvmci_bin += '/bin/java'
if os.path.isfile(local_jvmci_bin):
  java_bin = local_jvmci_bin

if not java_bin and JAVA_HOME:
  java_bin = JAVA_HOME + '/bin/java'

if not java_bin:
  java_bin = "java"


if args.use_embedded_graal is True:
  from subprocess import check_output, STDOUT, CalledProcessError
  try:
    java_bin = get_compiled_graalvm_java_bin(args.use_libgraal)
    if not java_bin or not os.path.isfile(java_bin):
      if args.use_libgraal:
        print("The use of LibGraal was requested, but it does not seem to be available.")
        print("To build it, please run `mx --env libgraal build`")
      else:
        print("The use of the Graal compiler was requested, but it does not seem to be available.")
        print("To build it, please run `mx --env graal build`")
      sys.exit(1)
  except CalledProcessError as e:
    print("Failed to determine location of libgraal")
    print(e.output.decode())
    sys.exit(1)
else:
    if args.use_embedded_graal is not False:
      if not os.path.isdir(args.use_embedded_graal):
        print("USE_EMBEDDED_GRAAL does not seem to exist: " + args.use_embedded_graal)
        sys.exit(1)

      if not os.path.isfile(args.use_embedded_graal + '/bin/java'):
        print("USE_EMBEDDED_GRAAL did not seem to have a bin/java: " + args.use_embedded_graal)
        sys.exit(1)

      java_bin = args.use_embedded_graal + '/bin/java'
    if not java_bin and JVMCI_BIN:
      java_bin = JVMCI_BIN
    if not java_bin and GRAAL_HOME and os.path.isfile(GRAAL_HOME + '/bin/java'):
      java_bin = GRAAL_HOME + '/bin/java'


##
## Defining Necessary Parameter Bits
##
LIBGRAAL_JAR    = TRUFFLE_DIR + '/compiler/mxbuild/dists/graal-truffle-compiler-libgraal.jar'
COVERAGE_JAR    = TRUFFLE_DIR + '/tools/mxbuild/dists/truffle-coverage.jar'
PROFILER_JAR    = TRUFFLE_DIR + '/tools/mxbuild/dists/truffle-profiler.jar'

if args.use_libgraal:
  GRAAL_JVMCI_FLAGS = ['-XX:+UnlockExperimentalVMOptions', '-XX:+EnableJVMCI', '-XX:+UseJVMCICompiler', '-XX:+UseJVMCINativeLibrary']
else:
  GRAAL_JVMCI_FLAGS = ['-XX:+UnlockExperimentalVMOptions', '-XX:+EnableJVMCI', '-XX:+UseJVMCICompiler', '-XX:-UseJVMCINativeLibrary']

MODULE_PATH_ENTRIES = [
  BASE_DIR + '/mxbuild/dists/trufflesom.jar',
  TRUFFLE_DIR + '/sdk/mxbuild/dists/graal-sdk.jar',
  TRUFFLE_DIR + '/sdk/mxbuild/dists/collections.jar',
  TRUFFLE_DIR + '/sdk/mxbuild/dists/polyglot.jar',
  TRUFFLE_DIR + '/sdk/mxbuild/dists/word.jar',
  TRUFFLE_DIR + '/sdk/mxbuild/dists/nativeimage.jar',
  TRUFFLE_DIR + '/truffle/mxbuild/dists/truffle-api.jar']

SOM_ARGS = ['--module', 'trufflesom/trufflesom.Launcher']

# == Compiler Settings
TWEAK_INLINING = ['-Dpolyglot.engine.CompilationThreshold=191',
                  '-Dpolyglot.engine.InliningMaxCallerSize=10000',
                  '-Dpolyglot.engine.SplittingMaxCalleeSize=100000']

JAVA_ARGS = ['-server', '-XX:+UseParallelGC']

##
## Processing Parameters and Assembling Command Line
##

if not args.interpreter and GRAAL_FLAGS:
    flags = shlex.split(str.strip(GRAAL_FLAGS))
else:
    flags = []

if args.dump_ir:
    SOM_ARGS += ['-di']

if args.interpreter:
    flags += ['-Dtruffle.TruffleRuntime=com.oracle.truffle.api.impl.DefaultTruffleRuntime',
              '-Dpolyglot.engine.WarnInterpreterOnly=false']
else:
    flags += GRAAL_JVMCI_FLAGS
    flags += ['-Dpolyglot.engine.MultiTier=false',
              '-Dpolyglot.engine.DynamicCompilationThresholds=false',
              '-Dpolyglot.engine.SingleTierCompilationThreshold=253',
              '-Dpolyglot.engine.Mode=throughput',
              '-Dpolyglot.engine.CompilationFailureAction=ExitVM']

if args.som_dnu:
    flags += ['-Dsom.printStackTraceOnDNU=true']

if args.only_igv:
    args.igv = True

if args.debug:
    flags += ['-Xdebug',
              '-Xrunjdwp:transport=dt_socket,server=y,suspend=y,address=8000']
if not args.interpreter and (args.igv or args.igv_to_file):
    flags += ['-Dgraal.Dump=Truffle,TruffleTree:2']
if not args.interpreter and args.only_igv:
    flags += ['-Dpolyglot.engine.MethodFilter=' + args.only_igv]
if not args.interpreter and args.igv_to_file:
    flags += ['-Dpolyglot.engine.PrintIdealGraphFile=true']
if args.low_level:
    flags += ['-XX:+UnlockDiagnosticVMOptions', '-XX:+LogCompilation',
              '-XX:+TraceDeoptimization']
if not args.interpreter and (args.graal_profile or args.graal_profile_allocations or args.graal_profile_timed):
    flags += ['-XX:JVMCICounterSize=5000', '-Dpolyglot.engine.ProfileCompiledMethods=true',
              '-DProfileCompiledMethodsPhase.WITH_SECTION_HEADER=true']
    if args.graal_profile_allocations:
        flags += ['-Dpolyglot.engine.ProfileAllocations=true']
    if args.graal_profile_timed:
        flags += ['-Dpolyglot.engine.TimedDynamicCounters=1000']
    if args.graal_profile:
        flags += ['-Dpolyglot.engine.BenchmarkDynamicCounters=out,completed,total']

if not args.interpreter and args.graal_branch_profile:
    flags += ['-Dpolyglot.engine.InstrumentBranches=true',
              '-Dpolyglot.engine.InstrumentBranchesFilter=*',
              '-Dpolyglot.engine.InstrumentBranchesPerInlineSite=true']

if not args.interpreter and args.perf_warnings:
    flags += ['-Dpolyglot.engine.CompilationExceptionsAreFatal=true',
              '-Dpolyglot.engine.TracePerformanceWarnings=all',
              '-Dpolyglot.engine.TraceCompilation=true',
              '-Dpolyglot.engine.TraceCompilationDetails=true']
if not args.interpreter and args.trace_invalidation:
    flags += ['-Dpolyglot.engine.TraceTransferToInterpreter=true',
              '-Dpolyglot.engine.TraceAssumptions=true']
if not args.interpreter and args.trace_splitting:
    flags += ['-Dpolyglot.engine.TraceSplittingSummary=true',
              '-Dpolyglot.engine.TraceSplitting=true']
if not args.interpreter and args.only_compile:
    flags.append("-Dpolyglot.engine.CompileOnly=%s" % args.only_compile)
if args.visual_vm:
    flags += ['-agentpath:/Users/smarr/Downloads/visualvm_205/visualvm/lib/deployed/jdk16/mac/libprofilerinterface.jnilib=/Users/smarr/Downloads/visualvm_205/visualvm/lib,5140']
if args.flight_recorder:
    flags += ['-XX:+UnlockCommercialFeatures', '-XX:+FlightRecorder',
              '-XX:StartFlightRecording=delay=10s,duration=10d,name=fr-recording2,filename=fr-recording2.jfr,settings=profile']

if args.nodestats:
    flags += ['-Dpolyglot.nodestats.OutputFile=' + args.nodestats, '-Dpolyglot.nodestats=true']

if args.coverage != False:
    MODULE_PATH_ENTRIES.append(COVERAGE_JAR)
    flags += ['-Dpolyglot.coverage=true',
              '-Dpolyglot.coverage.Count=true',
              '-Dpolyglot.coverage.StrictLines=false',
              '-Dpolyglot.coverage.Output=' + args.coverage_type]
    if args.coverage:
      flags += ['-Dpolyglot.coverage.OutputFile=' + args.coverage]

if args.cpusampler != False:
    MODULE_PATH_ENTRIES.append(PROFILER_JAR)
    flags += ['-Dpolyglot.cpusampler=true', '-Dpolyglot.cpusampler.Output=' + args.cpusampler]


if args.assert_:
    flags += ['-esa', '-ea']
else:
    flags += ['-dsa', '-da']
if not args.interpreter and not args.background_compilation:
    flags += ['-Dpolyglot.engine.BackgroundCompilation=false']
if not args.interpreter and args.no_compilation:
    flags.append('-Dpolyglot.engine.CompileOnly=__FAKE_METHOD_NON_EXISTING__')
if not args.interpreter and not args.graph_pe:
    flags += ['-Dpolyglot.engine.GraphPE=false']
if args.threads:
    flags += ['-Dsom.threads=%s' % args.threads ]
if args.java_interpreter:
    flags += ['-Xint']
if args.print_graal_options:
    flags += ['-XX:+JVMCIPrintProperties']

if args.java_properties:
    flags += ['-D' + property for property in args.java_properties]

if args.java_args:
    JAVA_ARGS += ['-' + property for property in args.java_args]

# HACK: Havlak needs a lot of stack to run reliably...
if 'Havlak' in args.args:
  # double the standard stack size
  # we don't do it always, because...
  JAVA_ARGS += ['-Xss3072k']

flags += [
  '--module-path=' + ':'.join(MODULE_PATH_ENTRIES),
  '--upgrade-module-path=' + LIBGRAAL_JAR
]

all_args = JAVA_ARGS + flags + SOM_ARGS + unknown_args + args.args

if args.verbose:
    print("CMD: " + java_bin + ' ' + ' '.join(all_args))

env = dict(os.environ)
env['JVMCI_VERSION_CHECK'] = 'ignore'

try:
  os.execvpe(java_bin, all_args, env)
except OSError as e:
  if e.errno == 2 and e.strerror == "No such file or directory":
    print("Could not find JVM: " + java_bin)
    sys.exit(1)