HitSet.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
#pragma once
 
#if (defined(__AUTORTFM) && __AUTORTFM)
 
#include "ExternAPI.h"
#include "Utils.h"
 
#include <cstdint>
 
namespace AutoRTFM
{
    struct FHitSetEntry
    {
        // The address of the write.
        uintptr_t Address : 48;
        // Size of the write.
        uintptr_t Size : 15;
        // If set, then the write shouldn't be considered by the memory validator.
        uintptr_t bNoMemoryValidation : 1;
        
        UE_AUTORTFM_FORCEINLINE uintptr_t Payload() const
        {
            // Note: Clang optimizes this away to a reinterpret cast.
            uintptr_t Result;
            memcpy(&Result, this, sizeof(Result));
            return Result;
        }
    };
 
    static_assert(sizeof(FHitSetEntry) == sizeof(uintptr_t));
    
    class FHitSet final
    {
    public:
        // Maximum size in bytes for a hit-set write record.
        // The cutoff here is arbitrarily any number less than UINT16_MAX, but its a
        // weigh up what a good size is. Because the hitset doesn't detect when you
        // are trying to write to a subregion of a previous hit (like memset something,
        // then write to an individual element), we've got to balance the cost of
        // recording meaningless hits, against the potential to hit again.
        static constexpr size_t MaxSize = 16;
 
    private:
        static constexpr uint8_t LinearProbeDepth = 32;
        static constexpr uint8_t LogInitialCapacity = 6;
        static constexpr uint64_t InitialCapacity = static_cast<uint64_t>(1) << LogInitialCapacity;
 
        // TODO: Revisit a good max capacity for the hitset.
        static constexpr uint64_t MaxCapacity = static_cast<uint64_t>(128 * 1024 * 1024) / sizeof(FHitSetEntry);
 
        static_assert(LinearProbeDepth <= InitialCapacity);
    public:
 
        ~FHitSet()
        {
            Reset();
        }
 
        enum class EInsertResult : uint8_t
        {
            Exists,
            Inserted,
            NotInserted
        };
 
        // Insert something in the HitSet, returning whether the key already
        // exists, it was inserted, or not.
        UE_AUTORTFM_FORCEINLINE EInsertResult FindOrTryInsertNoResize(FHitSetEntry Entry)
        {
            return InsertNoResize(Entry.Payload());
        }
 
        // Insert something in the HitSet, returning true if we found the key.
        // If false is returned, we could have inserted or not - depending on
        // whether the HitSet has reached its max capacity.
        UE_AUTORTFM_FORCEINLINE bool FindOrTryInsert(FHitSetEntry Entry)
        {
            while (true)
            {
                switch (InsertNoResize(Entry.Payload()))
                {
                default:
                    AutoRTFM::InternalUnreachable();
                case EInsertResult::Inserted:
                    return false;
                case EInsertResult::Exists:
                    return true;
                case EInsertResult::NotInserted:
                    if (!Resize())
                    {
                        return false;
                    }
                    continue;
                }
            }
        }
 
        UE_AUTORTFM_FORCEINLINE bool IsEmpty() const
        {
            return 0 == Count;
        }
 
        // Clear out the data stored in the set, but does not reduce the capacity.
        void Reset()
        {
            if (Count)
            {
                if (Payload != SmallPayload)
                {
                    AutoRTFM::Free(Payload);
                    Payload = SmallPayload;
                }
                memset(SmallPayload, 0, sizeof(SmallPayload));
                SixtyFourMinusLogCapacity = 64 - LogInitialCapacity;
                Count = 0;
            }
        }
 
        uint64_t GetCapacity() const
        {
            return Capacity();
        }
 
        uint64_t GetCount() const
        {
            return Count;
        }
 
    private:
        uintptr_t SmallPayload[InitialCapacity]{};
        uintptr_t* Payload = SmallPayload;
        uint64_t Count = 0;
        uint32_t SixtyFourMinusLogCapacity = 64 - LogInitialCapacity;
 
        UE_AUTORTFM_FORCEINLINE uint64_t Capacity() const
        {
            UE_AUTORTFM_ASSUME(SixtyFourMinusLogCapacity < 64);
            return static_cast<uint64_t>(1) << (64 - SixtyFourMinusLogCapacity);
        }
 
