WUMSLoader/wumsloader/src/utils/HeapMemoryFixedSize.cpp

#include "HeapMemoryFixedSize.h"

#include "logger.h"
#include "utils.h"

#include <cassert>

HeapMemoryFixedSizePool::MemorySegmentInfo::MemorySegmentInfo(void *data, const size_t size) : mData(data),
                                                                                               mSize(size) {}

void *HeapMemoryFixedSizePool::MemorySegmentInfo::data() const {
    return mData;
}

size_t HeapMemoryFixedSizePool::MemorySegmentInfo::size() const {
    return mSize;
}

HeapMemoryFixedSizePool::HeapMemoryFixedSizePool() = default;

HeapMemoryFixedSizePool::HeapMemoryFixedSizePool(const AllocFunc &allocFunc, FreeFunc freeFunc, const std::initializer_list<size_t> segmentSizes)
    : HeapMemoryFixedSizePool(allocFunc, std::move(freeFunc), std::span(segmentSizes.begin(), segmentSizes.size())) {
}

HeapMemoryFixedSizePool::HeapMemoryFixedSizePool(const AllocFunc &allocFunc, FreeFunc freeFunc, const std::span<const std::size_t> segmentSizes) {
    assert(!segmentSizes.empty());
    assert(allocFunc != nullptr);
    assert(freeFunc != nullptr);

    size_t totalSize = 0;
    for (const auto size : segmentSizes) {
        totalSize += size + 0x40; // add 0x40 bytes overhead for each entry to ensure padding to 0x40
    }

    // Call the injected allocator!
    void *raw_mem = allocFunc(totalSize, 0x40);

    if (!raw_mem) {
        return;
    }

    mData = std::unique_ptr<uint8_t[], CustomDeleter>(
            static_cast<uint8_t *>(raw_mem),
            CustomDeleter{std::move(freeFunc)});

    mTotalSize = (mData ? totalSize : 0);
    if (mData) {
        auto address = reinterpret_cast<uint32_t>(mData.get());
        for (const auto size : segmentSizes) {
            address = ROUNDUP(static_cast<int>(address), 0x40);
            assert(address >= reinterpret_cast<uint32_t>(mData.get()) && address < reinterpret_cast<uint32_t>(mData.get()) + totalSize);
            mSegmentInfos.emplace_back(reinterpret_cast<void *>(address), size);
            address += size;
        }
    }
}

HeapMemoryFixedSizePool::HeapMemoryFixedSizePool(HeapMemoryFixedSizePool &&other) noexcept
    : mData(std::move(other.mData)), mTotalSize(other.mTotalSize), mSegmentInfos(std::move(other.mSegmentInfos)) {
    other.mTotalSize = 0;
}

HeapMemoryFixedSizePool &HeapMemoryFixedSizePool::operator=(HeapMemoryFixedSizePool &&other) noexcept {
    if (this != &other) {
        mData            = std::move(other.mData);
        mTotalSize       = other.mTotalSize;
        mSegmentInfos    = std::move(other.mSegmentInfos);
        other.mTotalSize = 0;
        other.mSegmentInfos.clear();
    }
    return *this;
}

uint32_t HeapMemoryFixedSizePool::numberOfSegments() const {
    return mSegmentInfos.size();
}

HeapMemoryFixedSizePool::operator bool() const {
    return mData != nullptr && mTotalSize > 0;
}

HeapMemoryFixedSizePool::MemorySegmentInfo HeapMemoryFixedSizePool::operator[](const size_t idx) const {
    if (idx >= mSegmentInfos.size()) {
        DEBUG_FUNCTION_LINE_ERR("Out of bounce access (tried to access index %d; size is %d", idx, mSegmentInfos.size());
        return {nullptr, 0};
    }
    return mSegmentInfos[idx];
}

std::span<const uint8_t> HeapMemoryFixedSizePool::dataView() const {
    return {mData.get(), mTotalSize};
}