1#!/usr/bin/env python3
2'''A utility to update LLVM IR CHECK lines in C/C++ FileCheck test files.
3
4Example RUN lines in .c/.cc test files:
5
6// RUN: %clang -emit-llvm -S %s -o - -O2 | FileCheck %s
7// RUN: %clangxx -emit-llvm -S %s -o - -O2 | FileCheck -check-prefix=CHECK-A %s
8
9Usage:
10
11% utils/update_cc_test_checks.py --llvm-bin=release/bin test/a.cc
12% utils/update_cc_test_checks.py --c-index-test=release/bin/c-index-test \
13  --clang=release/bin/clang /tmp/c/a.cc
14'''
15
16import argparse
17import collections
18import distutils.spawn
19import os
20import shlex
21import string
22import subprocess
23import sys
24import re
25import tempfile
26
27from UpdateTestChecks import asm, common
28
29ADVERT = '// NOTE: Assertions have been autogenerated by '
30
31CHECK_RE = re.compile(r'^\s*//\s*([^:]+?)(?:-NEXT|-NOT|-DAG|-LABEL)?:')
32RUN_LINE_RE = re.compile(r'^//\s*RUN:\s*(.*)$')
33
34SUBST = {
35    '%clang': [],
36    '%clang_cc1': ['-cc1'],
37    '%clangxx': ['--driver-mode=g++'],
38}
39
40def get_line2spell_and_mangled(args, clang_args):
41  ret = {}
42  with tempfile.NamedTemporaryFile() as f:
43    # TODO Make c-index-test print mangled names without circumventing through precompiled headers
44    status = subprocess.run([args.c_index_test, '-write-pch', f.name, *clang_args],
45                            stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
46    if status.returncode:
47      sys.stderr.write(status.stdout.decode())
48      sys.exit(2)
49    output = subprocess.check_output([args.c_index_test,
50        '-test-print-mangle', f.name])
51    if sys.version_info[0] > 2:
52      output = output.decode()
53  DeclRE = re.compile(r'^FunctionDecl=(\w+):(\d+):\d+ \(Definition\)')
54  MangleRE = re.compile(r'.*\[mangled=([^]]+)\]')
55  MatchedDecl = False
56  for line in output.splitlines():
57    # Get the function source name, line number and mangled name.  Sometimes
58    # c-index-test outputs the mangled name on a separate line (this can happen
59    # with block comments in front of functions).  Keep scanning until we see
60    # the mangled name.
61    decl_m = DeclRE.match(line)
62    mangle_m = MangleRE.match(line)
63
64    if decl_m:
65      MatchedDecl = True
66      spell, lineno = decl_m.groups()
67    if MatchedDecl and mangle_m:
68      mangled = mangle_m.group(1)
69      MatchedDecl = False
70    else:
71      continue
72
73    if mangled == '_' + spell:
74      # HACK for MacOS (where the mangled name includes an _ for C but the IR won't):
75      mangled = spell
76    # Note -test-print-mangle does not print file names so if #include is used,
77    # the line number may come from an included file.
78    ret[int(lineno)-1] = (spell, mangled)
79  if args.verbose:
80    for line, func_name in sorted(ret.items()):
81      print('line {}: found function {}'.format(line+1, func_name), file=sys.stderr)
82  return ret
83
84
85def config():
86  parser = argparse.ArgumentParser(
87      description=__doc__,
88      formatter_class=argparse.RawTextHelpFormatter)
89  parser.add_argument('-v', '--verbose', action='store_true')
90  parser.add_argument('--llvm-bin', help='llvm $prefix/bin path')
91  parser.add_argument('--clang',
92                      help='"clang" executable, defaults to $llvm_bin/clang')
93  parser.add_argument('--clang-args',
94                      help='Space-separated extra args to clang, e.g. --clang-args=-v')
95  parser.add_argument('--c-index-test',
96                      help='"c-index-test" executable, defaults to $llvm_bin/c-index-test')
97  parser.add_argument('--opt',
98                      help='"opt" executable, defaults to $llvm_bin/opt')
99  parser.add_argument(
100      '--functions', nargs='+', help='A list of function name regexes. '
101      'If specified, update CHECK lines for functions matching at least one regex')
102  parser.add_argument(
103      '--x86_extra_scrub', action='store_true',
104      help='Use more regex for x86 matching to reduce diffs between various subtargets')
105  parser.add_argument('-u', '--update-only', action='store_true',
106                      help='Only update test if it was already autogened')
107  parser.add_argument('tests', nargs='+')
108  args = parser.parse_args()
109  args.clang_args = shlex.split(args.clang_args or '')
110
111  if args.clang is None:
112    if args.llvm_bin is None:
113      args.clang = 'clang'
114    else:
115      args.clang = os.path.join(args.llvm_bin, 'clang')
116  if not distutils.spawn.find_executable(args.clang):
117    print('Please specify --llvm-bin or --clang', file=sys.stderr)
118    sys.exit(1)
119
120  # Determine the builtin includes directory so that we can update tests that
121  # depend on the builtin headers. See get_clang_builtin_include_dir() and
122  # use_clang() in llvm/utils/lit/lit/llvm/config.py.
123  try:
124    builtin_include_dir = subprocess.check_output(
125      [args.clang, '-print-file-name=include']).decode().strip()
126    SUBST['%clang_cc1'] = ['-cc1', '-internal-isystem', builtin_include_dir,
127                           '-nostdsysteminc']
128  except subprocess.CalledProcessError:
129    common.warn('Could not determine clang builtins directory, some tests '
130                'might not update correctly.')
131
132  if args.opt is None:
133    if args.llvm_bin is None:
134      args.opt = 'opt'
135    else:
136      args.opt = os.path.join(args.llvm_bin, 'opt')
137  if not distutils.spawn.find_executable(args.opt):
138    # Many uses of this tool will not need an opt binary, because it's only
139    # needed for updating a test that runs clang | opt | FileCheck. So we
140    # defer this error message until we find that opt is actually needed.
141    args.opt = None
142
143  if args.c_index_test is None:
144    if args.llvm_bin is None:
145      args.c_index_test = 'c-index-test'
146    else:
147      args.c_index_test = os.path.join(args.llvm_bin, 'c-index-test')
148  if not distutils.spawn.find_executable(args.c_index_test):
149    print('Please specify --llvm-bin or --c-index-test', file=sys.stderr)
150    sys.exit(1)
151
152  return args
153
154
155def get_function_body(args, filename, clang_args, extra_commands, prefixes, triple_in_cmd, func_dict):
156  # TODO Clean up duplication of asm/common build_function_body_dictionary
157  # Invoke external tool and extract function bodies.
158  raw_tool_output = common.invoke_tool(args.clang, clang_args, filename)
159  for extra_command in extra_commands:
160    extra_args = shlex.split(extra_command)
161    with tempfile.NamedTemporaryFile() as f:
162      f.write(raw_tool_output.encode())
163      f.flush()
164      if extra_args[0] == 'opt':
165        if args.opt is None:
166          print(filename, 'needs to run opt. '
167                'Please specify --llvm-bin or --opt', file=sys.stderr)
168          sys.exit(1)
169        extra_args[0] = args.opt
170      raw_tool_output = common.invoke_tool(extra_args[0],
171                                           extra_args[1:], f.name)
172  if '-emit-llvm' in clang_args:
173    common.build_function_body_dictionary(
174            common.OPT_FUNCTION_RE, common.scrub_body, [],
175            raw_tool_output, prefixes, func_dict, args.verbose, False)
176  else:
177    print('The clang command line should include -emit-llvm as asm tests '
178          'are discouraged in Clang testsuite.', file=sys.stderr)
179    sys.exit(1)
180
181
182def main():
183  args = config()
184  script_name = os.path.basename(__file__)
185  autogenerated_note = (ADVERT + 'utils/' + script_name)
186
187  for filename in args.tests:
188    with open(filename) as f:
189      input_lines = [l.rstrip() for l in f]
190
191    first_line = input_lines[0] if input_lines else ""
192    if 'autogenerated' in first_line and script_name not in first_line:
193      common.warn("Skipping test which wasn't autogenerated by " + script_name, filename)
194      continue
195
196    if args.update_only:
197      if not first_line or 'autogenerated' not in first_line:
198        common.warn("Skipping test which isn't autogenerated: " + filename)
199        continue
200
201    # Extract RUN lines.
202    raw_lines = [m.group(1)
203                 for m in [RUN_LINE_RE.match(l) for l in input_lines] if m]
204    run_lines = [raw_lines[0]] if len(raw_lines) > 0 else []
205    for l in raw_lines[1:]:
206      if run_lines[-1].endswith("\\"):
207        run_lines[-1] = run_lines[-1].rstrip("\\") + " " + l
208      else:
209        run_lines.append(l)
210
211    if args.verbose:
212      print('Found {} RUN lines:'.format(len(run_lines)), file=sys.stderr)
213      for l in run_lines:
214        print('  RUN: ' + l, file=sys.stderr)
215
216    # Build a list of clang command lines and check prefixes from RUN lines.
217    run_list = []
218    line2spell_and_mangled_list = collections.defaultdict(list)
219    for l in run_lines:
220      commands = [cmd.strip() for cmd in l.split('|')]
221
222      triple_in_cmd = None
223      m = common.TRIPLE_ARG_RE.search(commands[0])
224      if m:
225        triple_in_cmd = m.groups()[0]
226
227      # Apply %clang substitution rule, replace %s by `filename`, and append args.clang_args
228      clang_args = shlex.split(commands[0])
229      if clang_args[0] not in SUBST:
230        print('WARNING: Skipping non-clang RUN line: ' + l, file=sys.stderr)
231        continue
232      clang_args[0:1] = SUBST[clang_args[0]]
233      clang_args = [filename if i == '%s' else i for i in clang_args] + args.clang_args
234
235      # Permit piping the output through opt
236      if not (len(commands) == 2 or
237              (len(commands) == 3 and commands[1].startswith('opt'))):
238        print('WARNING: Skipping non-clang RUN line: ' + l, file=sys.stderr)
239
240      # Extract -check-prefix in FileCheck args
241      filecheck_cmd = commands[-1]
242      common.verify_filecheck_prefixes(filecheck_cmd)
243      if not filecheck_cmd.startswith('FileCheck '):
244        print('WARNING: Skipping non-FileChecked RUN line: ' + l, file=sys.stderr)
245        continue
246      check_prefixes = [item for m in common.CHECK_PREFIX_RE.finditer(filecheck_cmd)
247                               for item in m.group(1).split(',')]
248      if not check_prefixes:
249        check_prefixes = ['CHECK']
250      run_list.append((check_prefixes, clang_args, commands[1:-1], triple_in_cmd))
251
252    # Strip CHECK lines which are in `prefix_set`, update test file.
253    prefix_set = set([prefix for p in run_list for prefix in p[0]])
254    input_lines = []
255    with open(filename, 'r+') as f:
256      for line in f:
257        m = CHECK_RE.match(line)
258        if not (m and m.group(1) in prefix_set) and line != '//\n':
259          input_lines.append(line)
260      f.seek(0)
261      f.writelines(input_lines)
262      f.truncate()
263
264    # Execute clang, generate LLVM IR, and extract functions.
265    func_dict = {}
266    for p in run_list:
267      prefixes = p[0]
268      for prefix in prefixes:
269        func_dict.update({prefix: dict()})
270    for prefixes, clang_args, extra_commands, triple_in_cmd in run_list:
271      if args.verbose:
272        print('Extracted clang cmd: clang {}'.format(clang_args), file=sys.stderr)
273        print('Extracted FileCheck prefixes: {}'.format(prefixes), file=sys.stderr)
274
275      get_function_body(args, filename, clang_args, extra_commands, prefixes, triple_in_cmd, func_dict)
276
277      # Invoke c-index-test to get mapping from start lines to mangled names.
278      # Forward all clang args for now.
279      for k, v in get_line2spell_and_mangled(args, clang_args).items():
280        line2spell_and_mangled_list[k].append(v)
281
282    output_lines = [autogenerated_note]
283    for idx, line in enumerate(input_lines):
284      # Discard any previous script advertising.
285      if line.startswith(ADVERT):
286        continue
287      if idx in line2spell_and_mangled_list:
288        added = set()
289        for spell, mangled in line2spell_and_mangled_list[idx]:
290          # One line may contain multiple function declarations.
291          # Skip if the mangled name has been added before.
292          # The line number may come from an included file,
293          # we simply require the spelling name to appear on the line
294          # to exclude functions from other files.
295          if mangled in added or spell not in line:
296            continue
297          if args.functions is None or any(re.search(regex, spell) for regex in args.functions):
298            if added:
299              output_lines.append('//')
300            added.add(mangled)
301            common.add_ir_checks(output_lines, '//', run_list, func_dict, mangled, False, False)
302      output_lines.append(line.rstrip('\n'))
303
304    # Update the test file.
305    with open(filename, 'w') as f:
306      for line in output_lines:
307        f.write(line + '\n')
308
309  return 0
310
311
312if __name__ == '__main__':
313  sys.exit(main())
314