NodeCodecQuadTreeHalfFloat.h

Go to the documentation of this file.

// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
// SPDX-FileCopyrightText: 2021 Jorrit Rouwe
// SPDX-License-Identifier: MIT
 
#pragma once
 
#include <Jolt/Core/ByteBuffer.h>
#include <Jolt/Math/HalfFloat.h>
#include <Jolt/AABBTree/AABBTreeBuilder.h>
 
JPH_NAMESPACE_BEGIN
 
class NodeCodecQuadTreeHalfFloat
{
public:
    static constexpr int                NumChildrenPerNode = 4;
 
    struct Header
    {
        Float3                          mRootBoundsMin;
        Float3                          mRootBoundsMax;
        uint32                          mRootProperties;
        uint8                           mBlockIDBits;           
        uint8                           mPadding[3] = { 0 };
    };
 
    static constexpr int                HeaderSize = sizeof(Header);
 
    static constexpr int                StackSize = 128;
 
    enum : uint32
    {
        TRIANGLE_COUNT_BITS             = 4,
        TRIANGLE_COUNT_SHIFT            = 28,
        TRIANGLE_COUNT_MASK             = (1 << TRIANGLE_COUNT_BITS) - 1,
        OFFSET_BITS                     = 28,
        OFFSET_MASK                     = (1 << OFFSET_BITS) - 1,
        OFFSET_NON_SIGNIFICANT_BITS     = 2,
        OFFSET_NON_SIGNIFICANT_MASK     = (1 << OFFSET_NON_SIGNIFICANT_BITS) - 1,
    };
 
    struct Node
    {
        HalfFloat                       mBoundsMinX[4];         
        HalfFloat                       mBoundsMinY[4];
        HalfFloat                       mBoundsMinZ[4];
        HalfFloat                       mBoundsMaxX[4];
        HalfFloat                       mBoundsMaxY[4];
        HalfFloat                       mBoundsMaxZ[4];
        uint32                          mNodeProperties[4];     
    };
 
    static_assert(sizeof(Node) == 64, "Node should be 64 bytes");
 
    class EncodingContext
    {
    public:
        void                            PrepareNodeAllocate(const AABBTreeBuilder::Node *inNode, uint64 &ioBufferSize) const
        {
            // We don't emit nodes for leafs
            if (!inNode->HasChildren())
                return;
 
            // Add size of node
            ioBufferSize += sizeof(Node);
        }
 
