imstkLooseOctree.h

/*
** This file is part of the Interactive Medical Simulation Toolkit (iMSTK)
** iMSTK is distributed under the Apache License, Version 2.0.
** See accompanying NOTICE for details.
*/

#pragma once

#include "imstkMath.h"
#include "imstkParallelUtils.h"

#include <array>
#include <unordered_set>

#ifdef WIN32
#pragma warning(disable : 4201)
#endif
namespace imstk
{
class OctreeNode;
class LooseOctree;
class Geometry;
class PointSet;
class SurfaceMesh;

enum OctreePrimitiveType
{
    Point = 0,
    Triangle,
    Analytical,
    NumPrimitiveTypes
};

struct OctreePrimitive
{
    OctreePrimitive(const OctreePrimitive&) = delete;
    OctreePrimitive& operator=(const OctreePrimitive&) = delete;

    OctreePrimitive() : m_pGeometry(nullptr), m_GeomIdx(0), m_Idx(0) {}
    OctreePrimitive(Geometry* const pGeometry, const uint32_t geomIdx, const uint32_t idx = 0) :
        m_pGeometry(pGeometry), m_GeomIdx(geomIdx), m_Idx(idx) {}

    Geometry* const m_pGeometry;        
    const uint32_t m_GeomIdx;           
    const uint32_t m_Idx;               

    OctreeNode* m_pNode      = nullptr; 
    OctreePrimitive* m_pNext = nullptr; 

    union
    {
        std::array<double, 3> m_Position;   
        struct
        {
            std::array<double, 3> m_LowerCorner;   
            std::array<double, 3> m_UpperCorner;   
        };
    };

    bool m_bValid = true;
};

struct OctreeNodeBlock;

class OctreeNode
{
friend class LooseOctree;
friend class OctreeBasedCD;
friend class LooseOctreeTest;
public:
    OctreeNode(const OctreeNode&) = delete;
    OctreeNode& operator=(const OctreeNode&) = delete;

    OctreeNode() :
        m_pTree(nullptr),
        m_pParent(nullptr),
        m_Center(Vec3d::Zero()),
        m_HalfWidth(0.0),
        m_Depth(0),
        m_MaxDepth(0),
        m_bIsLeaf(true) {}

    explicit OctreeNode(LooseOctree* const tree, OctreeNode* const pParent, const Vec3d& nodeCenter,
                        const double halfWidth, const uint32_t depth);

    bool isLeaf() const { return m_bIsLeaf; }

    OctreeNode* getChildNode(const uint32_t childIdx) const;

    OctreePrimitive* getPrimitiveList(const OctreePrimitiveType type) const { return m_pPrimitiveListHeads[type]; }

    uint32_t getPrimitiveCount(const OctreePrimitiveType type) const { return m_PrimitiveCounts[type]; }

    const Vec6d getBounds()
    {
        Vec6d bounds;
        bounds << m_LowerBound[0], m_UpperBound[0], m_LowerBound[1], m_UpperBound[1], m_LowerBound[2], m_UpperBound[2];
        return bounds;
    }

    const Vec6d getLooseBounds()
    {
        Vec6d bounds;
        bounds << m_LowerExtendedBound[0], m_UpperExtendedBound[0], m_LowerExtendedBound[1], m_UpperExtendedBound[1], m_LowerExtendedBound[2], m_UpperExtendedBound[2];
        return bounds;
    }

    void clearPrimitiveData(const OctreePrimitiveType type);

    void split();

    void removeAllDescendants();

    void removeEmptyDescendants();

    void keepPrimitive(OctreePrimitive* const pPrimitive, const OctreePrimitiveType type);

    void insertPoint(OctreePrimitive* const pPrimitive);

    void insertNonPointPrimitive(OctreePrimitive* const pPrimitive, const OctreePrimitiveType type);

