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