        size_t                          NodeAllocate(const AABBTreeBuilder::Node *inNode, Vec3Arg inNodeBoundsMin, Vec3Arg inNodeBoundsMax, Array<const AABBTreeBuilder::Node *> &ioChildren, Vec3 outChildBoundsMin[NumChildrenPerNode], Vec3 outChildBoundsMax[NumChildrenPerNode], ByteBuffer &ioBuffer, const char *&outError) const
        {
            // We don't emit nodes for leafs
            if (!inNode->HasChildren())
                return ioBuffer.size();
 
            // Remember the start of the node
            size_t node_start = ioBuffer.size();
 
            // Fill in bounds
            Node *node = ioBuffer.Allocate<Node>();
 
            for (size_t i = 0; i < 4; ++i)
            {
                if (i < ioChildren.size())
                {
                    const AABBTreeBuilder::Node *this_node = ioChildren[i];
 
                    // Copy bounding box
                    node->mBoundsMinX[i] = HalfFloatConversion::FromFloat<HalfFloatConversion::ROUND_TO_NEG_INF>(this_node->mBounds.mMin.GetX());
                    node->mBoundsMinY[i] = HalfFloatConversion::FromFloat<HalfFloatConversion::ROUND_TO_NEG_INF>(this_node->mBounds.mMin.GetY());
                    node->mBoundsMinZ[i] = HalfFloatConversion::FromFloat<HalfFloatConversion::ROUND_TO_NEG_INF>(this_node->mBounds.mMin.GetZ());
                    node->mBoundsMaxX[i] = HalfFloatConversion::FromFloat<HalfFloatConversion::ROUND_TO_POS_INF>(this_node->mBounds.mMax.GetX());
                    node->mBoundsMaxY[i] = HalfFloatConversion::FromFloat<HalfFloatConversion::ROUND_TO_POS_INF>(this_node->mBounds.mMax.GetY());
                    node->mBoundsMaxZ[i] = HalfFloatConversion::FromFloat<HalfFloatConversion::ROUND_TO_POS_INF>(this_node->mBounds.mMax.GetZ());
 
                    // Store triangle count
                    node->mNodeProperties[i] = this_node->GetTriangleCount() << TRIANGLE_COUNT_SHIFT;
                    if (this_node->GetTriangleCount() >= TRIANGLE_COUNT_MASK)
                    {
                        outError = "NodeCodecQuadTreeHalfFloat: Too many triangles";
                        return size_t(-1);
                    }
                }
                else
                {
                    // Make this an invalid triangle node
                    node->mNodeProperties[i] = uint32(TRIANGLE_COUNT_MASK) << TRIANGLE_COUNT_SHIFT;
 
                    // Make bounding box invalid
                    node->mBoundsMinX[i] = HALF_FLT_MAX;
                    node->mBoundsMinY[i] = HALF_FLT_MAX;
                    node->mBoundsMinZ[i] = HALF_FLT_MAX;
                    node->mBoundsMaxX[i] = HALF_FLT_MAX;
                    node->mBoundsMaxY[i] = HALF_FLT_MAX;
                    node->mBoundsMaxZ[i] = HALF_FLT_MAX;
                }
            }
 
            // Since we don't keep track of the bounding box while descending the tree, we keep the root bounds at all levels for triangle compression
            for (int i = 0; i < NumChildrenPerNode; ++i)
            {
                outChildBoundsMin[i] = inNodeBoundsMin;
                outChildBoundsMax[i] = inNodeBoundsMax;
            }
 
            return node_start;
        }
 
        bool                        NodeFinalize(const AABBTreeBuilder::Node *inNode, size_t inNodeStart, uint inNumChildren, const size_t *inChildrenNodeStart, const size_t *inChildrenTrianglesStart, ByteBuffer &ioBuffer, const char *&outError)
        {
            if (!inNode->HasChildren())
                return true;
 
            Node *node = ioBuffer.Get<Node>(inNodeStart);
            for (uint i = 0; i < inNumChildren; ++i)
            {
                size_t offset;
                if (node->mNodeProperties[i] != 0)
                {
                    // This is a triangle block
                    offset = inChildrenTrianglesStart[i];
 
                    // Store highest block with triangles so we can count the number of bits we need
                    mHighestTriangleBlock = max(mHighestTriangleBlock, offset);
                }
                else
                {
                    // This is a node block
                    offset = inChildrenNodeStart[i];
                }
 
                // Store offset of next node / triangles
                if (offset & OFFSET_NON_SIGNIFICANT_MASK)
                {
                    outError = "NodeCodecQuadTreeHalfFloat: Internal Error: Offset has non-significant bits set";
                    return false;
                }
                offset >>= OFFSET_NON_SIGNIFICANT_BITS;
                if (offset > OFFSET_MASK)
                {
                    outError = "NodeCodecQuadTreeHalfFloat: Offset too large. Too much data.";
                    return false;
                }
                node->mNodeProperties[i] |= uint32(offset);
            }
 
            return true;
        }
 
