SparseBitSet.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
#pragma once
 
#include "CoreTypes.h"
#include "Templates/UnrealTemplate.h"
#include "Containers/Array.h"
 
namespace UE::MovieScene
{
 
template<typename T, int32 InlineSize>
struct TDynamicSparseBitSetBucketStorage;
 
template<typename T>
struct TFixedSparseBitSetBucketStorage;
 
template<typename HashType, typename BucketStorage = TDynamicSparseBitSetBucketStorage<uint8, 4>>
struct TFixedSparseBitSet;
 
template<typename HashType, typename BucketStorage = TDynamicSparseBitSetBucketStorage<uint8, 4>>
struct TDynamicSparseBitSet;
 
enum class ESparseBitSetBitResult
{
    NewlySet,
    AlreadySet,
};
 
namespace Private
{
    static uint32 CountTrailingZeros(uint8 In)
    {
        const uint32 X = 0xFFFFFF00 | uint32(In);
        return FMath::CountTrailingZeros(X);
    }
    static uint32 CountTrailingZeros(uint16 In)
    {
        const uint32 X = 0xFFFF0000 | uint32(In);
        return FMath::CountTrailingZeros(X);
    }
    static uint32 CountTrailingZeros(uint32 In)
    {
        return FMath::CountTrailingZeros(In);
    }
    static uint32 CountTrailingZeros(uint64 In)
    {
        return static_cast<uint32>(FMath::CountTrailingZeros64(In));
    }
 
    template<typename T, typename U>
    static T BitOffsetToLowBitMask(U BitOffset)
    {
        constexpr T One(1);
        const T Index = static_cast<T>(BitOffset);
        return (One << Index)-One;
    }
 
    template<typename T, typename U>
    static T BitOffsetToHighBitMask(U BitOffset)
    {
        return ~BitOffsetToLowBitMask<T>(BitOffset);
    }
}
 
template<typename HashType, typename BucketStorage>
struct TFixedSparseBitSet
{
    using BucketType = typename BucketStorage::BucketType;
    static constexpr uint32 HashSize    = sizeof(HashType)*8;
    static constexpr uint32 BucketSize  = sizeof(typename BucketStorage::BucketType)*8;
    static constexpr uint32 MaxNumBits  = HashSize * BucketSize;
 
    explicit TFixedSparseBitSet()
        : BucketHash(0)
    {}
 
    template<typename ...StorageArgs>
    explicit TFixedSparseBitSet(StorageArgs&& ...Storage)
        : Buckets(Forward<StorageArgs>(Storage)...)
        , BucketHash(0)
    {}
 
    TFixedSparseBitSet(const TFixedSparseBitSet&) = default;
    TFixedSparseBitSet& operator=(const TFixedSparseBitSet&) = default;
 
    TFixedSparseBitSet(TFixedSparseBitSet&&) = default;
    TFixedSparseBitSet& operator=(TFixedSparseBitSet&&) = default;
 
    template<typename OtherHashType, typename OtherStorageType>
    void CopyTo(TFixedSparseBitSet<OtherHashType, OtherStorageType>& Other) const
    {
        static_assert(TFixedSparseBitSet<OtherHashType, OtherStorageType>::BucketSize == BucketSize, "Cannot copy sparse bitsets of different bucket sizes");
 
        // Copy the buckets
        const uint32 NumBuckets = FMath::CountBits(Other.BucketHash);
        Other.Buckets.SetNum(NumBuckets);
        CopyToUnsafe(Other, NumBuckets);
    }
 
    template<typename OtherHashType, typename OtherStorageType>
    void CopyToUnsafe(TFixedSparseBitSet<OtherHashType, OtherStorageType>& Other, uint32 OtherBucketCapacity) const
    {
        static_assert(TFixedSparseBitSet<OtherHashType, OtherStorageType>::BucketSize == BucketSize, "Cannot copy sparse bitsets of different bucket sizes");
 
        const uint32 ThisNumBuckets = this->NumBuckets();
        checkf(OtherBucketCapacity >= ThisNumBuckets, TEXT("Attempting to copy a sparse bitset without enough capacity in the destination (%d, required %d)"), OtherBucketCapacity, ThisNumBuckets);
 
        // Copy the hash
        Other.BucketHash = this->BucketHash;
 
        // Copy the buckets
        FMemory::Memcpy(Other.Buckets.GetData(), this->Buckets.Storage.GetData(), sizeof(typename BucketStorage::BucketType)*ThisNumBuckets);
    }
 
