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