SparseDynamicOctree3.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
#pragma once
 
 
#include "Async/ParallelFor.h"
#include "BoxTypes.h"
#include "Containers/Array.h"
#include "Containers/BitArray.h"
#include "Containers/Map.h"
#include "Containers/Set.h"
#include "Containers/StaticArray.h"
#include "Containers/UnrealString.h"
#include "CoreMinimal.h"
#include "GeometryTypes.h"
#include "HAL/PlatformCrt.h"
#include "IntVectorTypes.h"
#include "Intersection/IntrRay3AxisAlignedBox3.h"
#include "Math/NumericLimits.h"
#include "Math/Ray.h"
#include "Math/UnrealMathSSE.h"
#include "Math/Vector.h"
#include "MathUtil.h"
#include "Misc/AssertionMacros.h"
#include "Spatial/SparseGrid3.h"
#include "Util/DynamicVector.h"
#include "Util/RefCountVector.h"
#include "Util/SmallListSet.h"
 
template <typename FuncType> class TFunctionRef;
 
namespace UE
{
namespace Geometry
{
 
class FDynamicFlagArray
{
public:
    TBitArray<> BitArray;
    uint32 MaxIndex = 0;
    static constexpr int32 GrowChunkSize = 0xFFF;
 
    FDynamicFlagArray()
    {
        BitArray.Init(false, GrowChunkSize);
        MaxIndex = BitArray.Num();
    }
 
    void Set(uint32 BitIndex, bool bValue)
    {
        if (bValue && BitIndex >= MaxIndex)
        {
            int32 ToAdd = (BitIndex - MaxIndex + 1) | GrowChunkSize;
            BitArray.Add(false, ToAdd);
            MaxIndex = BitArray.Num();
            BitArray[BitIndex] = bValue;
        }
        else if (BitIndex < MaxIndex)
        {
            BitArray[BitIndex] = bValue;
        }
        // don't need to set anything for false bValue with BitIndex beyond bounds; Get() will return false beyond bounds
    }
 
    inline bool Get(uint32 BitIndex) const
    {
        return (BitIndex < MaxIndex) ? BitArray[BitIndex] : false;
    }
};
 
 
 
 
struct FSparseOctreeCell
{
    static constexpr uint32 InvalidID = TNumericLimits<uint32>::Max();;
    static constexpr uint8 InvalidLevel = TNumericLimits<uint8>::Max();
 
    uint32 CellID;
 
    uint8 Level;
 
    FVector3i Index;
 
    TStaticArray<uint32, 8> Children;
 
    FSparseOctreeCell()
        : CellID(InvalidID), Level(0), Index(FVector3i::Zero())
    {
        Children[0] = Children[1] = Children[2] = Children[3] = InvalidID;
        Children[4] = Children[5] = Children[6] = Children[7] = InvalidID;
    }
 
    FSparseOctreeCell(uint8 LevelIn, const FVector3i& IndexIn)
        : CellID(InvalidID), Level(LevelIn), Index(IndexIn)
    {
        Children[0] = Children[1] = Children[2] = Children[3] = InvalidID;
        Children[4] = Children[5] = Children[6] = Children[7] = InvalidID;
    }
 
    inline bool IsExistingCell() const
    {
        return CellID != InvalidID;
    }
 
    inline bool HasChild(int ChildIndex) const
    {
        return Children[ChildIndex] != InvalidID;
    }
 
    inline uint32 GetChildCellID(int ChildIndex) const
    {
        return Children[ChildIndex];
    }
 
    inline FSparseOctreeCell MakeChildCell(int ChildIndex)
    {
        FVector3i IndexOffset(
            ((ChildIndex & 1) != 0) ? 1 : 0,
            ((ChildIndex & 2) != 0) ? 1 : 0,
            ((ChildIndex & 4) != 0) ? 1 : 0);
        return FSparseOctreeCell(Level + 1, 2*Index + IndexOffset);
    }
 
    inline void SetChild(uint32 ChildIndex, const FSparseOctreeCell& ChildCell)
    {
        Children[ChildIndex] = ChildCell.CellID;
    }
 
};
 
 
class FSparseDynamicOctree3
{
    // potential optimizations/improvements
    //    - Cell sizes are known at each level...can keep a lookup table? will this improve performance?
    //    - Store cells for each level in separate TDynamicVectors. CellID is then [Level:8 | Index:24].
    //      This would allow level-grids to be processed separately / in-parallel (for example a cut at given level would be much faster)
    //    - Currently insertion is max-depth but we do not dynamically expand more than once. So early
    //      insertions end up in very large buckets. When a child expands we should check if any of its parents would fit.
    //    - Currently insertion is max-depth so we end up with a huge number of single-object cells. Should only go down a level
    //      if enough objects exist in current cell. Can do this in a greedy fashion, less optimal but still acceptable...
    //    - Store an expand-factor for each cell? or an actual AABB for each cell? this would allow for tighter bounds but
    //      requires accumulating expansion "up" the tree...
    //    - get rid of ValidObjectIDs, I don't think we need it?
 
public:
 
