#!/usr/bin/env python3

# Automatically formatted with yapf (https://github.com/google/yapf)

# Script for automatic 'opt' pipeline reduction for when using the new
# pass-manager (NPM). Based around the '-print-pipeline-passes' option.
#
# The reduction algorithm consists of several phases (steps).
#
# Step #0: Verify that input fails with the given pipeline and make note of the
# error code.
#
# Step #1: Split pipeline in two starting from front and move forward as long as
# first pipeline exits normally and the second pipeline fails with the expected
# error code. Move on to step #2 with the IR from the split point and the
# pipeline from the second invocation.
#
# Step #2: Remove passes from end of the pipeline as long as the pipeline fails
# with the expected error code.
#
# Step #3: Make several sweeps over the remaining pipeline trying to remove one
# pass at a time. Repeat sweeps until unable to remove any more passes.
#
# Usage example:
# reduce_pipeline.py --opt-binary=./build-all-Debug/bin/opt --input=input.ll --output=output.ll --passes=PIPELINE [EXTRA-OPT-ARGS ...]

import argparse
import pipeline
import shutil
import subprocess
import tempfile

parser = argparse.ArgumentParser(
    description=
    'Automatic opt pipeline reducer. Unrecognized arguments are forwarded to opt.'
)
parser.add_argument('--opt-binary',
                    action='store',
                    dest='opt_binary',
                    default='opt')
parser.add_argument('--passes', action='store', dest='passes', required=True)
parser.add_argument('--input', action='store', dest='input', required=True)
parser.add_argument('--output', action='store', dest='output')
parser.add_argument('--dont-expand-passes',
                    action='store_true',
                    dest='dont_expand_passes',
                    help='Do not expand pipeline before starting reduction.')
parser.add_argument(
    '--dont-remove-empty-pm',
    action='store_true',
    dest='dont_remove_empty_pm',
    help='Do not remove empty pass-managers from the pipeline during reduction.'
)
[args, extra_opt_args] = parser.parse_known_args()

print('The following extra args will be passed to opt: {}'.format(
    extra_opt_args))

lst = pipeline.fromStr(args.passes)
passes = '-passes={}'.format(pipeline.toStr(lst))
ll_input = args.input

# Step #-1
# Launch 'opt' once with '-print-pipeline-passes' to expand pipeline before
# starting reduction. Allows specifying a default pipelines (e.g.
# '-passes=default<O3>').
if not args.dont_expand_passes:
    run_args = [
        args.opt_binary, '-disable-symbolication', '-disable-output',
        '-print-pipeline-passes', passes, ll_input
    ]
    run_args.extend(extra_opt_args)
    opt = subprocess.run(run_args,
                         stdout=subprocess.PIPE,
                         stderr=subprocess.PIPE)
    if opt.returncode != 0:
        print('Failed to expand passes. Aborting.')
        print(run_args)
        print('exitcode: {}'.format(opt.returncode))
        print(opt.stderr.decode())
        exit(1)
    stdout = opt.stdout.decode()
    stdout = stdout[:stdout.rfind('\n')]
    print('Expanded pass sequence: {}'.format(stdout))
    passes = '-passes={}'.format(stdout)

