AABBTreeToBuffer.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/AABBTree/AABBTreeBuilder.h>
#include <Jolt/Core/ByteBuffer.h>
#include <Jolt/Geometry/IndexedTriangle.h>
 
JPH_NAMESPACE_BEGIN
 
template <class TriangleCodec, class NodeCodec>
class AABBTreeToBuffer
{
public:
    using NodeHeader = typename NodeCodec::Header;
 
    static const int HeaderSize = NodeCodec::HeaderSize;
 
    static const int NumChildrenPerNode = NodeCodec::NumChildrenPerNode;
 
    using TriangleHeader = typename TriangleCodec::TriangleHeader;
 
    static const int TriangleHeaderSize = TriangleCodec::TriangleHeaderSize;
 
    bool                            Convert(const Array<IndexedTriangle> &inTriangles, const Array<AABBTreeBuilder::Node> &inNodes, const VertexList &inVertices, const AABBTreeBuilder::Node *inRoot, bool inStoreUserData, const char *&outError)
    {
        const typename NodeCodec::EncodingContext node_ctx;
        typename TriangleCodec::EncodingContext tri_ctx(inVertices);
 
        // Estimate the amount of memory required
        uint tri_count = inRoot->GetTriangleCountInTree(inNodes);
        uint node_count = inRoot->GetNodeCount(inNodes);
        uint nodes_size = node_ctx.GetPessimisticMemoryEstimate(node_count);
        uint total_size = HeaderSize + TriangleHeaderSize + nodes_size + tri_ctx.GetPessimisticMemoryEstimate(tri_count, inStoreUserData);
        mTree.reserve(total_size);
 
        // Reset counters
        mNodesSize = 0;
 
        // Add headers
        NodeHeader *header = HeaderSize > 0? mTree.Allocate<NodeHeader>() : nullptr;
        TriangleHeader *triangle_header = TriangleHeaderSize > 0? mTree.Allocate<TriangleHeader>() : nullptr;
 
        struct NodeData
        {
            const AABBTreeBuilder::Node *   mNode = nullptr;                            // Node that this entry belongs to
            Vec3                            mNodeBoundsMin;                             // Quantized node bounds
            Vec3                            mNodeBoundsMax;
            uint                            mNodeStart = uint(-1);                      // Start of node in mTree
            uint                            mTriangleStart = uint(-1);                  // Start of the triangle data in mTree
            uint                            mNumChildren = 0;                           // Number of children
            uint                            mChildNodeStart[NumChildrenPerNode];        // Start of the children of the node in mTree
            uint                            mChildTrianglesStart[NumChildrenPerNode];   // Start of the triangle data in mTree
            uint *                          mParentChildNodeStart = nullptr;            // Where to store mNodeStart (to patch mChildNodeStart of my parent)
            uint *                          mParentTrianglesStart = nullptr;            // Where to store mTriangleStart (to patch mChildTrianglesStart of my parent)
        };
 
        Array<NodeData *> to_process;
        Array<NodeData *> to_process_triangles;
        Array<NodeData> node_list;
 
        node_list.reserve(node_count); // Needed to ensure that array is not reallocated, so we can keep pointers in the array
 
        NodeData root;
        root.mNode = inRoot;
        root.mNodeBoundsMin = inRoot->mBounds.mMin;
        root.mNodeBoundsMax = inRoot->mBounds.mMax;
        node_list.push_back(root);
        to_process.push_back(&node_list.back());
 
        // Child nodes out of loop so we don't constantly realloc it
        Array<const AABBTreeBuilder::Node *> child_nodes;
        child_nodes.reserve(NumChildrenPerNode);
 
        for (;;)
        {
            while (!to_process.empty())
            {
                // Get the next node to process
                NodeData *node_data = to_process.back();
                to_process.pop_back();
 
                // Due to quantization box could have become bigger, not smaller
                JPH_ASSERT(AABox(node_data->mNodeBoundsMin, node_data->mNodeBoundsMax).Contains(node_data->mNode->mBounds), "AABBTreeToBuffer: Bounding box became smaller!");
 
                // Collect the first NumChildrenPerNode sub-nodes in the tree
                child_nodes.clear(); // Won't free the memory
                node_data->mNode->GetNChildren(inNodes, NumChildrenPerNode, child_nodes);
                node_data->mNumChildren = (uint)child_nodes.size();
 
                // Fill in default child bounds
                Vec3 child_bounds_min[NumChildrenPerNode], child_bounds_max[NumChildrenPerNode];
                for (size_t i = 0; i < NumChildrenPerNode; ++i)
                    if (i < child_nodes.size())
                    {
                        child_bounds_min[i] = child_nodes[i]->mBounds.mMin;
                        child_bounds_max[i] = child_nodes[i]->mBounds.mMax;
                    }
                    else
                    {
                        child_bounds_min[i] = Vec3::sZero();
                        child_bounds_max[i] = Vec3::sZero();
                    }
 
                // Start a new node
                uint old_size = (uint)mTree.size();
                node_data->mNodeStart = node_ctx.NodeAllocate(node_data->mNode, node_data->mNodeBoundsMin, node_data->mNodeBoundsMax, child_nodes, child_bounds_min, child_bounds_max, mTree, outError);
                if (node_data->mNodeStart == uint(-1))
                    return false;
                mNodesSize += (uint)mTree.size() - old_size;
 
                if (node_data->mNode->HasChildren())
                {
                    // Insert in reverse order so we process left child first when taking nodes from the back
                    for (int idx = int(child_nodes.size()) - 1; idx >= 0; --idx)
                    {
                        // Due to quantization box could have become bigger, not smaller
                        JPH_ASSERT(AABox(child_bounds_min[idx], child_bounds_max[idx]).Contains(child_nodes[idx]->mBounds), "AABBTreeToBuffer: Bounding box became smaller!");
 
                        // Add child to list of nodes to be processed
                        NodeData child;
                        child.mNode = child_nodes[idx];
                        child.mNodeBoundsMin = child_bounds_min[idx];
                        child.mNodeBoundsMax = child_bounds_max[idx];
                        child.mParentChildNodeStart = &node_data->mChildNodeStart[idx];
                        child.mParentTrianglesStart = &node_data->mChildTrianglesStart[idx];
                        NodeData *old = &node_list[0];
                        node_list.push_back(child);
                        if (old != &node_list[0])
                        {
                            outError = "Internal Error: Array reallocated, memory corruption!";
                            return false;
                        }
 
                        // Store triangles in separate list so we process them last
                        if (node_list.back().mNode->HasChildren())
                            to_process.push_back(&node_list.back());
                        else
                            to_process_triangles.push_back(&node_list.back());
                    }
                }
                else
                {
                    // Add triangles
                    node_data->mTriangleStart = tri_ctx.Pack(&inTriangles[node_data->mNode->mTrianglesBegin], node_data->mNode->mNumTriangles, inStoreUserData, mTree, outError);
                    if (node_data->mTriangleStart == uint(-1))
                        return false;
                }
 
                // Patch offset into parent
                if (node_data->mParentChildNodeStart != nullptr)
                {
                    *node_data->mParentChildNodeStart = node_data->mNodeStart;
                    *node_data->mParentTrianglesStart = node_data->mTriangleStart;
                }
            }
 
            // If we've got triangles to process, loop again with just the triangles
            if (to_process_triangles.empty())
                break;
            else
                to_process.swap(to_process_triangles);
        }
 
        // Finalize all nodes
        for (NodeData &n : node_list)
            if (!node_ctx.NodeFinalize(n.mNode, n.mNodeStart, n.mNumChildren, n.mChildNodeStart, n.mChildTrianglesStart, mTree, outError))
                return false;
 
        // Finalize the triangles
        tri_ctx.Finalize(inVertices, triangle_header, mTree);
 
        // Validate that we reserved enough memory
        if (nodes_size < mNodesSize)
        {
            outError = "Internal Error: Not enough memory reserved for nodes!";
            return false;
        }
        if (total_size < (uint)mTree.size())
        {
            outError = "Internal Error: Not enough memory reserved for triangles!";
            return false;
        }
 
        // Finalize the nodes
        if (!node_ctx.Finalize(header, inRoot, node_list[0].mNodeStart, node_list[0].mTriangleStart, outError))
            return false;
 
        // Shrink the tree, this will invalidate the header and triangle_header variables
        mTree.shrink_to_fit();
 
        return true;
    }
 
    inline const ByteBuffer &       GetBuffer() const
    {
        return mTree;
    }
 
    inline ByteBuffer &             GetBuffer()
    {
        return mTree;
    }
 
    inline const NodeHeader *       GetNodeHeader() const
    {
        return mTree.Get<NodeHeader>(0);
    }
 
    inline const TriangleHeader *   GetTriangleHeader() const
    {
        return mTree.Get<TriangleHeader>(HeaderSize);
    }
 
    inline const void *             GetRoot() const
    {
        return mTree.Get<void>(HeaderSize + TriangleHeaderSize);
    }
 
private:
    ByteBuffer                      mTree;                                  
    uint                            mNodesSize;                             
};
 
JPH_NAMESPACE_END