# "Readable-tile" version of a fragment shader (merge B step 2): within one pass, when a later draw samples, at its own pixel, a color attachment written by an earlier draw,
# it is changed to read the color directly from the tile (framebuffer fetch); the two render passes become one, and the intermediate attachment need not be stored to main memory and loaded back just to be sampled.
# Built offline, as the entry point main0_f of the shader's own compile (fold); at runtime only function-constant specialization is done (no source compilation).
# Usage: fold(msl, bindings, om_only_ok=False) → (the source with main0_f, metadata) or (None, reason)
# Function constant numbering: FX_i (which color index the i-th rewritable sample reads, -1 = sample as usual) = i; FB_k (enable color input k) = 32 + k;
#               OM_j_k (the j-th output writes to attachment k) = 64 + 8 j + k
# Rewritable samples: .sample(sampler, select(_N, float2(0.0), isnan(_N))) on a 2D float texture, where _N = float2(in.m_X.x, in.m_X.y) (interpolated output unchanged),
# or _N = float2(in.m_X.x / in.m_X.w, in.m_X.y / in.m_X.w) (projective: 'proj', the screen position deferred decals sample their G-buffer at; screen_varyings gives the vertex side),
# or _N = float2(gl_FragCoord.x [/ render scale] * U, gl_FragCoord.y [/ render scale] * V) with U, V uniform components ('frag_pos': their [slot, vector, component]; Splatoon Raiders'
# ink stamps read their map at gl_FragCoord / 6400);
# whether the texel is really fetched at this pixel is proven at runtime (correspondence between vertex uv and position, a screen varying with the viewport over the whole target,
# or U and V = 1 / the target's size; texture size = target size); if it cannot be proven, the original version is used
import collections
import re
import dmsh_abi as abi
import dmsh_side

NFB = 8

def _once(msl, name):   # name is assigned exactly once in the source (its declaration): spirv-cross also declares variables first and assigns them in branches
    # Found as text, with a look at the character before: a pattern that starts with a lookbehind has no literal prefix to skip to, so re tries every character (4 ms each on Pokopia's 220 KB sources,
    # and 24 calls a pixel shader made this 38% of the builder's Python, under the GIL, GitHub issue #1036)
    text, n = name + ' = ', 0
    i = msl.find(text)
    while i >= 0:
        if not i or not (msl[i - 1] in '._' or msl[i - 1].isalnum()): n += 1
        i = msl.find(text, i + len(text))
    return n == 1

def _proj(msl, var, comp, locn):
    """var = in.m_X.<comp> / in.m_X.w, written as a product with 1.0 / w or as a quotient (each name assigned once) → m_X, or None"""
    def src(n, c):   # float n = in.m_X.c; → m_X
        m = re.search(r'float %s = in\.(m_\d+)\.%s;' % (re.escape(n), c), msl)
        return m.group(1) if m and _once(msl, n) and m.group(1) in locn else None
    if not _once(msl, var): return None
    m = re.search(r'float %s = (_\d+) \* (_\d+);' % re.escape(var), msl)
    if m:
        for x, r in ((m.group(1), m.group(2)), (m.group(2), m.group(1))):
            d = re.search(r'float %s = 1\.0 / (_\d+);' % re.escape(r), msl)
            if d and _once(msl, r) and src(x, comp) and src(x, comp) == src(d.group(1), 'w'): return src(x, comp)
        return None
    m = re.search(r'float %s = (_\d+) / (_\d+);' % re.escape(var), msl)
    return src(m.group(1), comp) if m and src(m.group(1), comp) == src(m.group(2), 'w') else None

_UNI = re.compile(r'\(\*spvDescriptorSet0\.fp_c(\d+)\)\._m0\[(\d+)\]\.([xyzw])')
_SCALE = '(*spvDescriptorSet0.support_buffer)._m5[0].x'   # the support buffer's render scale (Ryujinx's layout), which the host fills with 1 (nvn_mtl.m support buffer)

