TLAS

struct Geometry {
            glm::vec4 signature;
            glm::vec4 vertex1_or_sphere_center;
            glm::vec4 vertex2_or_sphere_radius;
            glm::vec4 vertex3_or_nothing;
        };

        struct GeometryCollection {
            Geometry geometry[1 << expOfTwoOfMaxGeometryElementsInCollection]; // geometry slots
        };

        struct AABB {
            glm::vec4 position; // this is the point of the aabb with min x, y and z components
            glm::vec4 dimensions; // this is the amount to be added to aabbVertex to obtain the point with max x, y and z in the aabb
        };

        struct alignas(16) NodeData {
            AABB aabb;
            alignas(sizeof(glm::uint32)) glm::uint32 left; // this can also be the content if right is zero
            alignas(sizeof(glm::uint32))glm::uint32 right; // this is zero if the node is a leaf
        };

        struct BLAS {
            glm::mat4 ModelMatrix;

            GeometryCollection collection[1 << expOfTwoOfMaxCollectionElementsInBLAS];

            NodeData tree[((1 << expOfTwoOfMaxCollectionElementsInBLAS) * 2) - 1];
        };

        struct TLAS {
            glm::mat4 ViewMatrix;
            NodeData tree[((1 << expOfTwoOfMaxBLASElementsInTLAS) * 2) - 1];
            BLAS blas[1 << expOfTwoOfMaxBLASElementsInTLAS];
        };


/******************************************************************************************/

#version 450

#define COMP_SIZE 32

#define expOfTwoOfMaxGeometryElementsInCollection 3
#define expOfTwoOfMaxCollectionElementsInBLAS 5
#define expOfTwoOfMaxBLASElementsInTLAS 5

const float PI = 3.14159265359;
const float infinity = 1. / 0.;
const mat4 identityTransform = mat4(1);

// This gets adjusted on each call by the OpenGL core
layout(local_size_x = COMP_SIZE, local_size_y = COMP_SIZE, local_size_z = 1) in;

layout(rgba32f, binding = 0) uniform image2D img_output;

uniform uint width; // Render target surface width
uniform uint height; // Render target surface height

struct Geometry {
    vec4 signature;
    vec4 vertex1_or_sphere_center;
    vec4 vertex2_or_sphere_radius;
    vec4 vertex3_or_nothing;
};

struct GeometryCollection {
    Geometry geometry[1 << expOfTwoOfMaxGeometryElementsInCollection]; // geometry slots
};

struct AABB {
    vec4 position; // this is the point of the aabb with min x, y and z components
    vec4 dimensions; // this is the amount to be added to aabbVertex to obtain the point with max x, y and z in the aabb
};

struct NodeData {
    AABB aabb;
    uint left; // this can also be the content if right is zero
    uint right; // this is zero if the node is a leaf
};

struct BLAS {
    mat4 ModelMatrix;

    GeometryCollection collection[1 << expOfTwoOfMaxCollectionElementsInBLAS];

    NodeData tree[((1 << expOfTwoOfMaxCollectionElementsInBLAS) * 2) - 1];
};

struct TLAS {
    mat4 ViewMatrix;
    NodeData tree[((1 << expOfTwoOfMaxBLASElementsInTLAS) * 2) - 1];
    BLAS blas[1 << expOfTwoOfMaxBLASElementsInTLAS];
};

// Output buffer object to be written
layout(std430, binding = 3) buffer scene {
    TLAS tlas[];
};

uniform uint shadingAlgorithm;

uniform vec3 mOrigin;
uniform vec3 mLowerLeftCorner;
uniform vec3 mHorizontal;
uniform vec3 mVertical;

/*************************************************************************************************************************
 *                                                  Algorithms                                                           *
 *************************************************************************************************************************/

struct RayGeometryIntersection {
    float dist;
    vec4 point;
    vec4 normal;
};

struct Ray {
    vec4 origin;
    vec4 direction;
};

// Empty geometry is a sphere with radius equals to 0.0 such geometry CANNOT be intersected
const Geometry emptyGeometry = Geometry(
    vec4(0, 0, 0, 0),
    vec4(0, 0, 0, 0),
    vec4(0, 0, 0, 0),
    vec4(0, 0, 0, 0)
);