    //
    // Tree configuration parameters. It is not safe to change these after tree initialization!
    // 
 
    double RootDimension = 1000.0;
 
    double MaxExpandFactor = 0.25;
 
    static constexpr uint32 MaxSupportedTreeDepth = 0x1F;
 
    int GetMaxTreeDepth()
    {
        return MaxTreeDepth;
    }
    void SetMaxTreeDepth(int MaxTreeDepthIn)
    {
        if (!ensure((uint32)MaxTreeDepthIn <= MaxSupportedTreeDepth))
        {
            MaxTreeDepthIn = (int32)MaxSupportedTreeDepth;
        }
 
        MaxTreeDepth = MaxTreeDepthIn;
    }
protected:
    int MaxTreeDepth = 10;
 
 
public:
 
    GEOMETRYCORE_API bool ContainsObject(int32 ObjectID) const;
 
    GEOMETRYCORE_API void InsertObject(int32 ObjectID, const FAxisAlignedBox3d& Bounds);
 
    GEOMETRYCORE_API bool RemoveObject(int32 ObjectID);
 
    GEOMETRYCORE_API bool ReinsertObject(int32 ObjectID, const FAxisAlignedBox3d& NewBounds, uint32 CellIDHint = InvalidCellID);
 
    GEOMETRYCORE_API bool CheckIfObjectNeedsReinsert(int32 ObjectID, const FAxisAlignedBox3d& NewBounds, uint32& CellIDOut) const;
 
 
    GEOMETRYCORE_API int32 FindNearestHitObject(const FRay3d& Ray,
        TFunctionRef<FAxisAlignedBox3d(int)> GetObjectBoundsFunc,
        TFunctionRef<double(int, const FRay3d&)> HitObjectDistFunc,
        double MaxDistance = TNumericLimits<double>::Max()) const;
 
    GEOMETRYCORE_API void ContainmentQuery(const FVector3d& Point,
        TFunctionRef<void(int)> ObjectIDFunc) const;
 
    GEOMETRYCORE_API bool ContainmentQueryCancellable(const FVector3d& Point,
        TFunctionRef<bool(int)> ObjectIDFunc) const;
 
    GEOMETRYCORE_API void RangeQuery(const FAxisAlignedBox3d& Bounds,
        TFunctionRef<void(int)> ObjectIDFunc) const;
 
    GEOMETRYCORE_API void RangeQuery(const FAxisAlignedBox3d& Bounds,
        TArray<int>& ObjectIDsOut ) const;
 
    GEOMETRYCORE_API void ParallelRangeQuery(const FAxisAlignedBox3d& Bounds,
        TArray<int>& ObjectIDsOut) const;
 
    GEOMETRYCORE_API void ParallelRangeQuery(
        const FAxisAlignedBox3d& RangeBoundsHint,
        TFunctionRef<bool (const FAxisAlignedBox3d& CellBounds)> BoundsOverlapFn,
        TArray<int32>& ObjectIDs) const;
 
    GEOMETRYCORE_API int ParallelOverlapAnyQuery(const FAxisAlignedBox3d& ShapeBounds,
        TFunctionRef<bool(int32)> ObjectOverlapFn, 
        TFunctionRef<bool(const FAxisAlignedBox3d&)> BoundsOverlapFn) const;
 
    int ParallelOverlapAnyQuery(const FAxisAlignedBox3d& ShapeBounds,
        TFunctionRef<bool(int32)> ObjectOverlapFn) const
    {
        return ParallelOverlapAnyQuery(ShapeBounds, ObjectOverlapFn, 
            // Note a bounds vs ShapeBounds test will always be performed automatically;
            // the optional bounds overlap fn here can just return true to indicate no additional filtering is performed
            [](const FAxisAlignedBox3d&) {return true;}
        );
    }
 
 
    GEOMETRYCORE_API void CheckValidity(
        TFunctionRef<bool(int)> IsValidObjectIDFunc,
        TFunctionRef<FAxisAlignedBox3d(int)> GetObjectBoundsFunc,
        EValidityCheckFailMode FailMode = EValidityCheckFailMode::Check,
        bool bVerbose = false,
        bool bFailOnMissingObjects = false) const;
 
