LumenSparseSpanArray.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
/*=============================================================================
    LumenSparseSpanArray.h:
=============================================================================*/
 
#pragma once
 
#include "CoreMinimal.h"
#include "Engine/EngineTypes.h"
#include "SpanAllocator.h"
 
// Sparse array with stable indices and contiguous span allocation
template <typename ElementType>
class TSparseSpanArray
{
public:
    int32 Num() const
    {
        return SpanAllocator.GetMaxSize();
    }
 
    void Reserve(int32 NumElements)
    {
        Elements.Reserve(NumElements);
    }
 
    int32 AddSpan(int32 NumElements)
    {
        check(NumElements > 0);
 
        const int32 InsertIndex = SpanAllocator.Allocate(NumElements);
        
        // Resize element array
        if (SpanAllocator.GetMaxSize() > Elements.Num())
        {
            const int32 NumElementsToAdd = SpanAllocator.GetMaxSize() - Elements.Num();
            Elements.AddDefaulted(NumElementsToAdd);
            AllocatedElementsBitArray.Add(false, NumElementsToAdd);
        }
 
        // Reuse existing elements
        for (int32 ElementIndex = InsertIndex; ElementIndex < InsertIndex + NumElements; ++ElementIndex)
        {
            checkSlow(!IsAllocated(ElementIndex));
            Elements[ElementIndex] = ElementType();
        }
 
        AllocatedElementsBitArray.SetRange(InsertIndex, NumElements, true);
 
        return InsertIndex;
    }
 
    void RemoveSpan(int32 FirstElementIndex, int32 NumElements)
    {
        check(NumElements > 0);
 
        for (int32 ElementIndex = FirstElementIndex; ElementIndex < FirstElementIndex + NumElements; ++ElementIndex)
        {
            checkSlow(IsAllocated(ElementIndex));
            Elements[ElementIndex] = ElementType();
        }
 
        SpanAllocator.Free(FirstElementIndex, NumElements);
        AllocatedElementsBitArray.SetRange(FirstElementIndex, NumElements, false);
    }
 
    void Consolidate()
    {
        SpanAllocator.Consolidate();
 
        if (Elements.Num() > SpanAllocator.GetMaxSize() * ShrinkThreshold)
        {
            Elements.SetNum(SpanAllocator.GetMaxSize());
            AllocatedElementsBitArray.SetNumUninitialized(SpanAllocator.GetMaxSize());
        }
    }
 
    void Reset()
    {
        Elements.Reset();
        SpanAllocator.Reset();
        AllocatedElementsBitArray.SetNumUninitialized(0);
    }
 
    ElementType& operator[](int32 Index)
    {
        checkSlow(IsAllocated(Index));
        return Elements[Index];
    }
 
    const ElementType& operator[](int32 Index) const
    {
        checkSlow(IsAllocated(Index));
        return Elements[Index];
    }
 
    bool IsAllocated(int32 ElementIndex) const
    {
        if (ElementIndex < Num())
        {
            return AllocatedElementsBitArray[ElementIndex];
        }
 
        return false;
    }
 
    SIZE_T GetAllocatedSize() const
    {
        return Elements.GetAllocatedSize() + AllocatedElementsBitArray.GetAllocatedSize() + SpanAllocator.GetAllocatedSize();
    }
 
    class TRangedForIterator
    {
    public:
        TRangedForIterator(TSparseSpanArray<ElementType>& InArray, int32 InElementIndex)
            : Array(InArray)
            , ElementIndex(InElementIndex)
        {
            // Scan for the first valid element.
            while (ElementIndex < Array.Num() && !Array.AllocatedElementsBitArray[ElementIndex])
            {
                ++ElementIndex;
            } 
        }
 
        TRangedForIterator operator++()
        {
            // Scan for the next first valid element.
            do
            {
                ++ElementIndex;
            } while (ElementIndex < Array.Num() && !Array.AllocatedElementsBitArray[ElementIndex]);
 
            return *this;
        }
 
        bool operator!=(const TRangedForIterator& Other) const
        {
            return ElementIndex != Other.ElementIndex;
        }
 
        ElementType& operator*()
        {
            return Array.Elements[ElementIndex];
        }
 
    private:
        TSparseSpanArray<ElementType>& Array;
        int32 ElementIndex;
    };
 