    TFixedSparseBitSet& operator|=(const TFixedSparseBitSet& Other)
    {
        using namespace Private;
 
        HashType One(1u);
        HashType NewHash   = Other.BucketHash | BucketHash;
        HashType OtherHash = Other.BucketHash;
 
        uint32 OtherBucketIndex    = 0;
        uint32 OtherBucketBitIndex = Private::CountTrailingZeros(OtherHash);
 
        while (OtherBucketBitIndex < HashSize)
        {
            const HashType HashBit = HashType(1) << OtherBucketBitIndex;
            const uint32 ThisBucketIndex = FMath::CountBits(BucketHash & (HashBit-1));
 
            if ((BucketHash & HashBit) == 0)
            {
                Buckets.Insert(Other.Buckets.Get(OtherBucketIndex), ThisBucketIndex);
            }
            else
            {
                Buckets.Get(ThisBucketIndex) |= Other.Buckets.Get(OtherBucketIndex);
            }
            
            BucketHash |= HashBit;
 
            ++OtherBucketIndex;
 
            // Mask out this bit and find the index of the next one
            OtherHash &= ~(One << OtherBucketBitIndex);
            OtherBucketBitIndex = Private::CountTrailingZeros(OtherHash);
        }
 
        return *this;
    }
 
    uint32 NumBuckets() const
    {
        return FMath::CountBits(this->BucketHash);
    }
 
    uint32 CountSetBits() const
    {
        uint32 Total = 0;
        const uint32 TotalNumBuckets = NumBuckets();
        for (uint32 Index = 0; Index < TotalNumBuckets; ++Index)
        {
            Total += FMath::CountBits(Buckets.Get(Index));
        }
        return Total;
    }
 
    uint32 GetMaxNumBits() const
    {
        return MaxNumBits;
    }
 
    bool IsEmpty() const
    {
        return this->BucketHash == 0;
    }
 
    ESparseBitSetBitResult SetBit(uint32 BitIndex)
    {
        CheckIndex(BitIndex);
 
        FBitOffsets Offsets(BucketHash, BitIndex);
 
        // Do we need to add a new bucket?
        if ( (BucketHash & Offsets.HashBit) == 0)
        {
            BucketHash |= Offsets.HashBit;
            Buckets.Insert(Offsets.BitMaskWithinBucket, Offsets.BucketIndex);
            return ESparseBitSetBitResult::NewlySet;
        }
        else if ((Buckets.Get(Offsets.BucketIndex) & Offsets.BitMaskWithinBucket) == 0)
        {
            Buckets.Get(Offsets.BucketIndex) |= Offsets.BitMaskWithinBucket;
            return ESparseBitSetBitResult::NewlySet;
        }
 
        return ESparseBitSetBitResult::AlreadySet;
    }
 
    bool IsBitSet(uint32 BitIndex) const
    {
        CheckIndex(BitIndex);
 
        const uint32   Hash    = BitIndex / BucketSize;
        const HashType HashBit = (HashType(1) << Hash);
        if (BucketHash & HashBit)
        {
            const uint32     BucketIndex  = FMath::CountBits(BucketHash & (HashBit-1));
            const uint32     ThisBitIndex = (BitIndex-BucketSize*Hash);
            const BucketType ThisBitMask  = BucketType(1u) << ThisBitIndex;
 
            return Buckets.Get(BucketIndex) & ThisBitMask;
        }
        return false;
    }
 
    int32 GetSparseBucketIndex(uint32 BitIndex) const
    {
        CheckIndex(BitIndex);
 
        const uint32   Hash    = BitIndex / BucketSize;
        const HashType HashBit = (HashType(1) << Hash);
        if (BucketHash & HashBit)
        {
            uint32 BucketIndex = FMath::CountBits(BucketHash & (HashBit-1));
 
            const uint32 ThisBitIndex = (BitIndex-BucketSize*Hash);
            const BucketType ThisBitMask = static_cast<BucketType>(BucketType(1u) << ThisBitIndex);
 
            BucketType ThisBucket = Buckets.Get(BucketIndex);
            if (ThisBucket & ThisBitMask)
            {
                // Compute the offset
                int32 SparseIndex = FMath::CountBits(ThisBucket & (ThisBitMask-1));
 
                // Count all the preceding buckets to find the final sparse index
                while (BucketIndex > 0)
                {
                    --BucketIndex;
                    SparseIndex += FMath::CountBits(Buckets.Get(BucketIndex));
                }
                return SparseIndex;
            }
        }
        return INDEX_NONE;
    }
 
