1"""util.py - General utilities for running, loading, and processing benchmarks
2"""
3import json
4import os
5import tempfile
6import subprocess
7import sys
8import functools
9
10# Input file type enumeration
11IT_Invalid = 0
12IT_JSON = 1
13IT_Executable = 2
14
15_num_magic_bytes = 2 if sys.platform.startswith('win') else 4
16
17
18def is_executable_file(filename):
19    """
20    Return 'True' if 'filename' names a valid file which is likely
21    an executable. A file is considered an executable if it starts with the
22    magic bytes for a EXE, Mach O, or ELF file.
23    """
24    if not os.path.isfile(filename):
25        return False
26    with open(filename, mode='rb') as f:
27        magic_bytes = f.read(_num_magic_bytes)
28    if sys.platform == 'darwin':
29        return magic_bytes in [
30            b'\xfe\xed\xfa\xce',  # MH_MAGIC
31            b'\xce\xfa\xed\xfe',  # MH_CIGAM
32            b'\xfe\xed\xfa\xcf',  # MH_MAGIC_64
33            b'\xcf\xfa\xed\xfe',  # MH_CIGAM_64
34            b'\xca\xfe\xba\xbe',  # FAT_MAGIC
35            b'\xbe\xba\xfe\xca'   # FAT_CIGAM
36        ]
37    elif sys.platform.startswith('win'):
38        return magic_bytes == b'MZ'
39    else:
40        return magic_bytes == b'\x7FELF'
41
42
43def is_json_file(filename):
44    """
45    Returns 'True' if 'filename' names a valid JSON output file.
46    'False' otherwise.
47    """
48    try:
49        with open(filename, 'r') as f:
50            json.load(f)
51        return True
52    except BaseException:
53        pass
54    return False
55
56
57def classify_input_file(filename):
58    """
59    Return a tuple (type, msg) where 'type' specifies the classified type
60    of 'filename'. If 'type' is 'IT_Invalid' then 'msg' is a human readable
61    string represeting the error.
62    """
63    ftype = IT_Invalid
64    err_msg = None
65    if not os.path.exists(filename):
66        err_msg = "'%s' does not exist" % filename
67    elif not os.path.isfile(filename):
68        err_msg = "'%s' does not name a file" % filename
69    elif is_executable_file(filename):
70        ftype = IT_Executable
71    elif is_json_file(filename):
72        ftype = IT_JSON
73    else:
74        err_msg = "'%s' does not name a valid benchmark executable or JSON file" % filename
75    return ftype, err_msg
76
77
78def check_input_file(filename):
79    """
80    Classify the file named by 'filename' and return the classification.
81    If the file is classified as 'IT_Invalid' print an error message and exit
82    the program.
83    """
84    ftype, msg = classify_input_file(filename)
85    if ftype == IT_Invalid:
86        print("Invalid input file: %s" % msg)
87        sys.exit(1)
88    return ftype
89
90
91def find_benchmark_flag(prefix, benchmark_flags):
92    """
93    Search the specified list of flags for a flag matching `<prefix><arg>` and
94    if it is found return the arg it specifies. If specified more than once the
95    last value is returned. If the flag is not found None is returned.
96    """
97    assert prefix.startswith('--') and prefix.endswith('=')
98    result = None
99    for f in benchmark_flags:
100        if f.startswith(prefix):
101            result = f[len(prefix):]
102    return result
103
104
105def remove_benchmark_flags(prefix, benchmark_flags):
106    """
107    Return a new list containing the specified benchmark_flags except those
108    with the specified prefix.
109    """
110    assert prefix.startswith('--') and prefix.endswith('=')
111    return [f for f in benchmark_flags if not f.startswith(prefix)]
112
113
114def load_benchmark_results(fname):
115    """
116    Read benchmark output from a file and return the JSON object.
117    REQUIRES: 'fname' names a file containing JSON benchmark output.
118    """
119    with open(fname, 'r') as f:
120        return json.load(f)
121
122
123def sort_benchmark_results(result):
124    benchmarks = result['benchmarks']
125
126    # From inner key to the outer key!
127    benchmarks = sorted(
128        benchmarks, key=lambda benchmark: benchmark['repetition_index'] if 'repetition_index' in benchmark else -1)
129    benchmarks = sorted(
130        benchmarks, key=lambda benchmark: 1 if 'run_type' in benchmark and benchmark['run_type'] == "aggregate" else 0)
131    benchmarks = sorted(
132        benchmarks, key=lambda benchmark: benchmark['per_family_instance_index'] if 'per_family_instance_index' in benchmark else -1)
133    benchmarks = sorted(
134        benchmarks, key=lambda benchmark: benchmark['family_index'] if 'family_index' in benchmark else -1)
135
136    result['benchmarks'] = benchmarks
137    return result
138
139
140def run_benchmark(exe_name, benchmark_flags):
141    """
142    Run a benchmark specified by 'exe_name' with the specified
143    'benchmark_flags'. The benchmark is run directly as a subprocess to preserve
144    real time console output.
145    RETURNS: A JSON object representing the benchmark output
146    """
147    output_name = find_benchmark_flag('--benchmark_out=',
148                                      benchmark_flags)
149    is_temp_output = False
150    if output_name is None:
151        is_temp_output = True
152        thandle, output_name = tempfile.mkstemp()
153        os.close(thandle)
154        benchmark_flags = list(benchmark_flags) + \
155            ['--benchmark_out=%s' % output_name]
156
157    cmd = [exe_name] + benchmark_flags
158    print("RUNNING: %s" % ' '.join(cmd))
159    exitCode = subprocess.call(cmd)
160    if exitCode != 0:
161        print('TEST FAILED...')
162        sys.exit(exitCode)
163    json_res = load_benchmark_results(output_name)
164    if is_temp_output:
165        os.unlink(output_name)
166    return json_res
167
168
169def run_or_load_benchmark(filename, benchmark_flags):
170    """
171    Get the results for a specified benchmark. If 'filename' specifies
172    an executable benchmark then the results are generated by running the
173    benchmark. Otherwise 'filename' must name a valid JSON output file,
174    which is loaded and the result returned.
175    """
176    ftype = check_input_file(filename)
177    if ftype == IT_JSON:
178        return load_benchmark_results(filename)
179    if ftype == IT_Executable:
180        return run_benchmark(filename, benchmark_flags)
181    raise ValueError('Unknown file type %s' % ftype)
182