def _frag_pos(msl, var, comp):
    """var = gl_FragCoord.<comp> [/ render scale] * U (either order), U a fragment uniform component, each name assigned once → [uniform slot, vector, component] (fp_cN is slot N − 3,
    as dmsh_build.py raster_facts reads vp_cN), or None"""
    def one(n, pat):
        m = re.search(r'float %s = %s;' % (re.escape(n), pat), msl)
        return m if m and _once(msl, n) else None
    def pixel(n):   # n = gl_FragCoord.<comp>, or that over the render scale
        if one(n, r'gl_FragCoord\.%s' % comp): return True
        q = one(n, r'(_\d+) / (_\d+)')
        return bool(q and one(q.group(1), r'gl_FragCoord\.%s' % comp) and one(q.group(2), re.escape(_SCALE)))
    m = one(var, r'(\S+) \* (\S+)')
    if not m: return None
    for a, u in ((m.group(1), m.group(2)), (m.group(2), m.group(1))):
        k = _UNI.fullmatch(u)
        if k and int(k.group(1)) >= 3 and re.fullmatch(r'_\d+', a) and pixel(a): return [int(k.group(1)) - 3, int(k.group(2)), 'xyzw'.index(k.group(3))]
    return None

# ---- Screen varyings (vertex side of a projective sample) ----
# Interpolated output locn K is the clip → screen varying when, in every vertex entry point (main0, main0_z, ...), its .x = (P.x + P.w) / 2, .y = (P.w − P.y) / 2 and .w = P.w,
# P being the final gl_Position: in.m_K.xy / in.m_K.w is then this fragment's position over the viewport (u from the left, v from the top, the convention of nvn_mtl.m uv_proof).
# Proven on the source: the final assignments of those components are at the function's top level after any branch, the main function returns only at its end, and every
# variable they read is assigned once, as a linear form of P's components with constant coefficients (+, -, *, / by a constant, fma). Breath of the Wild's and Tears of the
# Kingdom's G-buffer decals (095cd6ad97fd779a, bb19fab2f950e587)
_TOK = re.compile(r'\s*(?:(out\.gl_Position\.[xyzw])|(_\d+)|(\d+\.\d*(?:e[+-]?\d+)?|\d+)f?|(fma)|([-+*/(),]))')

def _lin(expr, env):
    """expr as {symbol: coefficient} ('1' the constant), or None: env(name) gives a variable's form"""
    toks, pos = [], 0
    while pos < len(expr):
        m = _TOK.match(expr, pos)
        if not m or m.end() == pos: return None
        toks.append(m.group(1) or m.group(2) or ('#' + m.group(3) if m.group(3) else None) or m.group(4) or m.group(5)); pos = m.end()
    toks = [t for t in toks if t]; i = [0]
    def add(a, b, s=1.0):
        r = dict(a)
        for k_, v in b.items(): r[k_] = r.get(k_, 0.0) + s * v
        return {k_: v for k_, v in r.items() if v}
    def const(a): return set(a) <= {'1'}
    def scale(a, c): return {k_: v * c for k_, v in a.items() if v * c}
    def peek(): return toks[i[0]] if i[0] < len(toks) else None
    def take(t=None):
        x = peek()
        if x is None or (t and x != t): raise ValueError
        i[0] += 1; return x
    def atom():
        t = take()
        if t == '-': return scale(unary(), -1.0)
        if t == '(': v = add_(); take(')'); return v
        if t == 'fma':
            take('('); a = add_(); take(','); b = add_(); take(','); c = add_(); take(')')
            return add(mul(a, b), c)
        if t.startswith('#'): return {'1': float(t[1:])} if float(t[1:]) else {}
        v = env(t)
        if v is None: raise ValueError
        return v
    def unary(): return atom()
    def mul(a, b):
        if const(a): return scale(b, a.get('1', 0.0))
        if const(b): return scale(a, b.get('1', 0.0))
        raise ValueError
    def mul_():
        v = unary()
        while peek() in ('*', '/'):
            op = take(); w = unary()
            if op == '*': v = mul(v, w)
            elif const(w) and w.get('1'): v = scale(v, 1.0 / w['1'])
            else: raise ValueError
        return v
    def add_():
        v = mul_()
        while peek() in ('+', '-'):
            op = take(); v = add(v, mul_(), 1.0 if op == '+' else -1.0)
        return v
    try:
        v = add_()
        return v if peek() is None else None
    except (ValueError, ZeroDivisionError):
        return None