    struct FIterator
    {
        FIterator()
            : BitSet(nullptr)
            , BucketBitIndex(0)
            , IndexWithinBucket(0)
            , CurrentBucket(0)
        {
        }
 
        static FIterator Begin(const TFixedSparseBitSet<HashType, BucketStorage>* InBitSet)
        {
            FIterator It;
            It.BitSet = InBitSet;
            It.CurrentBucket = 0;
 
            if (InBitSet->BucketHash != 0)
            {
                It.BucketBitIndex = Private::CountTrailingZeros(InBitSet->BucketHash);
                It.CurrentBucket = InBitSet->Buckets.Get(0);
                It.IndexWithinBucket = Private::CountTrailingZeros(It.CurrentBucket);
            }
            else
            {
                It.BucketBitIndex = HashSize;
                It.IndexWithinBucket = 0;
            }
            return It;
        }
        static FIterator End(const TFixedSparseBitSet<HashType, BucketStorage>* InBitSet)
        {
            FIterator It;
            It.BitSet = InBitSet;
            It.CurrentBucket = 0;
            It.BucketBitIndex = HashSize;
            It.IndexWithinBucket = 0;
            return It;
        }
 
        void operator++()
        {
            using namespace Private;
 
            // Clear the lowest 1 bit
            CurrentBucket = CurrentBucket & (CurrentBucket - 1);
 
            if (CurrentBucket != 0)
            {
                IndexWithinBucket = CountTrailingZeros(CurrentBucket);
            }
            else
            {
                // If this was the last bit, reset the iterator to end()
                if (BucketBitIndex == HashSize-1)
                {
                    IndexWithinBucket = 0;
                    BucketBitIndex = HashSize;
                    return;
                }
 
                HashType UnvisitedBucketBitMask = BitOffsetToHighBitMask<HashType>(BucketBitIndex+1);
                BucketBitIndex = CountTrailingZeros(HashType(BitSet->BucketHash & UnvisitedBucketBitMask));
 
                // Check whether we're at the end
                if (BucketBitIndex == HashSize)
                {
                    IndexWithinBucket = 0;
                }
                else
                {
                    const uint8 NextBucketIndex = FMath::CountBits(BitSet->BucketHash & BitOffsetToLowBitMask<HashType>(BucketBitIndex));
                    CurrentBucket = BitSet->Buckets.Get(NextBucketIndex);
                    IndexWithinBucket = CountTrailingZeros(CurrentBucket);
                }
            }
        }
 
        int32 operator*() const
        {
            return BucketSize*BucketBitIndex + IndexWithinBucket;
        }
 
        explicit operator bool() const
        {
            return BucketBitIndex < HashSize;
        }
 
        friend bool operator==(const FIterator& A, const FIterator& B)
        {
            return A.BitSet == B.BitSet && A.BucketBitIndex == B.BucketBitIndex && A.IndexWithinBucket == B.IndexWithinBucket;
        }
        friend bool operator!=(const FIterator& A, const FIterator& B)
        {
            return !(A == B);
        }
private:
 
        const TFixedSparseBitSet<HashType, BucketStorage>* BitSet;
        uint8 BucketBitIndex;
        uint8 IndexWithinBucket;
 
        BucketType CurrentBucket;
    };
 
    friend FIterator begin(const TFixedSparseBitSet<HashType, BucketStorage>& In) { return FIterator::Begin(&In); }
    friend FIterator end(const TFixedSparseBitSet<HashType, BucketStorage>& In)   { return FIterator::End(&In); }
 
private:
 
    template<typename, typename>
    friend struct TFixedSparseBitSet;
 
    inline void CheckIndex(uint32 BitIndex) const
    {
        checkfSlow(BitIndex < MaxNumBits, TEXT("Invalid index (%d) specified for a sparse bitset of maximum size (%d)"), BitIndex, MaxNumBits);
    }
 
    struct FBitOffsets
    {
        HashType HashBit;
        BucketType BitMaskWithinBucket;
        int32  BucketIndex;
        FBitOffsets(HashType InBucketHash, uint32 BitIndex)
        {
            const HashType Hash(BitIndex / BucketSize);
            HashBit = HashType(1) << Hash;
 
            BucketIndex = FMath::CountBits(InBucketHash & (HashBit-1u));
 
            const uint32 ThisBitIndex = (BitIndex-BucketSize*Hash);
            BitMaskWithinBucket = BucketType(1u) << ThisBitIndex;
        }
    };
 
