1#!/usr/bin/env python3
2import re, sys
3
4def fix_string(s):
5    TYPE = re.compile('\s*(i[0-9]+|float|double|x86_fp80|fp128|ppc_fp128|\[\[.*?\]\]|\[2 x \[\[[A-Z_0-9]+\]\]\]|<.*?>|{.*?}|\[[0-9]+ x .*?\]|%["a-z:A-Z0-9._]+({{.*?}})?|%{{.*?}}|{{.*?}}|\[\[.*?\]\])(\s*(\*|addrspace\(.*?\)|dereferenceable\(.*?\)|byval\(.*?\)|sret|zeroext|inreg|returned|signext|nocapture|align \d+|swiftself|swifterror|readonly|noalias|inalloca|nocapture))*\s*')
6
7    counter = 0
8    if 'i32{{.*}}' in s:
9        counter = 1
10
11    at_pos = s.find('@')
12    if at_pos == -1:
13        at_pos = 0
14
15    annoying_pos = s.find('{{[^(]+}}')
16    if annoying_pos != -1:
17        at_pos = annoying_pos + 9
18
19    paren_pos = s.find('(', at_pos)
20    if paren_pos == -1:
21        return s
22
23    res = s[:paren_pos+1]
24    s = s[paren_pos+1:]
25
26    m = TYPE.match(s)
27    while m:
28        res += m.group()
29        s = s[m.end():]
30        if s.startswith(',') or s.startswith(')'):
31            res += f' %{counter}'
32            counter += 1
33
34        next_arg = s.find(',')
35        if next_arg == -1:
36            break
37
38        res += s[:next_arg+1]
39        s = s[next_arg+1:]
40        m = TYPE.match(s)
41
42    return res+s
43
44def process_file(contents):
45    PREFIX = re.compile(r'check-prefix(es)?(=|\s+)([a-zA-Z0-9,]+)')
46    check_prefixes = ['CHECK']
47    result = ''
48    for line in contents.split('\n'):
49        if 'FileCheck' in line:
50            m = PREFIX.search(line)
51            if m:
52                check_prefixes.extend(m.group(3).split(','))
53
54        found_check = False
55        for prefix in check_prefixes:
56            if prefix in line:
57                found_check = True
58                break
59
60        if not found_check or 'define' not in line:
61            result += line + '\n'
62            continue
63
64        # We have a check for a function definition. Number the args.
65        line = fix_string(line)
66        result += line + '\n'
67    return result
68
69def main():
70    print(f'Processing {sys.argv[1]}')
71    f = open(sys.argv[1])
72    content = f.read()
73    f.close()
74
75    content = process_file(content)
76
77    f = open(sys.argv[1], 'w')
78    f.write(content)
79    f.close()
80
81if __name__ == '__main__':
82    main()
83