def _entry_screen(body, locn):
    """The locns (out struct members → K) that are screen varyings in one entry point's body (the text between its braces)"""
    if len(re.findall(r'\breturn\b', body)) != 1 or not re.search(r'\n    return out;\n$', body + '\n'): return set()
    depth, stmts = 0, []   # (depth, text) of each statement line
    lines = [ln.strip() for ln in body.split('\n')]
    for n, s in enumerate(lines):   # spirv-cross braces every block on lines of their own: a control statement without one is not read
        if re.match(r'(if|else|for|while|do|switch)\b', s) and (n + 1 >= len(lines) or lines[n + 1] != '{'): return set()
    for s in lines:
        if s == '{': depth += 1; continue
        if s == '}': depth -= 1; continue
        if s: stmts.append((depth, s))
    decl, last = {}, {}   # variable → (statement index, expression); output component → (statement index, depth, expression)
    nassign = collections.Counter(re.findall(r'(?<![\w.])(_\d+) = ', body))
    for n, (dp, s) in enumerate(stmts):
        m = re.fullmatch(r'(?:float|int|uint|bool|half) (_\d+) = (.+);', s)
        if m: decl[m.group(1)] = (n, dp, m.group(2)); continue
        m = re.fullmatch(r'out\.(gl_Position|m_\d+)\.([xyzw]) = (.+);', s)
        if m: last[(m.group(1), m.group(2))] = (n, dp, m.group(3)); continue
        m = re.fullmatch(r'out\.(gl_Position|m_\d+) = float4\((.+)\);', s)
        if m:
            parts = [p.strip() for p in m.group(2).split(',')]
            for c, p in zip('xyzw', parts if len(parts) == 4 else [None] * 4): last[(m.group(1), c)] = (n, dp, p)
    pos = {}
    for c in 'xyw':
        p = last.get(('gl_Position', c))
        if p is None or p[1] != 0 or p[2] is None: return set()
        pos[c] = p
    sym = {p[2]: 'P' + c for c, p in pos.items() if re.fullmatch(r'_\d+', p[2])}   # a final position component that is a variable: that variable is the symbol
    def env_at(at):
        memo = {}
        def env(name):
            if name.startswith('out.gl_Position.'):
                c = name[-1]; p = pos.get(c)
                return {'P' + c: 1.0} if p and p[0] < at else None   # read after its final assignment
            if name in sym: return {sym[name]: 1.0}
            if name in memo: return memo[name]
            d = decl.get(name)
            if not d or d[0] >= at or nassign[name] != 1: return None
            memo[name] = None; memo[name] = _lin(d[2], env_at(d[0])); return memo[name]
        return env
    want = {'x': {'Px': 0.5, 'Pw': 0.5}, 'y': {'Py': -0.5, 'Pw': 0.5}, 'w': {'Pw': 1.0}}
    got = set()
    for mem, k in locn.items():
        ok = True
        for c, w in want.items():
            a = last.get((mem, c))
            if not a or a[1] != 0 or a[2] is None or _lin(a[2], env_at(a[0])) != w: ok = False; break
        if ok: got.add(k)
    return got

def screen_varyings(msl):
    """The interpolated outputs (locn numbers, sorted) of a vertex shader that are its clip → screen varying in every entry point that returns the output struct"""
    ost = re.search(r'struct main0_out\s*\{(.*?)\};', msl, re.S)
    locn = {m.group(1): int(m.group(2)) for m in re.finditer(r'float4 (m_\d+) \[\[user\(locn(\d+)\)\]\]', ost.group(1) if ost else '')}
    entries = list(re.finditer(r'\nvertex main0_out \w+\(.*?\)\n\{\n(.*?)\n\}\n', msl, re.S))
    if not locn or not entries: return []
    got = set(locn.values())
    for e in entries: got &= _entry_screen(e.group(1), locn)
    return sorted(got)