const RayGeometryIntersection miss = RayGeometryIntersection(infinity, vec4(0, 0, 0, 1), vec4(0, 0, 0, 0));

bool isLeaf(NodeData node) {
    return node.right == 0;
}

bool hasMissed(RayGeometryIntersection test) {
    return isinf(test.dist) || isnan(test.dist);
}

vec3 rayPointAt(Ray ray, float coeff) {
    return vec3(ray.origin) + coeff * vec3(ray.direction);
}

vec3 applyShading(RayGeometryIntersection isect, uint algorithm) {
    if (algorithm == 0) { // hit or miss
        return hasMissed(isect) ? vec3(0,0,0) : vec3(1,1,1);
    } else if (algorithm == 1) {
        // Here the camera position is vec3(0, 0, 0) because all of this is in camera space
        return vec3(max(0, dot(isect.normal, normalize(/*mPosition*/ -isect.point))));
    }

    return vec3(0.5, 0.5, 0.5); // this is to get attention as this should never verify
}

vec4 getInvDirection(Ray ray) {
    return vec4(1.0/ray.direction.x, 1.0/ray.direction.y, 1.0/ray.direction.z, 0);
}

AABB transformAABB(AABB starting, mat4 transformMatrix) {
    vec4 v[8] = {
        vec4(starting.position.x, starting.position.y, starting.position.z, 1),
        vec4(starting.position.x, starting.position.y, starting.position.z, 1),
        vec4(starting.position.x, starting.position.y, starting.position.z, 1),
        vec4(starting.position.x, starting.position.y, starting.position.z, 1),
        vec4(starting.position.x, starting.position.y, starting.position.z, 1),
        vec4(starting.position.x, starting.position.y, starting.position.z, 1),
        vec4(starting.position.x, starting.position.y, starting.position.z, 1),
        vec4(starting.position.x, starting.position.y, starting.position.z, 1),
    };

    float maxX = 1.0/0.0, maxY = 1.0/0.0, maxZ=1.0/0.0, minX=-1.0/0.0, minY=-1.0/0.0, minZ=-1.0/0.0;

    for (uint i = 0; i < 8; ++i) {
        maxX = (v[i].x > maxX) ? v[i].x : maxX;
        maxY = (v[i].y > maxY) ? v[i].y : maxY;
        maxZ = (v[i].z > maxZ) ? v[i].z : maxZ;
        minX = (v[i].x < minX) ? v[i].x : minX;
        minY = (v[i].y < minY) ? v[i].y : minY;
        minZ = (v[i].z < minZ) ? v[i].z : minZ;
    }

    return AABB(vec4(minX, minY, minZ, 1), vec4(maxX-minX, maxY-minY, maxZ-minZ, 0));
}

struct raySigns {
    bool[3] signs;
};

bool intersectAABB(Ray ray, AABB aabb, mat4 transformMatrix) {
    AABB transformedAABB = /*(transformMatrix != identityTransform) ? */transformAABB(aabb, transformMatrix) /*: *this*/;

    float tmin, tmax, tymin, tymax, tzmin, tzmax;

    vec3 orig = ray.origin.xyz;
    vec3 invdir = getInvDirection(ray).xyz;

    vec3 bounds[2] = {
        transformedAABB.position.xyz,
        (transformedAABB.position.xyz) + (transformedAABB.dimensions.xyz)
    };

    const int[3] raySigns = int[3](
        (invdir.x < 0) ? 1 : 0,
        (invdir.y < 0) ? 1 : 0,
        (invdir.z < 0) ? 1 : 0
    );

    tmin = (bounds[raySigns[0]].x - orig.x) * invdir.x;
    tmax = (bounds[1 - raySigns[0]].x - orig.x) * invdir.x;
    tymin = (bounds[raySigns[1]].y - orig.y) * invdir.y;
    tymax = (bounds[1 - raySigns[1]].y - orig.y) * invdir.y;

    if ((tmin > tymax) || (tymin > tmax))
        return false;
    if (tymin > tmin)
        tmin = tymin;
    if (tymax < tmax)
        tmax = tymax;

    tzmin = (bounds[raySigns[2]].z - orig.z) * invdir.z;
    tzmax = (bounds[1 - raySigns[2]].z - orig.z) * invdir.z;

    if ((tmin > tzmax) || (tzmin > tmax))
        return false;
    if (tzmin > tmin)
        tmin = tzmin;
    if (tzmax < tmax)
        tmax = tzmax;

    return true;
}

