main.py

import os
import shutil
from pathlib import Path
from itertools import dropwhile
import sys
import time
from typing import List

import argparse

import steps
import tui
from steps import StepBase
from constants import *
from core import CONFIG, dump


os.environ['QEMU_LD_PREFIX'] = "/usr/riscv64-linux-gnu"

parser = argparse.ArgumentParser(description=ProgramDescription, add_help=False)

group_system = parser.add_argument_group("Basic Options")
group_system.add_argument("-h", "--help", action="help", help="show this help message and exit")
group_system.add_argument('-j', '--jobs', type=int, default=1, required=False, metavar='N', help="allow N jobs at once")
group_system.add_argument('--filter', type=str, default="", required=False, metavar='prefix', help="only cases starting with 'prefix' will be tested")
group_system.add_argument('-n', '--nums', type=int, default=1, required=False, metavar='N', help="run N times repeatedly for benchmark")

group_test_content = parser.add_argument_group("Test cases").add_mutually_exclusive_group()
group_test_content.add_argument('-p', '--perf', action='store_true', help="test performance cases")
group_test_content.add_argument('-f', '--func', action='store_true', help="test functional cases (default)")
group_test_content.add_argument('-a', '--all', action='store_true', help="test both functional and performance tests")

group_test_target = parser.add_argument_group("Test Target").add_mutually_exclusive_group()
group_test_target.add_argument('-s', '--asm', action='store_true', help="test the generated asm (default)")
group_test_target.add_argument('-r', '--ir', action='store_true', help="test the middle ir only")

group_test_mode = parser.add_argument_group("Test Mode").add_mutually_exclusive_group()
group_test_mode.add_argument('-O0', '--no-opt', action='store_true', help="turn off the optimization (default)")
group_test_mode.add_argument('-O1', '--opt', action='store_true', help="turn on the optimization")
group_test_mode.add_argument('-t', '--timeit', action='store_true', help="test both and compare the results")
group_test_mode.add_argument('-b', '--benchmark', action='store_true', help="compare with the clang (--ir or --asm)")

group_additional = parser.add_argument_group("Addition Arguments")
group_additional.add_argument('--', nargs='*', help="arguments after '--' will be passed to the compiler", dest="arg")


if __name__ == "__main__":
    args, remaining_args = parser.parse_known_args()
    remaining_args = list(dropwhile(lambda x: x != '--', remaining_args))
    if len(remaining_args) > 0:
        remaining_args = remaining_args[1:]

    steps_to_run: List[StepBase] = []
    timestamp = int(time.time())

    if args.all:
        args.func = True
        args.perf = True
    elif not args.perf:
        args.func = True

    if args.timeit:
        steps_to_run.append(steps.RunMain(
            max_workers=args.jobs,
            testcase_functional=args.func,
            testcase_performance=args.perf,
            test_ir_only=args.ir,
            test_with_opt=False,
            benchmark_with_llvm=False,
            prefix_filter=args.filter,
            tailing_args=remaining_args,
            repeated=args.nums,
        ))
        steps_to_run.append(steps.ResultCollect(
            test_with_opt=False,
            storage=f"{timestamp}0",
        ))
        steps_to_run.append(steps.RunMain(
            max_workers=args.jobs,
            testcase_functional=args.func,
            testcase_performance=args.perf,
            test_ir_only=args.ir,
            test_with_opt=True,
            benchmark_with_llvm=False,
            prefix_filter=args.filter,
            tailing_args=remaining_args,
            repeated=args.nums,
        ))
        steps_to_run.append(steps.ResultCollect(
            test_with_opt=True,
            storage=f"{timestamp}1",
        ))
        steps_to_run.append(steps.BenchmarkCollect(
            benchmark_self=True,
            storage=f"{timestamp}2",
            repeated=args.nums,
        ))
    elif args.benchmark:
        steps_to_run.append(steps.RunMain(
            max_workers=args.jobs,
            testcase_functional=args.func,
            testcase_performance=args.perf,
            test_ir_only=args.ir,
            test_with_opt=True,
            benchmark_with_llvm=False,
            prefix_filter=args.filter,
            tailing_args=remaining_args,
            repeated=args.nums,
        ))
        steps_to_run.append(steps.ResultCollect(
            test_with_opt=True,
            storage=f"{timestamp}1",
        ))
        steps_to_run.append(steps.RunMain(
            max_workers=args.jobs,
            testcase_functional=args.func,
            testcase_performance=args.perf,
            test_ir_only=args.ir,
            test_with_opt=True,
            benchmark_with_llvm=True,
            prefix_filter=args.filter,
            tailing_args=remaining_args,
            repeated=args.nums,
        ))
        steps_to_run.append(steps.BenchmarkCollect(
            benchmark_self=False,
            storage=f"{timestamp}2",
            repeated=args.nums,
        ))
    else:
        steps_to_run.append(steps.RunMain(
            max_workers=args.jobs,
            testcase_functional=args.func,
            testcase_performance=args.perf,
            test_ir_only=args.ir,
            test_with_opt=args.opt,
            benchmark_with_llvm=False,
            prefix_filter=args.filter,
            tailing_args=remaining_args,
        ))
        steps_to_run.append(steps.ResultCollect(
            test_with_opt=args.opt,
            storage=f"{timestamp}{int(args.opt)}",
        ))

    if (Path(CONFIG["paths.storage"]) / CONFIG["paths.names.cache"]).exists():
        shutil.rmtree(str(Path(CONFIG["paths.storage"]) / CONFIG["paths.names.cache"]))

    for i, step in enumerate(steps_to_run, 1):
        print(f"[{i}/{len(steps_to_run)}] {tui.C_CYAN(step.get_name())} ...")
        step.run()

    (Path(CONFIG["paths.storage"]) / CONFIG["paths.names.history"] / f"{timestamp}-.txt").write_text(dump(sys.argv))