    class TRangedForConstIterator
    {
    public:
        TRangedForConstIterator(const TSparseSpanArray<ElementType>& InArray, int32 InElementIndex)
            : Array(InArray)
            , ElementIndex(InElementIndex)
        {
            // Scan for the first valid element.
            while (ElementIndex < Array.Num() && !Array.AllocatedElementsBitArray[ElementIndex])
            {
                ++ElementIndex;
            }
        }
 
        TRangedForConstIterator operator++()
        { 
            // Scan for the next first valid element.
            do
            {
                ++ElementIndex;
            } while (ElementIndex < Array.Num() && !Array.AllocatedElementsBitArray[ElementIndex]);
 
            return *this;
        }
 
        bool operator!=(const TRangedForConstIterator& Other) const
        { 
            return ElementIndex != Other.ElementIndex;
        }
 
        const ElementType& operator*() const
        { 
            return Array.Elements[ElementIndex];
        }
 
    private:
        const TSparseSpanArray<ElementType>& Array;
        int32 ElementIndex;
    };
 
    // Iterate over all allocated elements (skip free ones)
    TRangedForIterator begin() { return TRangedForIterator(*this, 0); }
    TRangedForIterator end() { return TRangedForIterator(*this, Num()); }
    TRangedForConstIterator begin() const { return TRangedForConstIterator(*this, 0); }
    TRangedForConstIterator end() const { return TRangedForConstIterator(*this, Num()); }
 
private:
    // Allocated size needs to be this much bigger than used size before we shrink this array
    static constexpr int32 ShrinkThreshold = 2;
 
    TArray<ElementType> Elements;
    TBitArray<> AllocatedElementsBitArray;
    FSpanAllocator SpanAllocator;
};
 
template <typename ElementType, uint32 BytesPerChunk = 2 * 1024 * 1024>
class TChunkedSparseArray
{
public:
    TChunkedSparseArray() = default;
 
    TChunkedSparseArray(const TChunkedSparseArray& Other)
    {
        *this = Other;
    }
 
    TChunkedSparseArray(TChunkedSparseArray&& Other)
    {
        *this = MoveTemp(Other);
    }
 
    TChunkedSparseArray& operator=(const TChunkedSparseArray& Other)
    {
        if (&Other != this)
        {
            Empty();
 
            FreeElementIndexHint = Other.FreeElementIndexHint;
            NumAllocatedElements = Other.NumAllocatedElements;
            MaxAllocatedElementIndexPlusOne = Other.MaxAllocatedElementIndexPlusOne;
            AllocatedElementsBitArray = Other.AllocatedElementsBitArray;
 
            const int32 NumElementChunks = Other.ElementChunks.Num();
            ElementChunks.Reserve(NumElementChunks);
            for (int32 ChunkIndex = 0; ChunkIndex < NumElementChunks; ++ChunkIndex)
            {
                ElementChunks.Add((ElementType*)FMemory::Malloc(BytesPerChunk));
            }
 
            for (TConstSetBitIterator It(AllocatedElementsBitArray); It && It.GetIndex() < GetMaxSize(); ++It)
            {
                const int32 ElementIndex = It.GetIndex();
                ElementType& Element = (*this)[ElementIndex];
                new (&Element) ElementType(Other[ElementIndex]);
            }
        }
 
        return *this;
    }
 
    TChunkedSparseArray& operator=(TChunkedSparseArray&& Other)
    {
        if (&Other != this)
        {
            Empty();
            
            FreeElementIndexHint = Other.FreeElementIndexHint;
            NumAllocatedElements = Other.NumAllocatedElements;
            MaxAllocatedElementIndexPlusOne = Other.MaxAllocatedElementIndexPlusOne;
            Other.FreeElementIndexHint = 0;
            Other.NumAllocatedElements = 0;
            Other.MaxAllocatedElementIndexPlusOne = 0;
 
            AllocatedElementsBitArray = MoveTemp(Other.AllocatedElementsBitArray);
            ElementChunks = MoveTemp(Other.ElementChunks);
        }
 
        return *this;
    }
 
    ~TChunkedSparseArray()
    {
        Empty();
    }
 