RayGeometryIntersection intersect(Ray ray, Geometry geometry, mat4 transformMatrix, float minDistance, float maxDistance) {
    if (geometry.signature.w == 0.0) { // test ray-sphere intersection
        const vec3 center = vec3(transformMatrix * geometry.vertex1_or_sphere_center);
        const float radius = geometry.vertex2_or_sphere_radius.x;

        if (radius == 0) return miss;

        const vec3 origin = vec3(ray.origin);
        const vec3 direction = vec3(ray.direction);

        vec3 oc = origin - center;

        const float a = dot(direction, direction);
        const float b = dot(oc, direction);
        const float c = dot(oc, oc) - radius * radius;
        const float discriminant = b * b - a * c;

        if (discriminant > 0.0) { // if delta > 0 then we have two intersections (one for each side of the sphere)
            const float squareRoot = sqrt(discriminant);

            // x0 and x1 are the distances to the origin of the ray
            float x0 = (-b - squareRoot) / a, x1 = (-b + squareRoot) / a;

            // Use x0 and x1 to calculate intersection points
            vec3 point_x0 = rayPointAt(ray, x0), point_x1 = rayPointAt(ray, x1); // so if I use that distance as a coefficient I obtain the intersection point

            // Use intersecting points to calculate the surface normal at that point
            vec3 normal_x0 = (point_x0 - center) / radius, normal_x1 = (point_x1 - center) / radius; // and the normal for a point p is (point - center)/radius

            // No valid intersection? return a miss.
            if (((x0 <= minDistance) || (x0 >= maxDistance)) && ((x1 <= minDistance) || (x1 >= maxDistance))) {
                return miss;
            }

            // Choose the closest intersection point
            return (((x0 > minDistance) && (x0 < maxDistance)) && (x0 <= x1)) || ((x1 <= minDistance) || (x1 >= maxDistance)) ?
                RayGeometryIntersection(x0, vec4(point_x0, 1), vec4(normal_x0, 0)) : RayGeometryIntersection(x1, vec4(point_x1, 1), vec4(normal_x1, 0));
        }
    } else if (geometry.signature.w == 1.0) { // test ray-triangle intersection

    }

    return miss;
}

RayGeometryIntersection intersect(Ray ray, GeometryCollection collection, mat4 transformMatrix, float minDistance, float maxDistance) {
    RayGeometryIntersection bestHit = miss;

    for (uint i = 0; i < collection.geometry.length(); ++i) {
        // Execute the ray-geometry intersection algorithm
        const RayGeometryIntersection currentIntersectionInfo = intersect(ray, collection.geometry[i], transformMatrix, minDistance, maxDistance);

        // Check if this is a better hit than the former one
        if (currentIntersectionInfo.dist < bestHit.dist) {
            bestHit = currentIntersectionInfo;
        }
    }

    return bestHit;
}

RayGeometryIntersection intersect(Ray ray, BLAS blas, mat4 transformMatrix, float minDistance, float maxDistance) {
    // Adjust transformation matrix to consider the BLAS model matrix
    transformMatrix = transformMatrix * blas.ModelMatrix;

    RayGeometryIntersection bestHit = miss;

    int currentDepth = 0;
    uint currentPath = 0;

    while ((currentDepth < 32) && (currentDepth >= 0)) {
        uint currentNodeIndex = 0;

        // This for cycles is used to reach the interesting node
        for (uint i = 0; i < currentDepth; ++i) {
            currentNodeIndex = ((currentPath & (1 << currentDepth)) == 0) ? blas.tree[currentNodeIndex].left : blas.tree[currentNodeIndex].left;
        }

        if ((!isLeaf(blas.tree[currentNodeIndex])) && (intersectAABB(ray, blas.tree[currentNodeIndex].aabb, transformMatrix))) {
            // Go one level deeper
            currentDepth += 1;
        } else {
            // This is a leaf that should be tested...
            if ((isLeaf(blas.tree[currentNodeIndex])) && (intersectAABB(ray, blas.tree[currentNodeIndex].aabb, transformMatrix))) {
                RayGeometryIntersection currentHit = intersect(ray, blas.collection[blas.tree[currentNodeIndex].left], transformMatrix, minDistance, maxDistance);

                if (currentHit.dist < bestHit.dist) bestHit = currentHit;
            }

            // in either case: subtree discarded or visited leaf the algorithm should reduce the depth and change the path
            currentDepth -= 1;

            while (((currentPath & (1 << currentDepth)) == 1) && (currentDepth >= 0)) {
                //set the bit at currentDepth position to 0
                currentPath &= ~(1 << currentDepth);

                // and decrease currentDepth
                currentDepth -= 1;
            }

            if (currentDepth >= 0) {
                // set the bit at currentDepth position to 1 (the next visit will choose another path)
                currentPath |= (1 << currentDepth);
            } else {
                // the visited node was the very last one (the right-most) and the iteration should come to an end
                break;
            }
        }
    }

    return bestHit;
}