        UE_AUTORTFM_FORCEINLINE void IncreaseCapacity()
        {
            // It seems odd - but subtracting 1 here is totally intentional
            // because of how we store our capacity as (64 - log2(Capacity)).
            SixtyFourMinusLogCapacity -= 1;
 
            // Check that we haven't overflowed our capacity!
            AUTORTFM_ASSERT(0 != SixtyFourMinusLogCapacity);
        }
 
        bool Resize()
        {
            const uint64_t OldCapacity = Capacity();
 
            if (OldCapacity >= MaxCapacity)
            {
                return false;
            }
 
            uintptr_t* const OldPayload = Payload;
            const uint64_t OldCount = Count;
 
            while (true)
            {
                bool bNeedAnotherResize = false;
 
                Count = 0;
 
                IncreaseCapacity();
 
                Payload = static_cast<uintptr_t*>(AutoRTFM::AllocateZeroed(Capacity() * sizeof(uintptr_t), alignof(uintptr_t)));
                AUTORTFM_ASSERT(nullptr != Payload);
 
                // Now we need to rehash and reinsert all the items.
                for (size_t I = 0; I < OldCapacity; I++)
                {
                    // Skip empty locations.
                    if (0 == OldPayload[I])
                    {
                        continue;
                    }
 
                    const EInsertResult Result = InsertNoResize(OldPayload[I]);
 
                    if (EInsertResult::NotInserted == Result)
                    {
                        bNeedAnotherResize = true;
                        break;
                    }
                    else
                    {
                        AUTORTFM_ASSERT(EInsertResult::Inserted == Result);
                    }
                }
 
                if (bNeedAnotherResize)
                {
                    if (Payload != SmallPayload)
                    {
                        AutoRTFM::Free(Payload);
                    }
                    continue;
                }
 
                break;
            }
 
            AUTORTFM_ASSERT(OldCount == Count);
 
            if (OldPayload != SmallPayload)
            {
                AutoRTFM::Free(OldPayload);
            }
 
            return true;
        }
 
        UE_AUTORTFM_FORCEINLINE uintptr_t FibonacciHash(const uintptr_t Hashee) const
        {
            constexpr uintptr_t Fibonacci = 0x9E3779B97F4A7C15;
            return (Hashee * Fibonacci) >> SixtyFourMinusLogCapacity;
        }
 
        UE_AUTORTFM_FORCEINLINE EInsertResult TryInsertAtIndex(const uintptr_t Raw, const uintptr_t I)
        {
            // We have a free location in the set.
            if (0 == Payload[I])
            {
                Payload[I] = Raw;
                Count++;
                return EInsertResult::Inserted;
            }
 
            // We're already in the set.
            if (Raw == Payload[I])
            {
                return EInsertResult::Exists;
            }
 
            return EInsertResult::NotInserted;
        }
 
        // Insert something in the HitSet, returning true if the insert
        // succeeded (EG. the key was not already in the set).
        UE_AUTORTFM_FORCEINLINE EInsertResult InsertNoResize(const uintptr_t Raw)
        {
            const uintptr_t Hash = FibonacciHash(Raw);
 
            // The first index will always be in the range of capacity because we
            // are using a fibonacci hash. So just check if we can insert here.
            {
                const uintptr_t I = Hash;
 
                const EInsertResult R = TryInsertAtIndex(Raw, I);
 
                if (EInsertResult::NotInserted != R)
                {
                    return R;
                }
 
                // Otherwise we need to do a linear probe...
            }
 
            const uintptr_t Mask = Capacity() - 1;
 
            // Clang goes way over the top with loop unrolling, and makes this code noticably
            // slower as a result. So we disable it!
#ifdef __clang__
#pragma clang loop unroll(disable)
#endif
            for (uint8_t D = 1; D < LinearProbeDepth; D++)
            {
                // Capacity is always a power of 2, so we can just mask out the bits.
                const uintptr_t I = (Hash + D) & Mask;
 
                const EInsertResult R = TryInsertAtIndex(Raw, I);
 
                if (EInsertResult::NotInserted != R)
                {
                    return R;
                }
 
                // Otherwise we need to keep going with our linear probe...
            }
 
            return EInsertResult::NotInserted;
        }
    };
}
 
#endif // (defined(__AUTORTFM) && __AUTORTFM)