    void Empty()
    {
        for (ElementType& Element : *this)
        {
            Element.~ElementType();
        }
 
        for (ElementType* Chunk : ElementChunks)
        {
            FMemory::Free(Chunk);
        }
 
        ElementChunks.Empty();
        AllocatedElementsBitArray.Empty();
        FreeElementIndexHint = 0;
        NumAllocatedElements = 0;
        MaxAllocatedElementIndexPlusOne = 0;
    }
 
    int32 GetMaxSize() const
    {
        return MaxAllocatedElementIndexPlusOne;
    }
 
    void Reserve(int32 NumElements)
    {
        const int32 NewNumChunks = FMath::DivideAndRoundUp(NumElements, ElementsPerChunk);
        
        AllocatedElementsBitArray.SetNum(FMath::Max(ElementChunks.Num(), NewNumChunks) * ElementsPerChunk, false);
 
        while (ElementChunks.Num() < NewNumChunks)
        {
            ElementChunks.Add((ElementType*)FMemory::Malloc(BytesPerChunk));
        }
    }
 
    int32 AddDefaulted()
    {
        // Find the first unused element slot or add to the end
        int32 InsertIndex;
        if (NumAllocatedElements == MaxAllocatedElementIndexPlusOne)
        {
            check(MaxAllocatedElementIndexPlusOne <= AllocatedElementsBitArray.Num());
            InsertIndex = MaxAllocatedElementIndexPlusOne;
            ++MaxAllocatedElementIndexPlusOne;
 
            if (InsertIndex < AllocatedElementsBitArray.Num())
            {
                AllocatedElementsBitArray[InsertIndex] = true;
            }
        }
        else
        {
            check(NumAllocatedElements < MaxAllocatedElementIndexPlusOne);
            InsertIndex = AllocatedElementsBitArray.FindAndSetFirstZeroBit(FreeElementIndexHint);
            check(InsertIndex < MaxAllocatedElementIndexPlusOne);
        }
 
        FreeElementIndexHint = InsertIndex + 1;
        ++NumAllocatedElements;
 
        // Grow by one chunk if we run out of space
        if (MaxAllocatedElementIndexPlusOne > ElementChunks.Num() * ElementsPerChunk)
        {
            check(InsertIndex + 1 == MaxAllocatedElementIndexPlusOne);
            check(InsertIndex == ElementChunks.Num() * ElementsPerChunk);
            check(AllocatedElementsBitArray.Num() == ElementChunks.Num() * ElementsPerChunk);
 
            ElementChunks.Add((ElementType*)FMemory::Malloc(BytesPerChunk));
            AllocatedElementsBitArray.Add(false, ElementsPerChunk);
            AllocatedElementsBitArray[InsertIndex] = true;
        }
 
        check(IsAllocated(InsertIndex));
        const int32 ChunkIndex = InsertIndex / ElementsPerChunk;
        const int32 Index = InsertIndex - ChunkIndex * ElementsPerChunk;
        // We can construct in-place because the element is either uninitialized or has already been destructed
        new (&ElementChunks[ChunkIndex][Index]) ElementType;
 
        return InsertIndex;
    }
 
    void RemoveAt(int32 ElementIndex)
    {
        check(IsAllocated(ElementIndex));
        const int32 ChunkIndex = ElementIndex / ElementsPerChunk;
        const int32 Index = ElementIndex - ChunkIndex * ElementsPerChunk;
        ElementChunks[ChunkIndex][Index].~ElementType();
 
        AllocatedElementsBitArray[ElementIndex] = false;
 
        --NumAllocatedElements;
        FreeElementIndexHint = FMath::Min(FreeElementIndexHint, ElementIndex);
        if (ElementIndex + 1 == MaxAllocatedElementIndexPlusOne)
        {
            MaxAllocatedElementIndexPlusOne = AllocatedElementsBitArray.FindLastFrom(true, ElementIndex - 1) + 1;
        }
        check(FreeElementIndexHint <= MaxAllocatedElementIndexPlusOne);
    }
 
    void Shrink()
    {
        // Try to keep an additional chunk if the last used chunk is more than half full
        const int32 ChunksToKeep = FMath::DivideAndRoundUp(MaxAllocatedElementIndexPlusOne + ElementsPerChunk / 2, ElementsPerChunk);
        
        AllocatedElementsBitArray.SetNumUninitialized(FMath::Min(ElementChunks.Num(), ChunksToKeep) * ElementsPerChunk);
 
        while (ElementChunks.Num() > ChunksToKeep)
        {
            FMemory::Free(ElementChunks.Last());
            ElementChunks.Pop(EAllowShrinking::No);
        }
    }
 