    struct FStatistics
    {
        int32 Levels;
        TArray<int32> LevelBoxCounts;
        TArray<int32> LevelObjCounts;
        int32 SpillObjCount;
        FString ToString() const;
    };
 
    GEOMETRYCORE_API void ComputeStatistics(FStatistics& StatsOut) const;
 
protected:
    // this identifier is used for unknown cells
    static constexpr uint32 InvalidCellID = FSparseOctreeCell::InvalidID;
 
    // if an object is in the spill cell, that means it didn't fit in the tree
    static constexpr uint32 SpillCellID = InvalidCellID - 1;
 
    // reference counts for Cells list. We don't actually need reference counts here, but we need a free
    // list and iterators, and FRefCountVector provides this
    FRefCountVector CellRefCounts;
 
    // list of cells. Note that some cells may be unused, depending on CellRefCounts
    TDynamicVector<FSparseOctreeCell> Cells;
 
    // TODO: Consider switching FSmallListSet to TSparseListSet<int32> for more reliable performance in cases where we have more than 8 objects in many cells
    FSmallListSet CellObjectLists;          // per-cell object ID lists
    TSet<int32> SpillObjectSet;             // list of object IDs for objects that didn't fit in a root cell
 
    TDynamicVector<uint32> ObjectIDToCellMap;   // map from external Object IDs to which cell the object is in (or spill cell, or invalid)
    FDynamicFlagArray ValidObjectIDs;           // set of ObjectIDs in the tree. This is perhaps not necessary...couldn't we rely on ObjectIDToCellMap?
 
    // RootCells are the top-level cells of the octree, of size RootDimension. 
    // So the elements of this sparse grid are CellIDs
    TSparseGrid3<uint32> RootCells;
 
 
    // calculate the base width of a cell at a given level
    inline double GetCellWidth(uint32 Level) const
    {
        checkSlow(Level <= MaxSupportedTreeDepth);
        double Divisor = (double)( (uint64)1 << (Level & MaxSupportedTreeDepth) );
        double CellWidth = RootDimension / Divisor;
        return CellWidth;
    }
 
 
    FAxisAlignedBox3d GetBox(uint32 Level, const FVector3i& Index, double ExpandFactor) const
    {
        double CellWidth = GetCellWidth(Level);
        double ExpandDelta = CellWidth * ExpandFactor;
        double MinX = (CellWidth * (double)Index.X) - ExpandDelta;
        double MinY = (CellWidth * (double)Index.Y) - ExpandDelta;
        double MinZ = (CellWidth * (double)Index.Z) - ExpandDelta;
        CellWidth += 2.0 * ExpandDelta;
        return FAxisAlignedBox3d(
            FVector3d(MinX, MinY, MinZ),
            FVector3d(MinX + CellWidth, MinY + CellWidth, MinZ + CellWidth));
    }
    inline FAxisAlignedBox3d GetCellBox(const FSparseOctreeCell& Cell, double ExpandFactor = 0) const
    {
        return GetBox(Cell.Level, Cell.Index, ExpandFactor);
    }
    FVector3d GetCellCenter(const FSparseOctreeCell& Cell) const
    {
        double CellWidth = GetCellWidth(Cell.Level);
        double MinX = CellWidth * (double)Cell.Index.X;
        double MinY = CellWidth * (double)Cell.Index.Y;
        double MinZ = CellWidth * (double)Cell.Index.Z;
        CellWidth *= 0.5;
        return FVector3d(MinX + CellWidth, MinY + CellWidth, MinZ + CellWidth);
    }
 
 
    // warning: result here appears to be unstable (due to optimization?) if any of the position values are on the border of a cell
    // (in testing, pragma-optimization-disabled code produced off-by-one result from calling this function)
    FVector3i PointToIndex(uint32 Level, const FVector3d& Position) const
    {
        double CellWidth = GetCellWidth(Level);
 
        // Converts to an int32 in such a way that values outside of the representable
        //  range ends up at the min and max values of the representable range.
        auto SafeDoubleToInt32Floor = [](double Value)
        {
            return static_cast<int32>(FMath::Clamp(
                FMathd::Floor(Value),
                static_cast<double>(TNumericLimits<int32>::Lowest()), 
                static_cast<double>(TNumericLimits<int32>::Max())));
        };
 
        int32 i = SafeDoubleToInt32Floor(Position.X / CellWidth);
        int32 j = SafeDoubleToInt32Floor(Position.Y / CellWidth);
        int32 k = SafeDoubleToInt32Floor(Position.Z / CellWidth);
        return FVector3i(i, j, k);
    }
 
 
    int ToChildCellIndex(const FSparseOctreeCell& Cell, const FVector3d& Position) const
    {
        FVector3d Center = GetCellCenter(Cell);
        int ChildIndex =
            ((Position.X < Center.X) ? 0 : 1) +
            ((Position.Y < Center.Y) ? 0 : 2) +
            ((Position.Z < Center.Z) ? 0 : 4);
        return ChildIndex;
    }
 