# Step #0
# Confirm that the given input, passes and options result in failure.
print('---Starting step #0---')
run_args = [
    args.opt_binary, '-disable-symbolication', '-disable-output', passes,
    ll_input
]
run_args.extend(extra_opt_args)
opt = subprocess.run(run_args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
if opt.returncode >= 0:
    print('Input does not result in failure as expected. Aborting.')
    print(run_args)
    print('exitcode: {}'.format(opt.returncode))
    print(opt.stderr.decode())
    exit(1)

expected_error_returncode = opt.returncode
print('-passes="{}"'.format(pipeline.toStr(lst)))

# Step #1
# Try to narrow down the failing pass sequence by splitting the pipeline in two
# opt invocations (A and B) starting with invocation A only running the first
# pipeline pass and invocation B the remaining. Keep moving the split point
# forward as long as invocation A exits normally and invocation B fails with
# the expected error. This will accomplish two things first the input IR will be
# further reduced and second, with that IR, the reduced pipeline for invocation
# B will be sufficient to reproduce.
print('---Starting step #1---')
prevLstB = None
prevIntermediate = None
tmpd = tempfile.TemporaryDirectory()

for idx in range(pipeline.count(lst)):
    [lstA, lstB] = pipeline.split(lst, idx)
    if not args.dont_remove_empty_pm:
        lstA = pipeline.prune(lstA)
        lstB = pipeline.prune(lstB)
    passesA = '-passes=' + pipeline.toStr(lstA)
    passesB = '-passes=' + pipeline.toStr(lstB)

    intermediate = 'intermediate-0.ll' if idx % 2 else 'intermediate-1.ll'
    intermediate = tmpd.name + '/' + intermediate
    run_args = [
        args.opt_binary, '-disable-symbolication', '-S', '-o', intermediate,
        passesA, ll_input
    ]
    run_args.extend(extra_opt_args)
    optA = subprocess.run(run_args,
                          stdout=subprocess.PIPE,
                          stderr=subprocess.PIPE)
    run_args = [
        args.opt_binary, '-disable-symbolication', '-disable-output', passesB,
        intermediate
    ]
    run_args.extend(extra_opt_args)
    optB = subprocess.run(run_args,
                          stdout=subprocess.PIPE,
                          stderr=subprocess.PIPE)
    if not (optA.returncode == 0
            and optB.returncode == expected_error_returncode):
        break
    prevLstB = lstB
    prevIntermediate = intermediate
if prevLstB:
    lst = prevLstB
    ll_input = prevIntermediate
print('-passes="{}"'.format(pipeline.toStr(lst)))

# Step #2
# Try removing passes from the end of the remaining pipeline while still
# reproducing the error.
print('---Starting step #2---')
prevLstA = None
for idx in reversed(range(pipeline.count(lst))):
    [lstA, lstB] = pipeline.split(lst, idx)
    if not args.dont_remove_empty_pm:
        lstA = pipeline.prune(lstA)
    passesA = '-passes=' + pipeline.toStr(lstA)
    run_args = [
        args.opt_binary, '-disable-symbolication', '-disable-output', passesA,
        ll_input
    ]
    run_args.extend(extra_opt_args)
    optA = subprocess.run(run_args,
                          stdout=subprocess.PIPE,
                          stderr=subprocess.PIPE)
    if optA.returncode != expected_error_returncode:
        break
    prevLstA = lstA
if prevLstA:
    lst = prevLstA
print('-passes="{}"'.format(pipeline.toStr(lst)))

# Step #3
# Now that we have a pipeline that is reduced both front and back we do
# exhaustive sweeps over the remainder trying to remove one pass at a time.
# Repeat as long as reduction is possible.
print('---Starting step #3---')
while True:
    keepGoing = False
    for idx in range(pipeline.count(lst)):
        candLst = pipeline.remove(lst, idx)
        if not args.dont_remove_empty_pm:
            candLst = pipeline.prune(candLst)
        passes = '-passes=' + pipeline.toStr(candLst)
        run_args = [
            args.opt_binary, '-disable-symbolication', '-disable-output',
            passes, ll_input
        ]
        run_args.extend(extra_opt_args)
        opt = subprocess.run(run_args,
                             stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE)
        if opt.returncode == expected_error_returncode:
            lst = candLst
            keepGoing = True
    if not keepGoing:
        break
print('-passes="{}"'.format(pipeline.toStr(lst)))

print('---FINISHED---')
if args.output:
    shutil.copy(ll_input, args.output)
    print('Wrote output to \'{}\'.'.format(args.output))
print('-passes="{}"'.format(pipeline.toStr(lst)))
exit(0)