// SPDX-License-Identifier: Apache-2.0 OR MIT OR Unlicense

// The main phase for the bbox union computation

#version 450

#define LG_N_SEQ 0
#define N_SEQ (1 << LG_N_SEQ)
#define LG_WG_SIZE 9
#define WG_SIZE (1 << LG_WG_SIZE)
#define PART_SIZE (WG_SIZE * N_SEQ)

layout(local_size_x = WG_SIZE, local_size_y = 1) in;

// The bicyclic monoid
struct Bic {
    uint a;
    uint b;
};

struct BicBbox {
    Bic bic;
    uint pad2;
    uint pad3;
    vec4 bbox;
};

Bic bic_combine(Bic x, Bic y) {
    uint m = min(x.b, y.a);
    return Bic(x.a + y.a - m, x.b + y.b - m);
}

// Node type values
#define OPEN_PAREN 0
#define CLOSE_PAREN 1

struct Node {
    uint node_type;
    uint pad1;
    uint pad2;
    uint pad3;
    vec4 bbox;
};

#define BBOX_IDENTITY vec4(1e9, 1e9, -1e9, -1e9)

vec4 bbox_union(vec4 a, vec4 b) {
    return vec4(min(a.xy, b.xy), max(a.zw, b.zw));
}

layout(binding = 0) readonly buffer InBuf {
    Node[] inbuf;
};

layout(binding = 1) readonly buffer BicBuf {
    BicBbox[] bicbuf;
};

layout(binding = 2) readonly buffer StackBuf {
    vec4[] stack;
};

layout(binding = 3) buffer OutBuf {
    vec4[] outbuf;
};

shared Bic sh_bic[WG_SIZE * 2 - 2];
shared vec4 sh_stack[WG_SIZE];
shared vec4 sh_bbox[WG_SIZE * 2 - 2];

void main() {
    uint th = gl_LocalInvocationID.x;
    // materialize stack snapshot as of the start of this partition
    Bic bic = Bic(0, 0);
    vec4 bbox = BBOX_IDENTITY;
    if (th < gl_WorkGroupID.x) {
        bic = bicbuf[th].bic;
        bbox = bicbuf[th].bbox;
    }
    sh_bic[th] = bic;
    sh_bbox[th] = bbox;
    for (uint i = 0; i < LG_WG_SIZE; i++) {
        barrier();
        uint other_ix = th + (1u << i);
        if (other_ix < WG_SIZE) {
            bic = bic_combine(bic, sh_bic[other_ix]);
            bbox = bbox_union(bbox, sh_bbox[other_ix]);
        }
        barrier();
        sh_bic[th] = bic;
        sh_bbox[th] = bbox;
    }

    barrier();
    uint size = sh_bic[0].b;
    // binary search in stack to do stream compaction
    uint sp = PART_SIZE - 1 - th;
    uint ix = 0;
    for (uint i = 0; i < LG_WG_SIZE; i++) {
        uint probe = ix + (uint(PART_SIZE / 2) >> i);
        if (sp < sh_bic[probe].b) {
            ix = probe;
        }
    }
    // ix is the largest value such that sp < sh_bic[ix].b (if any)
    uint b = sh_bic[ix].b;
    if (sp < b) {
        vec4 bbox = stack[ix * PART_SIZE + b - sp - 1];
        if (ix + 1 < WG_SIZE) {
            bbox = bbox_union(bbox, sh_bbox[ix + 1]);
        }
        sh_stack[th] = bbox;
    }
    barrier();

    // Do tree reduction of bicyclic semigroups and bounding boxes
    Node inp = inbuf[gl_GlobalInvocationID.x];
    uint node_type = inp.node_type;
    bic = Bic(uint(node_type == CLOSE_PAREN), uint(node_type == OPEN_PAREN));
    sh_bic[th] = bic;
    sh_bbox[th] = inp.bbox;
    uint inbase = 0;
    for (uint i = 0; i < LG_WG_SIZE - 1; i++) {
        uint outbase = 2 * WG_SIZE - (1u << (LG_WG_SIZE - i));
        barrier();
        if (th < (1u << (LG_WG_SIZE - 1 - i))) {
            sh_bic[outbase + th] = bic_combine(sh_bic[inbase + th * 2], sh_bic[inbase + th * 2 + 1]);
            sh_bbox[outbase + th] = bbox_union(sh_bbox[inbase + th * 2], sh_bbox[inbase + th * 2 + 1]);
        }
        inbase = outbase;
    }
    barrier();

    // Search for predecessor node and aggregate bounding boxes
    ix = th;
    bbox = inp.bbox;
    bic = Bic(0, 0);
    if (node_type == CLOSE_PAREN) {
        // only aggregate bbox for nodes with children; ascribe to close paren
        // Note: if we also maintained the link to the open paren, we could
        // store the bbox there. This will depend on the use case.
        uint j = 0;
        while (j < LG_WG_SIZE) {
            uint base = 2 * WG_SIZE - (2u << (LG_WG_SIZE - j));
            if (((ix >> j) & 1) != 0) {
                Bic test = bic_combine(sh_bic[base + (ix >> j) - 1], bic);
                if (test.b > 0) {
                    break;
                }
                bic = test;
                bbox = bbox_union(sh_bbox[base + (ix >> j) - 1], bbox);
                ix -= 1u << j;
            }
            j++;
        }
        if (ix > 0) {
            while (j > 0) {
                j--;
                uint base = 2 * WG_SIZE - (2u << (LG_WG_SIZE - j));
                Bic test = bic_combine(sh_bic[base + (ix >> j) - 1], bic);
                if (test.b == 0) {
                    bic = test;
                    bbox = bbox_union(sh_bbox[base + (ix >> j) - 1], bbox);
                    ix -= 1u << j;
                }
            }
        }
        // ix is the smallest value such that reduce(ix..th).b == 0
        if (ix == 0 && bic.a < size) {
            bbox = bbox_union(sh_stack[PART_SIZE - 1 - bic.a], bbox);
        }
    }
    outbuf[gl_GlobalInvocationID.x] = bbox;
}