    bool CanFit(const FSparseOctreeCell& Cell, const FAxisAlignedBox3d& Bounds) const
    {
        FAxisAlignedBox3d CellBox = GetCellBox(Cell, MaxExpandFactor);
        return CellBox.Contains(Bounds);
    }
 
    uint32 GetCellForObject(int32 ObjectID) const
    {
        if (ObjectID >= 0 && static_cast<size_t>(ObjectID) < ObjectIDToCellMap.Num())
        {
            return ObjectIDToCellMap[ObjectID];
        }
        return InvalidCellID;
    }
 
    GEOMETRYCORE_API FSparseOctreeCell FindCurrentContainingCell(const FAxisAlignedBox3d& Bounds) const;
 
 
    GEOMETRYCORE_API void Insert_Spill(int32 ObjectID, const FAxisAlignedBox3d& Bounds);
    GEOMETRYCORE_API void Insert_NewRoot(int32 ObjectID, const FAxisAlignedBox3d& Bounds, FSparseOctreeCell NewRootCell);
    GEOMETRYCORE_API void Insert_ToCell(int32 ObjectID, const FAxisAlignedBox3d& Bounds, const FSparseOctreeCell& ExistingCell);
    GEOMETRYCORE_API void Insert_NewChildCell(int32 ObjectID, const FAxisAlignedBox3d& Bounds, int ParentCellID, FSparseOctreeCell NewChildCell, int ChildIdx);
 
    GEOMETRYCORE_API double FindNearestRayCellIntersection(const FSparseOctreeCell& Cell, const FRay3d& Ray) const;
 
 
    GEOMETRYCORE_API void BranchRangeQuery(const FSparseOctreeCell* ParentCell,
        const FAxisAlignedBox3d& Bounds,
        TArray<int>& ObjectIDs) const;
 
private:
    // Gather objects in all cells that pass BoundsOverlapFn, starting from a set of root cells
    void ParallelCollectObjectsInOverlappingCells(
        TArray<const FSparseOctreeCell*, TInlineAllocator<32>> RelevantRootCells,
        TFunctionRef<bool(const FAxisAlignedBox3d& CellBounds)> BoundsOverlapFn,
        TArray<int32>& ObjectIDs) const;
    // Go down subtree of ParentCell and collect all objects in cells that pass BoundsOverlapFn.
    void BranchRangeQuery(const FSparseOctreeCell* ParentCell,
        TFunctionRef<bool(const FAxisAlignedBox3d&)> BoundsOverlapFn,
        TArray<int>& ObjectIDs) const;
 
    int BranchCustomOverlapAnyQuery(const FSparseOctreeCell* ParentCell,
        const FAxisAlignedBox3d& Bounds,
        TFunctionRef<bool(int32)> ObjectOverlapFn, TFunctionRef<bool(const FAxisAlignedBox3d&)> BoundsOverlapFn) const;
 
    TArray<const FSparseOctreeCell*, TInlineAllocator<32>> InitializeQueryQueue(
        const FAxisAlignedBox3d& RangeBounds, 
        TFunctionRef<bool(const FAxisAlignedBox3d& CellBounds)> BoundsOverlapFn, 
        bool bCheckOverlapFnInRangeBounds = true) const;
 
    // helper to find the root-level cells that could intersect with a given query object, specialized for point containment queries
    TArray<const FSparseOctreeCell*, TInlineAllocator<32>> InitializeQueryQueue(const FVector3d& Point) const
    {
        return InitializeQueryQueue(FAxisAlignedBox3d(Point, Point),
            [&Point](const FAxisAlignedBox3d& CellBounds) { return CellBounds.Contains(Point); }, 
            false);
    }
 
    // helper to find the root-level cells that could intersect with a given query object, specialized for box queries
    TArray<const FSparseOctreeCell*, TInlineAllocator<32>> InitializeQueryQueue(const FAxisAlignedBox3d& Bounds) const
    {
        return InitializeQueryQueue(Bounds,
            [&Bounds](const FAxisAlignedBox3d& CellBounds) { return CellBounds.Intersects(Bounds); },
            false);
    }
};
 
 
} // end namespace UE::Geometry
} // end namespace UE