Files
UnrealEngine/Engine/Source/Runtime/AutoRTFM/Private/HitSet.h
2025-05-18 13:04:45 +08:00

294 lines
6.8 KiB
C++

// 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)