    bool contains(const Vec3d& point) { return contains(point[0], point[1], point[2]); }
    bool contains(const std::array<double, 3>& point) { return contains(point[0], point[1], point[2]); }
    bool contains(const double x, const double y, const double z)
    {
        return x >= m_LowerBound[0]
               && y >= m_LowerBound[1]
               && z >= m_LowerBound[2]
               && x <= m_UpperBound[0]
               && y <= m_UpperBound[1]
               && z <= m_UpperBound[2];
    }

    bool contains(const std::array<double, 3>& lowerCorner, const std::array<double, 3>& upperCorner)
    {
        return lowerCorner[0] >= m_LowerBound[0]
               && lowerCorner[1] >= m_LowerBound[1]
               && lowerCorner[2] >= m_LowerBound[2]
               && upperCorner[0] <= m_UpperBound[0]
               && upperCorner[1] <= m_UpperBound[1]
               && upperCorner[2] <= m_UpperBound[2];
    }

    bool looselyContains(const Vec3d& point) { return looselyContains(point[0], point[1], point[2]); }
    bool looselyContains(const std::array<double, 3>& point) { return looselyContains(point[0], point[1], point[2]); }
    bool looselyContains(const double x, const double y, const double z)
    {
        return x >= m_LowerExtendedBound[0]
               && y >= m_LowerExtendedBound[1]
               && z >= m_LowerExtendedBound[2]
               && x <= m_UpperExtendedBound[0]
               && y <= m_UpperExtendedBound[1]
               && z <= m_UpperExtendedBound[2];
    }

    bool looselyContains(const std::array<double, 3>& lowerCorner, const std::array<double, 3>& upperCorner)
    {
        return lowerCorner[0] >= m_LowerExtendedBound[0]
               && lowerCorner[1] >= m_LowerExtendedBound[1]
               && lowerCorner[2] >= m_LowerExtendedBound[2]
               && upperCorner[0] <= m_UpperExtendedBound[0]
               && upperCorner[1] <= m_UpperExtendedBound[1]
               && upperCorner[2] <= m_UpperExtendedBound[2];
    }

    bool looselyOverlaps(const std::array<double, 3>& lowerCorner, const std::array<double, 3>& upperCorner)
    {
        return upperCorner[0] >= m_LowerExtendedBound[0]
               && upperCorner[1] >= m_LowerExtendedBound[1]
               && upperCorner[2] >= m_LowerExtendedBound[2]
               && lowerCorner[0] <= m_UpperExtendedBound[0]
               && lowerCorner[1] <= m_UpperExtendedBound[1]
               && lowerCorner[2] <= m_UpperExtendedBound[2];
    }

// \todo: Make private or protected once a proper public element iteration structure is implemented
public:
    LooseOctree*     m_pTree;               
    OctreeNode*      m_pParent;             
    OctreeNodeBlock* m_pChildren = nullptr; 

    const Vec3d    m_Center;                
    const Vec3d    m_LowerBound;            
    const Vec3d    m_UpperBound;            
    const Vec3d    m_LowerExtendedBound;    
    const Vec3d    m_UpperExtendedBound;    
    const double   m_HalfWidth;             
    const uint32_t m_Depth;                 
    uint32_t       m_MaxDepth;              
    bool m_bIsLeaf = true;                  

    OctreePrimitive* m_pPrimitiveListHeads[OctreePrimitiveType::NumPrimitiveTypes];

    uint32_t m_PrimitiveCounts[OctreePrimitiveType::NumPrimitiveTypes];

    ParallelUtils::SpinLock m_PrimitiveLock[OctreePrimitiveType::NumPrimitiveTypes];

    ParallelUtils::SpinLock m_NodeSplitingLock;
};

struct OctreeNodeBlock
{
    OctreeNode m_Nodes[8];
    OctreeNodeBlock* m_NextBlock = nullptr; 
};

class LooseOctree
{
friend class OctreeNode;
friend class LooseOctreeTest;
public:
    explicit LooseOctree(const Vec3d& center, const double width, const double minWidth,
                         const double minWidthRatio = 1.0, const std::string name = "LooseOctree");