    ElementType& operator[](int32 ElementIndex)
    {
        check(IsAllocated(ElementIndex));
        const int32 ChunkIndex = ElementIndex / ElementsPerChunk;
        const int32 Index = ElementIndex - ChunkIndex * ElementsPerChunk;
        return ElementChunks[ChunkIndex][Index];
    }
 
    const ElementType& operator[](int32 ElementIndex) const
    {
        check(IsAllocated(ElementIndex));
        const int32 ChunkIndex = ElementIndex / ElementsPerChunk;
        const int32 Index = ElementIndex - ChunkIndex * ElementsPerChunk;
        return ElementChunks[ChunkIndex][Index];
    }
 
    bool IsAllocated(int32 ElementIndex) const
    {
        if (ElementIndex < GetMaxSize())
        {
            return AllocatedElementsBitArray[ElementIndex];
        }
 
        return false;
    }
 
    SIZE_T GetAllocatedSize() const
    {
        return ElementChunks.Num() * BytesPerChunk + ElementChunks.GetAllocatedSize() + AllocatedElementsBitArray.GetAllocatedSize() + sizeof(TChunkedSparseArray);
    }
 
    class TRangedForIterator
    {
    public:
        TRangedForIterator(TChunkedSparseArray<ElementType>& InArray, int32 InElementIndex)
            : Array(InArray)
            , ElementIndex(InElementIndex)
        {
            // Scan for the first valid element.
            while (ElementIndex < Array.GetMaxSize() && !Array.AllocatedElementsBitArray[ElementIndex])
            {
                ++ElementIndex;
            }
        }
 
        TRangedForIterator operator++()
        {
            // Scan for the next first valid element.
            do
            {
                ++ElementIndex;
            } while (ElementIndex < Array.GetMaxSize() && !Array.AllocatedElementsBitArray[ElementIndex]);
 
            return *this;
        }
 
        bool operator!=(const TRangedForIterator& Other) const
        {
            return ElementIndex != Other.ElementIndex;
        }
 
        ElementType& operator*()
        {
            return Array[ElementIndex];
        }
 
    private:
        TChunkedSparseArray<ElementType>& Array;
        int32 ElementIndex;
    };
 
    class TRangedForConstIterator
    {
    public:
        TRangedForConstIterator(const TChunkedSparseArray<ElementType>& InArray, int32 InElementIndex)
            : Array(InArray)
            , ElementIndex(InElementIndex)
        {
            // Scan for the first valid element.
            while (ElementIndex < Array.GetMaxSize() && !Array.AllocatedElementsBitArray[ElementIndex])
            {
                ++ElementIndex;
            }
        }
 
        TRangedForConstIterator operator++()
        {
            // Scan for the next first valid element.
            do
            {
                ++ElementIndex;
            } while (ElementIndex < Array.GetMaxSize() && !Array.AllocatedElementsBitArray[ElementIndex]);
 
            return *this;
        }
 
        bool operator!=(const TRangedForConstIterator& Other) const
        {
            return ElementIndex != Other.ElementIndex;
        }
 
        const ElementType& operator*() const
        {
            return Array[ElementIndex];
        }
 
    private:
        const TChunkedSparseArray<ElementType>& Array;
        int32 ElementIndex;
    };
 
    // Iterate over all allocated elements (skip free ones)
    TRangedForIterator begin() { return TRangedForIterator(*this, 0); }
    TRangedForIterator end() { return TRangedForIterator(*this, GetMaxSize()); }
    TRangedForConstIterator begin() const { return TRangedForConstIterator(*this, 0); }
    TRangedForConstIterator end() const { return TRangedForConstIterator(*this, GetMaxSize()); }
 
private:
    static constexpr int32 ElementsPerChunk = BytesPerChunk / sizeof(ElementType);
 
    int32 FreeElementIndexHint = 0;
    int32 NumAllocatedElements = 0;
    int32 MaxAllocatedElementIndexPlusOne = 0;
 
    TArray<ElementType*> ElementChunks;
    TBitArray<> AllocatedElementsBitArray;
};