    BucketStorage Buckets;
    HashType      BucketHash;
};
 
 
 
template<typename T, int32 InlineSize = 8>
struct TDynamicSparseBitSetBucketStorage
{
    using BucketType = T;
 
    TArray<BucketType, TInlineAllocator<InlineSize>> Storage;
 
    void Insert(BucketType InitialValue, int32 Index)
    {
        Storage.Insert(InitialValue, Index);
    }
 
    BucketType* GetData()              { return Storage.GetData(); }
 
    BucketType& Get(int32 Index)       { return Storage[Index]; }
    BucketType  Get(int32 Index) const { return Storage[Index]; }
};
 
template<typename T>
struct TDynamicSparseBitSetBucketStorage<T, 0>
{
    using BucketType = T;
 
    TArray<BucketType> Storage;
 
    void Insert(BucketType InitialValue, int32 Index)
    {
        Storage.Insert(InitialValue, Index);
    }
 
    BucketType* GetData()              { return Storage.GetData(); }
 
    BucketType& Get(int32 Index)       { return Storage[Index]; }
    BucketType  Get(int32 Index) const { return Storage[Index]; }
};
 
template<typename T>
struct TFixedSparseBitSetBucketStorage
{
    using BucketType = T;
 
    explicit TFixedSparseBitSetBucketStorage()
        : Storage(nullptr)
    {}
 
    explicit TFixedSparseBitSetBucketStorage(BucketType* StoragePtr)
        : Storage(StoragePtr)
    {}
 
    TFixedSparseBitSetBucketStorage(const TFixedSparseBitSetBucketStorage&) = delete;
    void operator=(const TFixedSparseBitSetBucketStorage&) = delete;
 
    TFixedSparseBitSetBucketStorage(TFixedSparseBitSetBucketStorage&&) = delete;
    void operator=(TFixedSparseBitSetBucketStorage&&) = delete;
 
    BucketType* Storage;
 
    BucketType* GetData()              { return Storage; }
 
    BucketType& Get(int32 Index)       { return Storage[Index]; }
    BucketType  Get(int32 Index) const { return Storage[Index]; }
};
 
template<typename HashType, typename BucketStorage>
struct TDynamicSparseBitSet
{
    uint32 GetMaxNumBits() const
    {
        return MAX_uint32;
    }
 
 
    ESparseBitSetBitResult SetBit(uint32 Bit)
    {
        const uint32 Bucket = Bit / NumBitsInBucket;
 
        Bit -= Bucket*NumBitsInBucket;
 
        FEntry* EntrPtr = Entries.GetData();
 
        const int32 Num = Entries.Num();
        for (int32 EntryIndex = 0; EntryIndex < Num; ++EntryIndex)
        {
            if (EntrPtr[EntryIndex].Offset == Bucket)
            {
                return EntrPtr[EntryIndex].Bits.SetBit(Bit);
            }
            else if (EntrPtr[EntryIndex].Offset > Bucket)
            {
                Entries.InsertUninitialized(EntryIndex);
                new(&Entries[EntryIndex]) FEntry(Bucket, Bit);
                return ESparseBitSetBitResult::NewlySet;
            }
        }
 
        Entries.Emplace(Bucket, Bit);
        return ESparseBitSetBitResult::NewlySet;
    }
 
 
    bool IsEmpty() const
    {
        return Entries.Num() == 0;
    }
 
 
    bool IsBitSet(uint32 Bit) const
    {
        const uint32 Bucket = Bit / NumBitsInBucket;
 
        const FEntry* EntrPtr = Entries.GetData();
 
        const int32 Num = Entries.Num();
        for (int32 EntryIndex = 0; EntryIndex < Num; ++EntryIndex)
        {
            if (EntrPtr[EntryIndex].Offset == Bucket)
            {
                return EntrPtr[EntryIndex].Bits.IsBitSet(Bit);
            }
            if (EntrPtr[EntryIndex].Offset > Bucket)
            {
                return false;
            }
        }
 
        return false;
    }
 
