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)
    {
        typename NodeCodec::EncodingContext node_ctx;
        typename TriangleCodec::EncodingContext tri_ctx(inVertices);
 
        // Child nodes out of loop so we don't constantly realloc it
        Array<const AABBTreeBuilder::Node *> child_nodes;
        child_nodes.reserve(NumChildrenPerNode);
 
        // First calculate how big the tree is going to be.
        // Since the tree can be huge for very large meshes, we don't want
        // to reallocate the buffer as it may cause out of memory situations.
        // This loop mimics the construction loop below.
        uint64 total_size = HeaderSize + TriangleHeaderSize;
        size_t node_count = 1; // Start with root node
        size_t to_process_max_size = 1; // Track size of queues so we can do a single reserve below
        size_t to_process_triangles_max_size = 0;
        {   // A scope to free the memory associated with to_estimate and to_estimate_triangles
            Array<const AABBTreeBuilder::Node *> to_estimate;
            Array<const AABBTreeBuilder::Node *> to_estimate_triangles;
            to_estimate.push_back(inRoot);
            for (;;)
            {
                while (!to_estimate.empty())
                {
                    // Get the next node to process
                    const AABBTreeBuilder::Node *node = to_estimate.back();
                    to_estimate.pop_back();
 
                    // Update total size
                    node_ctx.PrepareNodeAllocate(node, total_size);
 
                    if (node->HasChildren())
                    {
                        // Collect the first NumChildrenPerNode sub-nodes in the tree
                        child_nodes.clear(); // Won't free the memory
                        node->GetNChildren(inNodes, NumChildrenPerNode, child_nodes);
 
                        // Increment the number of nodes we're going to store
                        node_count += child_nodes.size();
 
                        // Insert in reverse order so we estimate left child first when taking nodes from the back
                        for (int idx = int(child_nodes.size()) - 1; idx >= 0; --idx)
                        {
                            // Store triangles in separate list so we process them last
                            const AABBTreeBuilder::Node *child = child_nodes[idx];
                            if (child->HasChildren())
                            {
                                to_estimate.push_back(child);
                                to_process_max_size = max(to_estimate.size(), to_process_max_size);
                            }
                            else
                            {
                                to_estimate_triangles.push_back(child);
                                to_process_triangles_max_size = max(to_estimate_triangles.size(), to_process_triangles_max_size);
                            }
                        }
                    }
                    else
                    {
                        // Update total size
                        tri_ctx.PreparePack(&inTriangles[node->mTrianglesBegin], node->mNumTriangles, inStoreUserData, total_size);
                    }
                }
 
                // If we've got triangles to estimate, loop again with just the triangles
                if (to_estimate_triangles.empty())
                    break;
                else
                    to_estimate.swap(to_estimate_triangles);
            }
        }
 
        // Finalize the prepare stage for the triangle context
        tri_ctx.FinalizePreparePack(total_size);
 
        // Reserve the buffer
        if (size_t(total_size) != total_size)
        {
            outError = "AABBTreeToBuffer: Out of memory!";
            return false;
        }
        mTree.reserve(size_t(total_size));
 
        // 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;
            size_t                          mNodeStart = size_t(-1);                    // Start of node in mTree
            size_t                          mTriangleStart = size_t(-1);                // Start of the triangle data in mTree
            size_t                          mChildNodeStart[NumChildrenPerNode];        // Start of the children of the node in mTree
            size_t                          mChildTrianglesStart[NumChildrenPerNode];   // Start of the triangle data in mTree
            size_t *                        mParentChildNodeStart = nullptr;            // Where to store mNodeStart (to patch mChildNodeStart of my parent)
            size_t *                        mParentTrianglesStart = nullptr;            // Where to store mTriangleStart (to patch mChildTrianglesStart of my parent)
            uint                            mNumChildren = 0;                           // Number of children
        };
 
        Array<NodeData *> to_process;
        to_process.reserve(to_process_max_size);
        Array<NodeData *> to_process_triangles;
        to_process_triangles.reserve(to_process_triangles_max_size);
        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());
 
        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
                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 == size_t(-1))
                    return false;
 
                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)
                    {
                        const AABBTreeBuilder::Node *child_node = child_nodes[idx];
 
                        // Due to quantization box could have become bigger, not smaller
                        JPH_ASSERT(AABox(child_bounds_min[idx], child_bounds_max[idx]).Contains(child_node->mBounds), "AABBTreeToBuffer: Bounding box became smaller!");
 
                        // Add child to list of nodes to be processed
                        NodeData child;
                        child.mNode = child_node;
                        child.mNodeBoundsMin = child_bounds_min[idx];
                        child.mNodeBoundsMax = child_bounds_max[idx];
                        child.mParentChildNodeStart = &node_data->mChildNodeStart[idx];
                        child.mParentTrianglesStart = &node_data->mChildTrianglesStart[idx];
                        node_list.push_back(child);
 
                        // Store triangles in separate list so we process them last
                        if (child_node->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 == size_t(-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);
        }
 
        // Assert that our reservation was correct (we don't know if we swapped the arrays or not)
        JPH_ASSERT(to_process_max_size == to_process.capacity() || to_process_triangles_max_size == to_process.capacity());
        JPH_ASSERT(to_process_max_size == to_process_triangles.capacity() || to_process_triangles_max_size == to_process_triangles.capacity());
 
        // 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 our reservations were correct
        if (node_count != node_list.size())
        {
            outError = "Internal Error: Node memory estimate was incorrect, memory corruption!";
            return false;
        }
        if (total_size != mTree.size())
        {
            outError = "Internal Error: Tree memory estimate was incorrect, memory corruption!";
            return false;
        }
 
        // Finalize the nodes
        return node_ctx.Finalize(header, inRoot, node_list[0].mNodeStart, node_list[0].mTriangleStart, outError);
    }
 
    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;                                  
};
 
JPH_NAMESPACE_END