[compute]

#version 450

VERSION_DEFINES

layout(local_size_x = 4, local_size_y = 4, local_size_z = 4) in;

#define MAX_DISTANCE 100000

#define NO_CHILDREN 0xFFFFFFFF
#define GREY_VEC vec3(0.33333,0.33333,0.33333)

struct CellChildren {
    uint children[8];
};

layout(set=0,binding=1,std430) buffer CellChildrenBuffer {
    CellChildren data[];
} cell_children;


struct CellData {
    uint position; // xyz 10 bits
    uint albedo; //rgb albedo
    uint emission; //rgb normalized with e as multiplier
    uint normal; //RGB normal encoded
};

layout(set=0,binding=2,std430) buffer CellDataBuffer {
    CellData data[];
} cell_data;

layout (r8ui,set=0,binding=3) uniform restrict writeonly uimage3D sdf_tex;

layout(push_constant, binding = 0, std430) uniform Params {

    ivec3 limits;
    uint stack_size;
} params;


float distance_to_aabb(ivec3 pos, ivec3 aabb_pos, ivec3 aabb_size) {

    vec3 delta = vec3(max(ivec3(0),max(aabb_pos - pos, pos - (aabb_pos + aabb_size - ivec3(1)))));
    return length(delta);
}

void main() {

    ivec3 pos = ivec3(gl_GlobalInvocationID);

    uint stack[10]=uint[](0,0,0,0,0,0,0,0,0,0);
    uint stack_indices[10]=uint[](0,0,0,0,0,0,0,0,0,0);
    ivec3 stack_positions[10]=ivec3[](ivec3(0),ivec3(0),ivec3(0),ivec3(0),ivec3(0),ivec3(0),ivec3(0),ivec3(0),ivec3(0),ivec3(0));

    bool cell_found = false;
    bool cell_found_exact = false;
    ivec3 closest_cell_pos;
    float closest_distance = MAX_DISTANCE;
    int stack_pos = 0;

    while(true) {

        uint index = stack_indices[stack_pos];

        ivec3 cell_size = params.limits >> stack_pos;
        ivec3 cell_pos = stack_positions[stack_pos];

        if (index==0) {
            //enter stack (going down), need to set up this stack level
#if 0

            uint found = 0;
            for(uint i=0;i<8;i++) {
                if (cell_children.data[stack[stack_pos]].children[i] != NO_CHILDREN) {
                    index|=(i<<(found*3));
                    found++;
                }
            }

            index |= found<<24;
#else

            uint order[8]=uint[](0,1,2,3,4,5,6,7);

            ivec3 dir = ( pos - ( cell_pos + (cell_size >>1) ) );
            if (dir.x > 0) {
                uint tmp = order[0];
                order[0]=order[1];
                order[1]=tmp;

                tmp = order[2];
                order[2]=order[3];
                order[3]=tmp;

                tmp = order[4];
                order[4]=order[5];
                order[5]=tmp;

                tmp = order[6];
                order[6]=order[7];
                order[7]=tmp;

            }

            if (dir.y > 0) {
                uint tmp = order[0];
                order[0]=order[2];
                order[2]=tmp;

                tmp = order[1];
                order[1]=order[3];
                order[3]=tmp;

                tmp = order[4];
                order[4]=order[6];
                order[6]=tmp;

                tmp = order[5];
                order[5]=order[7];
                order[7]=tmp;

            }


            if (dir.z > 0) {
                uint tmp = order[0];
                order[0]=order[4];
                order[4]=tmp;

                tmp = order[1];
                order[1]=order[5];
                order[5]=tmp;

                tmp = order[2];
                order[2]=order[6];
                order[6]=tmp;

                tmp = order[3];
                order[3]=order[7];
                order[7]=tmp;

            }

            uint found = 0;
            for(uint i=0;i<8;i++) {
                uint idx = order[i];
                if (cell_children.data[stack[stack_pos]].children[idx] != NO_CHILDREN) {
                    index|=(idx<<(found*3));
                    found++;
                }
            }

            index |= found<<24;

#endif
        }

        //octree traversal
        uint current_child = (index>>27)&0x7;
        uint current_count = (index>>24)&0x7;

        if (current_child == current_count) {
            //go up
            if (stack_pos==0) {
                break; //done going through octree
            }
            stack_pos--;
            continue;
        }

        index=(index&((1<<27)-1))|(current_child+1); //increment current child

        stack_indices[stack_pos] = index; //store back index

        uint child_index = (index>>(current_child*3))&0x7;
        //create child AABB

        ivec3 child_cell_size = cell_size >> 1;
        ivec3 child_cell_pos = cell_pos;
        child_cell_pos+=mix(ivec3(0),child_cell_size,bvec3(uvec3(child_index&1,child_index&2,child_index&4)!=uvec3(0)));

        bool is_leaf = stack_pos == (params.stack_size-2);

        if (child_cell_pos==pos && is_leaf) {
            //we may actually end up in the exact cell.
            //if this happens, just abort
            cell_found_exact=true;
            break;
        }

        if (cell_found) {
            //discard by distance
            float distance = distance_to_aabb(pos,child_cell_pos,child_cell_size);
            if (distance >= closest_distance) {
                continue; //pointless, just test next child
            } else if (is_leaf) {
                //closer than what we have AND end of stack, save and continue
                closest_cell_pos = child_cell_pos;
                closest_distance = distance;
                continue;
            }
        } else if (is_leaf) {
            //first solid cell we find, save and continue
            closest_distance = distance_to_aabb(pos,child_cell_pos,child_cell_size);
            closest_cell_pos = child_cell_pos;
            cell_found=true;
            continue;
        }


        stack[stack_pos+1]=cell_children.data[stack[stack_pos]].children[child_index];
        stack_indices[stack_pos+1]=0; //forces regen
        stack_positions[stack_pos+1]=child_cell_pos;
        stack_pos++; //go up stack

    }

    uint dist_8;

    if (cell_found_exact) {
        dist_8=0; //equals to -1
    } else {
        float closest_distance = length(vec3(pos-closest_cell_pos));
        dist_8 = clamp(uint(closest_distance),0,254) + 1; //conservative, 0 is 1, so <1 is considered solid
    }

    imageStore(sdf_tex,pos,uvec4(dist_8));

}