        bool                        Finalize(Header *outHeader, const AABBTreeBuilder::Node *inRoot, size_t inRootNodeStart, size_t inRootTrianglesStart, const char *&outError) const
        {
            // Check if we can address the root node
            size_t offset = inRoot->HasChildren()? inRootNodeStart : inRootTrianglesStart;
            if (offset & OFFSET_NON_SIGNIFICANT_MASK)
            {
                outError = "NodeCodecQuadTreeHalfFloat: Internal Error: Offset has non-significant bits set";
                return false;
            }
            offset >>= OFFSET_NON_SIGNIFICANT_BITS;
            if (offset > OFFSET_MASK)
            {
                outError = "NodeCodecQuadTreeHalfFloat: Offset too large. Too much data.";
                return false;
            }
 
            // If the root has triangles, we need to take that offset instead since the mHighestTriangleBlock will be zero
            size_t highest_triangle_block = inRootTrianglesStart != size_t(-1)? inRootTrianglesStart : mHighestTriangleBlock;
            highest_triangle_block >>= OFFSET_NON_SIGNIFICANT_BITS;
 
            inRoot->mBounds.mMin.StoreFloat3(&outHeader->mRootBoundsMin);
            inRoot->mBounds.mMax.StoreFloat3(&outHeader->mRootBoundsMax);
            outHeader->mRootProperties = uint32(offset) + (inRoot->GetTriangleCount() << TRIANGLE_COUNT_SHIFT);
            outHeader->mBlockIDBits = uint8(32 - CountLeadingZeros(uint32(highest_triangle_block)));
            if (inRoot->GetTriangleCount() >= TRIANGLE_COUNT_MASK)
            {
                outError = "NodeCodecQuadTreeHalfFloat: Too many triangles";
                return false;
            }
 
            return true;
        }
 
    private:
        size_t                      mHighestTriangleBlock = 0;
    };
 
    class DecodingContext
    {
    public:
        inline static uint          sTriangleBlockIDBits(const Header *inHeader)
        {
            return inHeader->mBlockIDBits;
        }
 
        inline static const void *  sGetTriangleBlockStart(const uint8 *inBufferStart, uint inTriangleBlockID)
        {
            return inBufferStart + (inTriangleBlockID << OFFSET_NON_SIGNIFICANT_BITS);
        }
 
        JPH_INLINE explicit         DecodingContext(const Header *inHeader)
        {
            // Start with the root node on the stack
            mNodeStack[0] = inHeader->mRootProperties;
        }
 
