428 lines
12 KiB
C++
428 lines
12 KiB
C++
// Copyright Epic Games, Inc. All Rights Reserved.
|
|
|
|
#pragma once
|
|
|
|
#if (defined(__AUTORTFM) && __AUTORTFM)
|
|
|
|
#include "BitStack.h"
|
|
#include "ContainerValidation.h"
|
|
#include "Stack.h"
|
|
#include "Utils.h"
|
|
|
|
namespace AutoRTFM
|
|
{
|
|
|
|
struct FIntervalTree final
|
|
{
|
|
FIntervalTree() = default;
|
|
|
|
FIntervalTree(const FIntervalTree&) = delete;
|
|
FIntervalTree& operator=(const FIntervalTree&) = delete;
|
|
|
|
UE_AUTORTFM_FORCENOINLINE bool Insert(const void* const Address, const size_t Size)
|
|
{
|
|
FRange NewRange(Address, Size);
|
|
return Insert(std::move(NewRange));
|
|
}
|
|
|
|
UE_AUTORTFM_FORCENOINLINE bool Contains(const void* const Address, const size_t Size) const
|
|
{
|
|
const FRange NewRange(Address, Size);
|
|
|
|
if (AUTORTFM_UNLIKELY(IntervalTreeNodeIndexNone == Root))
|
|
{
|
|
return false;
|
|
}
|
|
|
|
FIntervalTreeNodeIndex Current = Root;
|
|
|
|
do
|
|
{
|
|
const FRange Range = NodeRanges[Current];
|
|
|
|
// This check does not need to prove that NewRange is entirely
|
|
// enclosed within Range, because if any byte of NewRange was in the
|
|
// original Range then it **must** already have been new memory.
|
|
if ((NewRange.Start < Range.End) && (Range.Start < NewRange.End))
|
|
{
|
|
return true;
|
|
}
|
|
|
|
const StackType<FIntervalTreeNodeIndex>& Next = (NewRange.Start < Range.End) ? NodeLefts : NodeRights;
|
|
Current = Next[Current];
|
|
} while (IntervalTreeNodeIndexNone != Current);
|
|
|
|
return false;
|
|
}
|
|
|
|
bool IsEmpty() const
|
|
{
|
|
return IntervalTreeNodeIndexNone == Root;
|
|
}
|
|
|
|
void Reset()
|
|
{
|
|
Root = IntervalTreeNodeIndexNone;
|
|
NodeRanges.Reset();
|
|
NodeLefts.Reset();
|
|
NodeRights.Reset();
|
|
NodeParents.Reset();
|
|
NodeColors.Reset();
|
|
}
|
|
|
|
void Merge(const FIntervalTree& Other)
|
|
{
|
|
if (Other.IsEmpty())
|
|
{
|
|
return;
|
|
}
|
|
|
|
StackType<FIntervalTreeNodeIndex> ToProcess;
|
|
|
|
ToProcess.Push(Other.Root);
|
|
|
|
do
|
|
{
|
|
const FIntervalTreeNodeIndex Current = ToProcess.Back();
|
|
ToProcess.Pop();
|
|
|
|
FRange Range = Other.NodeRanges[Current];
|
|
|
|
Insert(std::move(Range));
|
|
|
|
const FIntervalTreeNodeIndex Left = Other.NodeLefts[Current];
|
|
const FIntervalTreeNodeIndex Right = Other.NodeRights[Current];
|
|
|
|
if (IntervalTreeNodeIndexNone != Left)
|
|
{
|
|
ToProcess.Push(Left);
|
|
}
|
|
|
|
if (IntervalTreeNodeIndexNone != Right)
|
|
{
|
|
ToProcess.Push(Right);
|
|
}
|
|
} while (!ToProcess.IsEmpty());
|
|
}
|
|
|
|
private:
|
|
static constexpr size_t InlineArraySize = 8;
|
|
template<typename T> using StackType = TStack<T, InlineArraySize, EContainerValidation::Disabled>;
|
|
using BitStackType = TBitStack<InlineArraySize, EContainerValidation::Disabled>;
|
|
|
|
struct FRange final
|
|
{
|
|
FRange(const void* const Address, const size_t Size) :
|
|
Start(reinterpret_cast<uintptr_t>(Address)), End(reinterpret_cast<uintptr_t>(Address) + Size) {}
|
|
|
|
uintptr_t Start;
|
|
uintptr_t End;
|
|
};
|
|
|
|
using FIntervalTreeNodeIndex = uint32_t;
|
|
|
|
static constexpr FIntervalTreeNodeIndex IntervalTreeNodeIndexNone = UINT32_MAX;
|
|
|
|
enum EColor : bool
|
|
{
|
|
Red = false,
|
|
Black = true
|
|
};
|
|
|
|
FIntervalTreeNodeIndex Root = IntervalTreeNodeIndexNone;
|
|
StackType<FRange> NodeRanges;
|
|
StackType<FIntervalTreeNodeIndex> NodeLefts;
|
|
StackType<FIntervalTreeNodeIndex> NodeRights;
|
|
StackType<FIntervalTreeNodeIndex> NodeParents;
|
|
BitStackType NodeColors;
|
|
|
|
static constexpr bool bExtraDebugging = false;
|
|
|
|
UE_AUTORTFM_FORCEINLINE FIntervalTreeNodeIndex AddNode(FRange&& NewRange)
|
|
{
|
|
const FIntervalTreeNodeIndex Index = static_cast<FIntervalTreeNodeIndex>(NodeRanges.Num());
|
|
NodeRanges.Push(NewRange);
|
|
NodeLefts.Push(IntervalTreeNodeIndexNone);
|
|
NodeRights.Push(IntervalTreeNodeIndexNone);
|
|
NodeParents.Push(IntervalTreeNodeIndexNone);
|
|
NodeColors.Push(EColor::Black);
|
|
return Index;
|
|
}
|
|
|
|
UE_AUTORTFM_FORCEINLINE FIntervalTreeNodeIndex AddNode(FRange&& NewRange, FIntervalTreeNodeIndex Parent)
|
|
{
|
|
const FIntervalTreeNodeIndex Index = static_cast<FIntervalTreeNodeIndex>(NodeRanges.Num());
|
|
NodeRanges.Push(NewRange);
|
|
NodeLefts.Push(IntervalTreeNodeIndexNone);
|
|
NodeRights.Push(IntervalTreeNodeIndexNone);
|
|
NodeParents.Push(Parent);
|
|
NodeColors.Push(EColor::Red);
|
|
return Index;
|
|
}
|
|
|
|
UE_AUTORTFM_FORCENOINLINE bool Insert(FRange&& NewRange)
|
|
{
|
|
if (AUTORTFM_UNLIKELY(IntervalTreeNodeIndexNone == Root))
|
|
{
|
|
AUTORTFM_ASSERT(NodeRanges.IsEmpty());
|
|
Root = AddNode(std::move(NewRange));
|
|
return true;
|
|
}
|
|
|
|
FIntervalTreeNodeIndex Current = Root;
|
|
|
|
for(;;)
|
|
{
|
|
const FRange Range = NodeRanges[Current];
|
|
|
|
if (AUTORTFM_UNLIKELY((NewRange.Start < Range.End) && (Range.Start < NewRange.End)))
|
|
{
|
|
return false;
|
|
}
|
|
|
|
if (NewRange.Start < Range.Start)
|
|
{
|
|
if (AUTORTFM_UNLIKELY(NewRange.End == Range.Start))
|
|
{
|
|
AUTORTFM_ASSERT(NewRange.Start < Range.Start);
|
|
|
|
// We can just modify the existing node in place.
|
|
NodeRanges[Current].Start = NewRange.Start;
|
|
return true;
|
|
}
|
|
else if (IntervalTreeNodeIndexNone == NodeLefts[Current])
|
|
{
|
|
const FIntervalTreeNodeIndex Index = AddNode(std::move(NewRange), Current);
|
|
NodeLefts[Current] = Index;
|
|
Current = Index;
|
|
break;
|
|
}
|
|
|
|
Current = NodeLefts[Current];
|
|
}
|
|
else
|
|
{
|
|
if (AUTORTFM_UNLIKELY(NewRange.Start == Range.End))
|
|
{
|
|
AUTORTFM_ASSERT(NewRange.End > Range.End);
|
|
|
|
// We can just modify the existing node in place.
|
|
NodeRanges[Current].End = NewRange.End;
|
|
return true;
|
|
}
|
|
else if (IntervalTreeNodeIndexNone == NodeRights[Current])
|
|
{
|
|
const FIntervalTreeNodeIndex Index = AddNode(std::move(NewRange), Current);
|
|
NodeRights[Current] = Index;
|
|
Current = Index;
|
|
break;
|
|
}
|
|
|
|
Current = NodeRights[Current];
|
|
}
|
|
|
|
AUTORTFM_ASSERT(Root != Current);
|
|
}
|
|
|
|
auto IsBlack = [this](FIntervalTreeNodeIndex Index)
|
|
{
|
|
return (IntervalTreeNodeIndexNone == Index) || (EColor::Black == NodeColors[Index]);
|
|
};
|
|
|
|
for(;;)
|
|
{
|
|
FIntervalTreeNodeIndex Parent = NodeParents[Current];
|
|
|
|
UE_AUTORTFM_ASSUME(Current != Parent);
|
|
|
|
// The root will always have a black parent, so this check covers both.
|
|
if (Parent == Root)
|
|
{
|
|
NodeColors[Parent] = EColor::Black;
|
|
break;
|
|
}
|
|
else if (IsBlack(Parent))
|
|
{
|
|
break;
|
|
}
|
|
|
|
const FIntervalTreeNodeIndex GrandParent = NodeParents[Parent];
|
|
|
|
UE_AUTORTFM_ASSUME((Parent != GrandParent) && (Current != GrandParent));
|
|
|
|
const bool bParentIsLeft = NodeLefts[GrandParent] == Parent;
|
|
|
|
// The uncle is the other node of our parent.
|
|
const FIntervalTreeNodeIndex Uncle = bParentIsLeft ? NodeRights[GrandParent] : NodeLefts[GrandParent];
|
|
|
|
UE_AUTORTFM_ASSUME((GrandParent != Uncle) && (Parent != Uncle) && (Current != Uncle));
|
|
|
|
if (!IsBlack(Uncle))
|
|
{
|
|
AUTORTFM_ASSERT(IsBlack(GrandParent));
|
|
NodeColors[Parent] = EColor::Black;
|
|
NodeColors[Uncle] = EColor::Black;
|
|
|
|
if (GrandParent == Root)
|
|
{
|
|
break;
|
|
}
|
|
else
|
|
{
|
|
NodeColors[GrandParent] = EColor::Red;
|
|
}
|
|
|
|
Current = GrandParent;
|
|
continue;
|
|
}
|
|
|
|
// Our uncle is black, so we need to swizzle around.
|
|
const bool bCurrentIsLeft = NodeLefts[Parent] == Current;
|
|
|
|
if (bParentIsLeft)
|
|
{
|
|
if (!bCurrentIsLeft)
|
|
{
|
|
NodeRights[Parent] = NodeLefts[Current];
|
|
|
|
if (IntervalTreeNodeIndexNone != NodeRights[Parent])
|
|
{
|
|
NodeParents[NodeRights[Parent]] = Parent;
|
|
}
|
|
|
|
NodeParents[Parent] = Current;
|
|
NodeParents[Current] = GrandParent;
|
|
NodeLefts[Current] = Parent;
|
|
NodeLefts[GrandParent] = Current;
|
|
std::swap(Parent, Current);
|
|
}
|
|
|
|
NodeParents[Parent] = NodeParents[GrandParent];
|
|
NodeLefts[GrandParent] = NodeRights[Parent];
|
|
|
|
if (IntervalTreeNodeIndexNone != NodeLefts[GrandParent])
|
|
{
|
|
NodeParents[NodeLefts[GrandParent]] = GrandParent;
|
|
}
|
|
|
|
NodeRights[Parent] = GrandParent;
|
|
NodeParents[GrandParent] = Parent;
|
|
|
|
const bool bGrandParentIsBlack = NodeColors[GrandParent];
|
|
NodeColors[GrandParent] = NodeColors[Parent];
|
|
NodeColors[Parent] = bGrandParentIsBlack;
|
|
}
|
|
else
|
|
{
|
|
if (bCurrentIsLeft)
|
|
{
|
|
NodeLefts[Parent] = NodeRights[Current];
|
|
|
|
if (IntervalTreeNodeIndexNone != NodeLefts[Parent])
|
|
{
|
|
NodeParents[NodeLefts[Parent]] = Parent;
|
|
}
|
|
|
|
NodeParents[Parent] = Current;
|
|
NodeParents[Current] = GrandParent;
|
|
NodeRights[Current] = Parent;
|
|
NodeRights[GrandParent] = Current;
|
|
std::swap(Parent, Current);
|
|
}
|
|
|
|
NodeParents[Parent] = NodeParents[GrandParent];
|
|
NodeRights[GrandParent] = NodeLefts[Parent];
|
|
|
|
if (IntervalTreeNodeIndexNone != NodeRights[GrandParent])
|
|
{
|
|
NodeParents[NodeRights[GrandParent]] = GrandParent;
|
|
}
|
|
|
|
NodeLefts[Parent] = GrandParent;
|
|
NodeParents[GrandParent] = Parent;
|
|
|
|
const bool bGrandParentIsBlack = NodeColors[GrandParent];
|
|
NodeColors[GrandParent] = NodeColors[Parent];
|
|
NodeColors[Parent] = bGrandParentIsBlack;
|
|
}
|
|
|
|
if (IntervalTreeNodeIndexNone == NodeParents[Parent])
|
|
{
|
|
Root = Parent;
|
|
}
|
|
else
|
|
{
|
|
const FIntervalTreeNodeIndex GreatGrandParent = NodeParents[Parent];
|
|
|
|
if (GrandParent == NodeLefts[GreatGrandParent])
|
|
{
|
|
NodeLefts[GreatGrandParent] = Parent;
|
|
}
|
|
else
|
|
{
|
|
NodeRights[GreatGrandParent] = Parent;
|
|
}
|
|
}
|
|
|
|
break;
|
|
}
|
|
|
|
AssertStructureIsOk();
|
|
|
|
return true;
|
|
}
|
|
|
|
UE_AUTORTFM_FORCEINLINE void AssertStructureIsOk() const
|
|
{
|
|
if constexpr (bExtraDebugging)
|
|
{
|
|
if (IntervalTreeNodeIndexNone != Root)
|
|
{
|
|
AssertNodeIsOk(Root);
|
|
}
|
|
}
|
|
}
|
|
|
|
UE_AUTORTFM_FORCENOINLINE void AssertNodeIsOk(FIntervalTreeNodeIndex Index) const
|
|
{
|
|
// We need to use recursion to check because we cannot have any
|
|
// allocations within the checker itself!
|
|
if constexpr (bExtraDebugging)
|
|
{
|
|
AUTORTFM_ASSERT(IntervalTreeNodeIndexNone != Index);
|
|
AUTORTFM_ASSERT(Index < NodeRanges.Num());
|
|
|
|
const FIntervalTreeNodeIndex Parent = NodeParents[Index];
|
|
const FIntervalTreeNodeIndex Left = NodeLefts[Index];
|
|
const FIntervalTreeNodeIndex Right = NodeRights[Index];
|
|
|
|
if (IntervalTreeNodeIndexNone == Parent)
|
|
{
|
|
AUTORTFM_ASSERT(Root == Index);
|
|
}
|
|
else
|
|
{
|
|
AUTORTFM_ASSERT(NodeColors[Parent] || NodeColors[Index]);
|
|
AUTORTFM_ASSERT((NodeLefts[Parent] == Index) ^ (NodeRights[Parent] == Index));
|
|
}
|
|
|
|
AUTORTFM_ASSERT(Left != Index);
|
|
AUTORTFM_ASSERT(Right != Index);
|
|
|
|
if (IntervalTreeNodeIndexNone != Left)
|
|
{
|
|
AssertNodeIsOk(Left);
|
|
}
|
|
|
|
if (IntervalTreeNodeIndexNone != Right)
|
|
{
|
|
AssertNodeIsOk(Right);
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
} // namespace AutoRTFM
|
|
|
|
#endif // (defined(__AUTORTFM) && __AUTORTFM)
|