def fold(msl, bindings, om_only_ok=False):
    """The readable-tile version main0_f added to msl (a pixel shader's source, or the compile unit dmsh_side.unit made of it), to be cut from the one compile of the shader's other entry
    points (dmsh_build.py ps_libraries); it was a source and a compile of its own. main0's body is dm_body_main0 (dmsh_side.shared_body), and each rewritable sample in it becomes
    (fx_i == k ? fb_k : the sample), fx_i and fb0-fb7 parameters of the body: main0_f passes its function constants FX_i and the tile's colours, every other caller -1 and zeros, which
    the always-inline body folds back to the plain sample, so main0 and the side entry points stay functions without function constants. main0_f writes each output to the
    attachments its OM_j_k name. Only the body is read and rewritten: a copy of a sample in another function would have no fb_k (483 of Breath of the Wild's variants once failed to
    compile on main0_u's copy: use of undeclared identifier 'fb0'). → (the source with main0_f, {'calls', 'mixed_regs', 'outputs'}), or (None, why)"""
    src = dmsh_side.shared_body(msl)
    if src is None: return None, 'main function signature not recognized'
    try: where, e = dmsh_side.body(src), dmsh_side.entry(src, 'main0')
    except ValueError as x: return None, str(x)
    if not where or not e: return None, 'main function signature not recognized'
    start, head, end = where; text = src[head:end]
    tex = {m.group(1) for m in re.finditer(r'texture2d<float> (\w+) \[\[id\(\d+\)\]\];', src)}
    locn = {m.group(1): int(m.group(2)) for m in re.finditer(r'float4 (m_\d+) \[\[user\(locn(\d+)\)\]\];', src.split('struct main0_in', 1)[-1].split('};', 1)[0])}
    reg = {b['msl_name']: b['reg'] for b in bindings if b.get('kind') == abi.KIND_TEX and b.get('msl_name') and b.get('reg') is not None}
    # A sample the border colour correction wrapped (dmsh_build.py border_emu) is taken whole: a tile read replaces the corrected sample, never feeds the correction
    pat = re.compile(r'(dm_bs\()?spvDescriptorSet0\.(\w+)\.sample\(spvDescriptorSet0\.(\w+)Smplr, select\((_\d+), float2\(0\.0\), isnan\(\4\)\)\)'
                     r'(?(1), spvDescriptorSet0\.%s, \d+u, spvDescriptorSet0\.\2, spvDescriptorSet0\.\3Smplr%s, select\(\4, float2\(0\.0\), isnan\(\4\)\)\))' % (abi.BORDER_MEMBER, abi.BORDER_TWIN))
    calls, out_parts, pos = [], [], 0
    for m in pat.finditer(text):
        member, var = m.group(2), m.group(4)
        if member not in tex or member != m.group(3) or member not in reg: continue
        d = re.search(r'float2 %s = float2\((_\d+), (_\d+)\);' % var, text)
        a = d and re.search(r'float %s = in\.(m_\d+)\.x;' % d.group(1), text); b = d and re.search(r'float %s = in\.(m_\d+)\.y;' % d.group(2), text)
        if a and b and a.group(1) == b.group(1) and a.group(1) in locn: how = {'varying_locn': locn[a.group(1)]}
        else:
            pa, pb = (_proj(text, d.group(1), 'x', locn), _proj(text, d.group(2), 'y', locn)) if d else (None, None)
            fx, fy = (_frag_pos(text, d.group(1), 'x'), _frag_pos(text, d.group(2), 'y')) if d else (None, None)
            if pa and pa == pb: how = {'varying_locn': locn[pa], 'proj': 1}
            elif fx and fy: how = {'frag_pos': [fx, fy]}
            else: continue
        i = len(calls)
        if i >= 32: break
        sel = ''.join(f'fx{i} == {k} ? fb{k} : ' for k in range(NFB))
        out_parts.append(text[pos:m.start()]); out_parts.append(f'({sel}{m.group(0)})'); pos = m.end()
        calls.append({'fc': i, 'tex_reg': reg[member], **how})
    if not calls and not om_only_ok:
        return None, 'no rewritable sample'
    # Mixed-use slots: the same texture also has other (non-rewritable) uses - at runtime this texture must not be switched to tile reads (the rewritten reads and the unrewritten ones would not see the same data)
    conv = collections.Counter(c['tex_reg'] for c in calls)
    mixed = sorted({reg[mb] for mb in reg if mb in tex and text.count(f'spvDescriptorSet0.{mb}.') > conv[reg[mb]]})
    om = re.search(r'struct main0_out\s*\{(.*?)\};', src, re.S)
    if not om: return None, 'no output struct'
    outs = [(t, n, int(c)) for t, n, c in re.findall(r'(\w+) (m_\d+) \[\[color\((\d+)\)\]\];', om.group(1))]
    if not outs or len(outs) > 4 or re.sub(r'\w+ m_\d+ \[\[color\(\d+\)\]\];', '', om.group(1)).strip(): return None, 'the output struct is not plain color outputs (or has more than 4)'
    rest = src[end:]
    if calls:   # the body's new parameters, and every call of it (the definition comes first) passing what turns the tile read off
        params = [f'int fx{i}' for i in range(len(calls))] + [f'float4 fb{k}' for k in range(NFB)]
        off = ', '.join(['-1'] * len(calls) + ['float4(0.0)'] * NFB)
        sig = src[start:head - 1]; sig += (', ' if not sig.endswith('(') else '') + ', '.join(params) + ')'
        parts, at = [], 0
        for m in re.finditer(r'\bdm_body_main0\(', rest):
            close = dmsh_side._close(rest, m.end(), '(', ')') - 1
            parts += [rest[at:close], (', ' if rest[m.end():close].strip() else '') + off]; at = close
        rest = ''.join(parts) + rest[at:]
    else: sig = src[start:head]
    consts = ''.join(f'constant int FX_{c["fc"]} [[function_constant({c["fc"]})]];\n' for c in calls) + \
             ''.join(f'constant bool FB_{k} [[function_constant({32 + k})]];\n' for k in range(NFB)) + \
             ''.join(f'constant bool OM_{j}_{k} [[function_constant({64 + 8 * j + k})]];\n' for j in range(len(outs)) for k in range(NFB))
    lead = src[:start]
    if 'using namespace metal;\n' not in lead: return None, 'no using namespace metal'
    lead = lead.replace('using namespace metal;\n', 'using namespace metal;\n\n' + consts, 1)
    # main0_f: main0's parameters and the tile's colours; a colour the pipeline does not read (FB_k false) is not passed to the body
    members = ''.join(f'    {t} {n}_at{k} [[color({k}), function_constant(OM_{j}_{k})]];\n' for j, (t, n, c) in enumerate(outs) for k in range(NFB))
    fbs = [f'float4 fb{k} [[color({k}), function_constant(FB_{k})]]' for k in range(NFB)]
    args = ', '.join(([e['args']] if e['args'] else []) + ([f'FX_{i}' for i in range(len(calls))] + [f'_fb{k}' for k in range(NFB)] if calls else []))
    load = ''.join(f'    float4 _fb{k} = float4(0.0);\n    if (FB_{k}) _fb{k} = fb{k};\n' for k in range(NFB)) if calls else ''
    put = ''.join(f'    if (OM_{j}_{k}) out.{n}_at{k} = r.{n};\n' for j, (t, n, c) in enumerate(outs) for k in range(NFB))
    main0_f = (f'\nstruct main0_out_f\n{{\n{members}}};\n\n{e["attr"]}fragment main0_out_f main0_f({", ".join(e["params"] + fbs)})\n{{\n'
               f'{load}    main0_out r = dm_body_main0({args});\n    main0_out_f out = {{}};\n{put}    return out;\n}}\n')
    return lead + sig + ''.join(out_parts) + text[pos:] + rest + main0_f, {'calls': calls, 'mixed_regs': mixed, 'outputs': [{'color': c, 'type': t} for t, n, c in outs]}