diff --git a/Source/Vmt/VmtCopy.h b/Source/Vmt/VmtCopy.h index 046f2189786..08fa6b8bd3f 100644 --- a/Source/Vmt/VmtCopy.h +++ b/Source/Vmt/VmtCopy.h @@ -4,7 +4,7 @@ #include #include -#include +#include #include #include "VmtLength.h" @@ -12,8 +12,7 @@ class VmtCopy { public: VmtCopy(std::uintptr_t* vmt, VmtLength length) noexcept : originalVmt{ vmt } - , length{ static_cast(length) } - , replacementVmtWithTypeInfo{ allocateReplacementVmtWithTypeInfo() } + , replacementVmtWithTypeInfo{ mem::makeUniqueForOverwrite(static_cast(length) + platform::lengthOfTypeInfoPrecedingVmt) } { copyOriginalVmt(); } @@ -21,7 +20,7 @@ class VmtCopy { [[nodiscard]] std::uintptr_t* getReplacementVmt() const noexcept { if (replacementVmtWithTypeInfo) [[likely]] - return replacementVmtWithTypeInfo + platform::lengthOfTypeInfoPrecedingVmt; + return replacementVmtWithTypeInfo.get() + platform::lengthOfTypeInfoPrecedingVmt; return nullptr; } @@ -30,30 +29,18 @@ class VmtCopy { return originalVmt; } - ~VmtCopy() noexcept - { - if (replacementVmtWithTypeInfo) - MemoryAllocator::deallocate(reinterpret_cast(replacementVmtWithTypeInfo), MemoryAllocator::memoryFor(lengthWithTypeInfo())); - } - private: - [[nodiscard]] std::uintptr_t* allocateReplacementVmtWithTypeInfo() const noexcept - { - return new (MemoryAllocator::allocate(MemoryAllocator::memoryFor(lengthWithTypeInfo()))) std::uintptr_t[lengthWithTypeInfo()]; - } - void copyOriginalVmt() const noexcept { if (replacementVmtWithTypeInfo) [[likely]] - std::copy_n(originalVmt - platform::lengthOfTypeInfoPrecedingVmt, lengthWithTypeInfo(), replacementVmtWithTypeInfo); + std::copy_n(originalVmt - platform::lengthOfTypeInfoPrecedingVmt, lengthWithTypeInfo(), replacementVmtWithTypeInfo.get()); } [[nodiscard]] std::size_t lengthWithTypeInfo() const noexcept { - return length + platform::lengthOfTypeInfoPrecedingVmt; + return replacementVmtWithTypeInfo.get_deleter().getNumberOfElements(); } std::uintptr_t* originalVmt; - std::size_t length; - std::uintptr_t* replacementVmtWithTypeInfo; + UniquePtr replacementVmtWithTypeInfo; };