    uint32 CountSetBits() const
    {
        uint32 SetBits = 0;
        for (const FEntry& Entry : Entries)
        {
            SetBits += Entry.Bits.CountSetBits();
        }
        return SetBits;
    }
 
 
    TDynamicSparseBitSet<HashType, BucketStorage>& operator|=(const TDynamicSparseBitSet<HashType, BucketStorage>& Other)
    {
        if (Other.Entries.Num() == 0)
        {
            return *this;
        }
 
        if (Entries.Num() == 0)
        {
            *this = Other;
            return *this;
        }
 
 
        int32 ThisIndex = 0;
        int32 OtherIndex = 0;
 
        while (OtherIndex < Other.Entries.Num() && Other.Entries[OtherIndex].Offset < Entries[0].Offset)
        {
            ++OtherIndex;
        }
 
        if (OtherIndex > 0)
        {
            Entries.Insert(Other.Entries.GetData(), OtherIndex, 0);
        }
 
        while (OtherIndex < Other.Entries.Num() && ThisIndex < Entries.Num())
        {
            if (Other.Entries[OtherIndex].Offset < Entries[ThisIndex].Offset)
            {
                Entries.Insert(Other.Entries[OtherIndex], ThisIndex);
                ++OtherIndex;
            }
            else if (Other.Entries[OtherIndex].Offset == Entries[ThisIndex].Offset)
            {
                Entries[ThisIndex].Bits |= Other.Entries[OtherIndex].Bits;
                ++OtherIndex;
                ++ThisIndex;
            }
            else
            {
                ++ThisIndex;
            }
        }
 
        return *this;
    }
 
 
private:
 
    using FBucketBitSet = TFixedSparseBitSet<HashType, BucketStorage>;
 
    static constexpr uint32 NumBitsInBucket = FBucketBitSet::MaxNumBits;
 
    struct FEntry
    {
        FEntry(uint32 InOffset)
            : Offset(InOffset)
        {
        }
        FEntry(uint32 InOffset, uint32 InBit)
            : Offset(InOffset)
        {
            checkSlow(InBit < FBucketBitSet::MaxNumBits);
            Bits.SetBit(InBit);
        }
 
        FBucketBitSet Bits;
        uint32 Offset;
    };
 
public:
 
    struct FIterator
    {
        static FIterator Begin(const TDynamicSparseBitSet<HashType, BucketStorage>* InBitSet)
        {
            FIterator It;
            It.Entries = InBitSet->Entries.GetData();
            It.NumEntries = InBitSet->Entries.Num();
            It.EntryIndex = 0;
            It.CurrentOffsetInBits = 0;
 
            if (It.NumEntries != 0)
            {
                It.CurrentOffsetInBits = It.Entries[0].Offset * NumBitsInBucket;
                It.BucketIt = BucketIterator::Begin(&It.Entries[0].Bits);
            }
            return It;
        }
        static FIterator End(const TDynamicSparseBitSet<HashType, BucketStorage>* InBitSet)
        {
            FIterator It;
            It.Entries = InBitSet->Entries.GetData();
            It.NumEntries = InBitSet->Entries.Num();
            It.EntryIndex = It.NumEntries;
            It.CurrentOffsetInBits = 0;
            return It;
        }
 
        void operator++()
        {
            using namespace Private;
 
            ++BucketIt;
            if (!BucketIt)
            {
                ++EntryIndex;
                if (EntryIndex < NumEntries)
                {
                    CurrentOffsetInBits = Entries[EntryIndex].Offset * NumBitsInBucket;
                    BucketIt = BucketIterator::Begin(&Entries[EntryIndex].Bits);
                }
                else
                {
                    CurrentOffsetInBits = 0;
                    BucketIt = BucketIterator();
                }
            }
        }
 
        int32 operator*() const
        {
            return CurrentOffsetInBits + *BucketIt;
        }
 
        explicit operator bool() const
        {
            return EntryIndex < NumEntries;
        }
 
        friend bool operator==(const FIterator& A, const FIterator& B)
        {
            return A.Entries == B.Entries && A.EntryIndex == B.EntryIndex && A.BucketIt == B.BucketIt;
        }
        friend bool operator!=(const FIterator& A, const FIterator& B)
        {
            return !(A == B);
        }
private:
        FIterator() = default;
 
        using BucketIterator = typename TFixedSparseBitSet<HashType, BucketStorage>::FIterator;
 
        const typename TDynamicSparseBitSet<HashType, BucketStorage>::FEntry* Entries;
        BucketIterator BucketIt;
        int32 NumEntries;
        int32 EntryIndex;
        int32 CurrentOffsetInBits;
    };
 
    friend FIterator begin(const TDynamicSparseBitSet<HashType, BucketStorage>& In) { return FIterator::Begin(&In); }
    friend FIterator end(const TDynamicSparseBitSet<HashType, BucketStorage>& In)   { return FIterator::End(&In); }
 
    TArray<FEntry> Entries;
};
 
} // namespace UE::MovieScene