    virtual ~LooseOctree();

    virtual void clear();

    void clearPrimitive(const OctreePrimitiveType type);

    const Vec3d getCenter() const { return m_Center; }

    double getWidth() const { return m_Width; }

    double getMinWidth() const { return m_MinWidth; }

    uint32_t getMaxDepth() const { return m_MaxDepth; }

    uint32_t getNumAllocatedNodes() const { return m_NumAllocatedNodes; }

    uint32_t getNumActiveNodes() const { return m_NumAllocatedNodes - m_NumAvaiableBlocksInPool * 8u; }

    OctreeNode* getRootNode() const { return m_pRootNode; }

    size_t getPrimitiveCount(const OctreePrimitiveType type) const { return m_vPrimitivePtrs[type].size(); }

    uint32_t getMaxNumPrimitivesInNodes() const;

    size_t getNumGeometries() const { return m_sGeometryIndices.size(); }

    bool hasGeometry(uint32_t geomIdx) const { return m_sGeometryIndices.find(geomIdx) != m_sGeometryIndices.end(); }

    uint32_t addPointSet(const std::shared_ptr<PointSet>& pointset);

    uint32_t addTriangleMesh(const std::shared_ptr<SurfaceMesh>& surfMesh);

    uint32_t addAnalyticalGeometry(const std::shared_ptr<Geometry>& geometry);

    void setAlwaysRebuild(const bool bAlwaysRebuild) { m_bAlwaysRebuild = bAlwaysRebuild; }

    void build();

    void update();

protected:
    void addGeometry(const uint32_t geomIdx);

    void removeGeometry(const uint32_t geomIdx);

    void rebuild();

    void populatePointPrimitives();

    void populateNonPointPrimitives(const OctreePrimitiveType type);

    void incrementalUpdate();

    void updatePositionAndCheckValidity();

    void updateBoundingBoxAndCheckValidity(const OctreePrimitiveType type);

    void removeInvalidPrimitivesFromNodes();

    void reinsertInvalidPrimitives(const OctreePrimitiveType type);

    void computePrimitiveBoundingBox(OctreePrimitive* const pPrimitive, const OctreePrimitiveType type);

    OctreeNodeBlock* requestChildrenFromPool();

    void returnChildrenToPool(OctreeNodeBlock* const pNodeBlock);

    void allocateMoreNodeBlock(const uint32_t numBlocks);

    void deallocateMemoryPool();

    const std::string m_Name;   
    const Vec3d       m_Center; 
    const double      m_Width;  

    const double m_MinWidthRatio;

    double   m_MinWidth;                                         
    uint32_t m_MaxDepth;                                         
    bool     m_useMaxDepth;                                      

    OctreeNode* const m_pRootNode;                               
    OctreeNodeBlock*  m_pNodeBlockPoolHead = nullptr;            
    uint32_t m_NumAvaiableBlocksInPool     = 0;                  
    uint32_t m_NumAllocatedNodes;                                
    ParallelUtils::SpinLock m_PoolLock;                          

    tbb::concurrent_unordered_set<OctreeNodeBlock*> m_sActiveTreeNodeBlocks;

    std::vector<OctreeNodeBlock*> m_pNodeBigBlocks;

    std::vector<OctreePrimitive*> m_vPrimitivePtrs[OctreePrimitiveType::NumPrimitiveTypes];

    std::vector<OctreePrimitive*> m_pPrimitiveBlocks[OctreePrimitiveType::NumPrimitiveTypes];

    std::unordered_set<uint32_t> m_sGeometryIndices;

    bool m_bAlwaysRebuild = false;                        
    bool m_bCompleteBuild = false;                        

    bool m_bDrawNonEmptyParent = true;
};
} // namespace imstk
#ifdef WIN32
#pragma warning(default : 4201)
#endif