        template <class TriangleContext, class Visitor>
        JPH_INLINE void             WalkTree(const uint8 *inBufferStart, const TriangleContext &inTriangleContext, Visitor &ioVisitor)
        {
            do
            {
                // Test if node contains triangles
                uint32 node_properties = mNodeStack[mTop];
                uint32 tri_count = node_properties >> TRIANGLE_COUNT_SHIFT;
                if (tri_count == 0)
                {
                    const Node *node = reinterpret_cast<const Node *>(inBufferStart + (node_properties << OFFSET_NON_SIGNIFICANT_BITS));
 
                    // Unpack bounds
                #ifdef JPH_CPU_BIG_ENDIAN
                    Vec4 bounds_minx = HalfFloatConversion::ToFloat(UVec4(node->mBoundsMinX[0] + (node->mBoundsMinX[1] << 16), node->mBoundsMinX[2] + (node->mBoundsMinX[3] << 16), 0, 0));
                    Vec4 bounds_miny = HalfFloatConversion::ToFloat(UVec4(node->mBoundsMinY[0] + (node->mBoundsMinY[1] << 16), node->mBoundsMinY[2] + (node->mBoundsMinY[3] << 16), 0, 0));
                    Vec4 bounds_minz = HalfFloatConversion::ToFloat(UVec4(node->mBoundsMinZ[0] + (node->mBoundsMinZ[1] << 16), node->mBoundsMinZ[2] + (node->mBoundsMinZ[3] << 16), 0, 0));
 
                    Vec4 bounds_maxx = HalfFloatConversion::ToFloat(UVec4(node->mBoundsMaxX[0] + (node->mBoundsMaxX[1] << 16), node->mBoundsMaxX[2] + (node->mBoundsMaxX[3] << 16), 0, 0));
                    Vec4 bounds_maxy = HalfFloatConversion::ToFloat(UVec4(node->mBoundsMaxY[0] + (node->mBoundsMaxY[1] << 16), node->mBoundsMaxY[2] + (node->mBoundsMaxY[3] << 16), 0, 0));
                    Vec4 bounds_maxz = HalfFloatConversion::ToFloat(UVec4(node->mBoundsMaxZ[0] + (node->mBoundsMaxZ[1] << 16), node->mBoundsMaxZ[2] + (node->mBoundsMaxZ[3] << 16), 0, 0));
                #else
                    UVec4 bounds_minxy = UVec4::sLoadInt4(reinterpret_cast<const uint32 *>(&node->mBoundsMinX[0]));
                    Vec4 bounds_minx = HalfFloatConversion::ToFloat(bounds_minxy);
                    Vec4 bounds_miny = HalfFloatConversion::ToFloat(bounds_minxy.Swizzle<SWIZZLE_Z, SWIZZLE_W, SWIZZLE_UNUSED, SWIZZLE_UNUSED>());
 
                    UVec4 bounds_minzmaxx = UVec4::sLoadInt4(reinterpret_cast<const uint32 *>(&node->mBoundsMinZ[0]));
                    Vec4 bounds_minz = HalfFloatConversion::ToFloat(bounds_minzmaxx);
                    Vec4 bounds_maxx = HalfFloatConversion::ToFloat(bounds_minzmaxx.Swizzle<SWIZZLE_Z, SWIZZLE_W, SWIZZLE_UNUSED, SWIZZLE_UNUSED>());
 
                    UVec4 bounds_maxyz = UVec4::sLoadInt4(reinterpret_cast<const uint32 *>(&node->mBoundsMaxY[0]));
                    Vec4 bounds_maxy = HalfFloatConversion::ToFloat(bounds_maxyz);
                    Vec4 bounds_maxz = HalfFloatConversion::ToFloat(bounds_maxyz.Swizzle<SWIZZLE_Z, SWIZZLE_W, SWIZZLE_UNUSED, SWIZZLE_UNUSED>());
                #endif
 
                    // Load properties for 4 children
                    UVec4 properties = UVec4::sLoadInt4(&node->mNodeProperties[0]);
 
                    // Check which sub nodes to visit
                    int num_results = ioVisitor.VisitNodes(bounds_minx, bounds_miny, bounds_minz, bounds_maxx, bounds_maxy, bounds_maxz, properties, mTop);
 
                    // Push them onto the stack
                    JPH_ASSERT(mTop + 4 < StackSize);
                    properties.StoreInt4(&mNodeStack[mTop]);
                    mTop += num_results;
                }
                else if (tri_count != TRIANGLE_COUNT_MASK) // TRIANGLE_COUNT_MASK indicates a padding node, normally we shouldn't visit these nodes but when querying with a big enough box you could touch HALF_FLT_MAX (about 65K)
                {
                    // Node contains triangles, do individual tests
                    uint32 triangle_block_id = node_properties & OFFSET_MASK;
                    const void *triangles = sGetTriangleBlockStart(inBufferStart, triangle_block_id);
 
                    ioVisitor.VisitTriangles(inTriangleContext, triangles, tri_count, triangle_block_id);
                }
 
                // Check if we're done
                if (ioVisitor.ShouldAbort())
                    break;
 
                // Fetch next node until we find one that the visitor wants to see
                do
                    --mTop;
                while (mTop >= 0 && !ioVisitor.ShouldVisitNode(mTop));
            }
            while (mTop >= 0);
        }
 
        bool                        IsDoneWalking() const
        {
            return mTop < 0;
        }
 
    private:
        uint32                      mNodeStack[StackSize];
        int                         mTop = 0;
    };
};
 
JPH_NAMESPACE_END