#!/usr/bin/env python3
"""Condense a ptx2glsl.py translation for reading: drops declarations, prints float constants as decimals and
substitutes registers that are assigned once and used once into their use. The result is not meant to compile.

usage: condense.py KERNEL.comp > KERNEL.txt
"""
import re, struct, sys
src = open(sys.argv[1]).read().split('\n')
start = next(i for i, l in enumerate(src) if l.startswith('void main'))
body = [l for l in src[start:] if not re.match(r'\s*(precise )?(uint|float|bool|uint64_t|float16_t|f16vec2|int) [a-z0-9_]+ = [^;]*(, [a-z0-9_]+ = [^;,]*)*;\s*$', l) or '(' in l.split('=', 1)[1][:3]]
def fl(m):
    v = struct.unpack('<f', struct.pack('<I', int(m.group(1), 16)))[0]
    return repr(round(v, 7)) if abs(v) < 1e7 and (abs(v) > 1e-5 or v == 0) else f'{v:.6g}'
body = [re.sub(r'uintBitsToFloat\(0x([0-9A-Fa-f]{8})u\)', fl, l) for l in body]
body = [re.sub(r'\b(\d+)u\b', r'\1', l) for l in body]
body = [re.sub(r' \((\d+) & 31\)', r' \1', l) for l in body]
tok = re.compile(r'\b(?:r|f|rs|rd|p)\d+\b')
for _ in range(12):
    assigns, uses = {}, {}
    for i, l in enumerate(body):
        m = re.match(r'(\s*)((?:r|f|rs|rd|p)\d+) = (.*);\s*$', l)
        rhs = l
        if m:
            assigns.setdefault(m.group(2), []).append(i); rhs = m.group(3)
        for t in tok.findall(rhs): uses.setdefault(t, []).append(i)
    changed = False
    for name, at in assigns.items():
        u = uses.get(name, [])
        if len(at) != 1 or len(u) != 1 or u[0] <= at[0]: continue
        a, b = at[0], u[0]
        if body[a] is None or body[b] is None: continue
        ind = lambda s: len(s) - len(s.lstrip())
        # only within straight-line code at the same nesting, and short enough to stay readable
        if any(body[k] is not None and ('{' in body[k] or '}' in body[k]) for k in range(a, b)): continue
        m = re.match(r'(\s*)((?:r|f|rs|rd|p)\d+) = (.*);\s*$', body[a])
        expr = m.group(3)
        if len(expr) > 90: continue
        # the expression's inputs must not be reassigned in between
        ins = set(tok.findall(expr))
        if any(body[k] is not None and re.match(r'\s*(%s) = ' % '|'.join(map(re.escape, ins)) if ins else r'$^', body[k]) for k in range(a + 1, b)): continue
        simple = re.fullmatch(r'[\w.]+|\w+\([^()]*\)', expr)
        body[b] = re.sub(r'\b%s\b' % name, expr if simple else f'({expr})', body[b], count=1)
        body[a] = None; changed = True
    body = [l for l in body if l is not None]
    if not changed: break
for a, b in ((r'uint\(uint\(([^()]*)\)\)', r'\1'), (r'\buint\(((?:r|rs)\d+)\)', r'\1'), (r'uint64_t\((rd\d+)\)', r'\1')):
    body = [re.sub(a, b, l) for l in body]
# scripts/check_redistributable.py refuses files carrying this line: the input was NVIDIA's code.
print(f'// Generated by vulkan/condense.py from NVIDIA PTX ({sys.argv[1]}); do not commit or redistribute.')
print('\n'.join(l.rstrip() for l in body if l.strip()))