RayGeometryIntersection intersect(Ray ray, TLAS tlas, mat4 transformMatrix, float minDistance, float maxDistance) {
    // Adjust transformation matrix to consider the BLAS model matrix
    transformMatrix = transformMatrix * tlas.ViewMatrix;

    RayGeometryIntersection bestHit = miss;

    int currentDepth = 0;
    uint currentPath = 0;

    while ((currentDepth < 32) && (currentDepth >= 0)) {
        uint currentNodeIndex = 0;

        // This for cycles is used to reach the interesting node
        for (uint i = 0; i < currentDepth; ++i) {
            currentNodeIndex = ((currentPath & (1 << currentDepth)) == 0) ? tlas.tree[currentNodeIndex].left : tlas.tree[currentNodeIndex].left;
        }

        if ((!isLeaf(tlas.tree[currentNodeIndex])) && (intersectAABB(ray, tlas.tree[currentNodeIndex].aabb, transformMatrix))) {
            // Go one level deeper
            currentDepth += 1;
        } else {
            // This is a leaf that should be tested...
            if ((isLeaf(tlas.tree[currentNodeIndex])) && (intersectAABB(ray, tlas.tree[currentNodeIndex].aabb, transformMatrix))) {
                RayGeometryIntersection currentHit = intersect(ray, tlas.blas[tlas.tree[currentNodeIndex].left], transformMatrix, minDistance, maxDistance);

                if (currentHit.dist < bestHit.dist) bestHit = currentHit;
            }

            // in either case: subtree discarded or visited leaf the algorithm should reduce the depth and change the path
            currentDepth -= 1;

            while (((currentPath & (1 << currentDepth)) == 1) && (currentDepth >= 0)) {
                //set the bit at currentDepth position to 0
                currentPath &= ~(1 << currentDepth);

                // and decrease currentDepth
                currentDepth -= 1;
            }

            if (currentDepth >= 0) {
                // set the bit at currentDepth position to 1 (the next visit will choose another path)
                currentPath |= (1 << currentDepth);
            } else {
                // the visited node was the very last one (the right-most) and the iteration should come to an end
                break;
            }
        }
    }

    return bestHit;
}

Ray generateCameraRay(float s, float t) {
    return Ray(vec4(mOrigin, 1), vec4(mLowerLeftCorner + s * mHorizontal + t * mVertical - mOrigin, 0));
}

void main() {
    // base pixel colour for image
    vec4 pixel = vec4(0, 0, 0, 0);
    // get index in global work group i.e x,y position
    ivec2 pixel_coords = ivec2(gl_GlobalInvocationID.xy);

    // Avoid calculating useless pixels
    if ((gl_GlobalInvocationID.x >= width) || (gl_GlobalInvocationID.y >= height)) {
        return;
    }

    const float u = float(gl_GlobalInvocationID.x) / float(width);
    const float v = float(gl_GlobalInvocationID.y) / float(height);

    // Generate camera ray
    Ray cameraRay = generateCameraRay(u, v);

    pixel = vec4(applyShading(intersect(cameraRay, tlas[0], identityTransform, 0.001, 1000.0), shadingAlgorithm), 1.0);

    // output to a specific pixel in the image
    imageStore(img_output, pixel_coords, pixel);
}