#!/usr/bin/env python3
import re, sys

def fix_string(s):
    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*')

    counter = 0
    if 'i32{{.*}}' in s:
        counter = 1

    at_pos = s.find('@')
    if at_pos == -1:
        at_pos = 0

    annoying_pos = s.find('{{[^(]+}}')
    if annoying_pos != -1:
        at_pos = annoying_pos + 9

    paren_pos = s.find('(', at_pos)
    if paren_pos == -1:
        return s

    res = s[:paren_pos+1]
    s = s[paren_pos+1:]

    m = TYPE.match(s)
    while m:
        res += m.group()
        s = s[m.end():]
        if s.startswith(',') or s.startswith(')'):
            res += f' %{counter}'
            counter += 1

        next_arg = s.find(',')
        if next_arg == -1:
            break

        res += s[:next_arg+1]
        s = s[next_arg+1:]
        m = TYPE.match(s)

    return res+s

def process_file(contents):
    PREFIX = re.compile(r'check-prefix(es)?(=|\s+)([a-zA-Z0-9,]+)')
    check_prefixes = ['CHECK']
    result = ''
    for line in contents.split('\n'):
        if 'FileCheck' in line:
            m = PREFIX.search(line)
            if m:
                check_prefixes.extend(m.group(3).split(','))

        found_check = False
        for prefix in check_prefixes:
            if prefix in line:
                found_check = True
                break

        if not found_check or 'define' not in line:
            result += line + '\n'
            continue

        # We have a check for a function definition. Number the args.
        line = fix_string(line)
        result += line + '\n'
    return result

def main():
    print(f'Processing {sys.argv[1]}')
    f = open(sys.argv[1])
    content = f.read()
    f.close()

    content = process_file(content)

    f = open(sys.argv[1], 'w')
    f.write(content)
    f.close()

if __name__ == '__main__':
    main()