SIMDCache.hpp

#pragma once
#include <immintrin.h>
#include <vector>
#include <cstdint>
#include <memory>
#include <Windows.h>

class SIMDCache {
public:
    static constexpr size_t CACHE_LINE_SIZE = 64;
    static constexpr size_t AVX2_ALIGNMENT = 32;

    SIMDCache() : cacheSize(0), cacheData(nullptr) {}

    bool Initialize(size_t size) {
        cacheSize = AlignUp(size, AVX2_ALIGNMENT);

        cacheData = static_cast<uint8_t*>(_aligned_malloc(cacheSize, AVX2_ALIGNMENT));

        if (!cacheData) {
            cacheSize = 0;
            return false;
        }

        return true;
    }

    ~SIMDCache() {
        if (cacheData) {
            _aligned_free(cacheData);
            cacheData = nullptr;
        }
        cacheSize = 0;
    }

    bool ReadMemory(void* address, size_t size, HANDLE process = nullptr) {
        if (!cacheData || size > cacheSize) return false;

        if (!process) process = GetCurrentProcess();

        SIZE_T bytesRead;
        if (!ReadProcessMemory(process, address, cacheData, size, &bytesRead)) {
            return false;
        }

        currentSize = bytesRead;
        return true;
    }

    // Write cache back to memory
    bool WriteMemory(void* address, size_t size, HANDLE process = nullptr) {
        if (!cacheData || size > currentSize) return false;

        if (!process) process = GetCurrentProcess();

        SIZE_T bytesWritten;
        if (!WriteProcessMemory(process, address, cacheData, size, &bytesWritten)) {
            return false;
        }

        return bytesWritten == size;
    }

    std::vector<uintptr_t> ScanPattern(const std::vector<uint8_t>& pattern,
        const std::vector<uint8_t>& mask) {
        std::vector<uintptr_t> results;

        if (!cacheData || pattern.empty() || pattern.size() != mask.size()) {
            return results;
        }

        const size_t patternSize = pattern.size();
        const size_t scanSize = currentSize - patternSize;

        if (patternSize >= 32) {
            ScanAVX2(pattern, mask, results, scanSize, patternSize);
        }
        else {
            ScanScalar(pattern, mask, results, scanSize, patternSize);
        }

        return results;
    }

    uint8_t* GetData() { return cacheData; }
    size_t GetSize() const { return currentSize; }

private:
    uint8_t* cacheData;
    size_t cacheSize;
    size_t currentSize = 0;

    static size_t AlignUp(size_t value, size_t alignment) {
        return (value + alignment - 1) & ~(alignment - 1);
    }

    void ScanAVX2(const std::vector<uint8_t>& pattern,
        const std::vector<uint8_t>& mask,
        std::vector<uintptr_t>& results,
        size_t scanSize,
        size_t patternSize) {

        const size_t avxPatternSize = AlignUp(patternSize, 32);

        __m256i patternReg = _mm256_loadu_si256(
            reinterpret_cast<const __m256i*>(pattern.data()));
        __m256i maskReg = _mm256_loadu_si256(
            reinterpret_cast<const __m256i*>(mask.data()));

        for (size_t i = 0; i <= scanSize; i++) {
            __m256i dataReg = _mm256_loadu_si256(
                reinterpret_cast<const __m256i*>(cacheData + i));

            __m256i maskedData = _mm256_and_si256(dataReg, maskReg);
            __m256i cmpResult = _mm256_cmpeq_epi8(maskedData, patternReg);

            int maskResult = _mm256_movemask_epi8(cmpResult);

            if (maskResult == -1) {
                if (patternSize <= 32 ||
                    CheckRemainingBytes(i, pattern, mask, patternSize)) {
                    results.push_back(i);
                }
            }
        }
    }

    void ScanScalar(const std::vector<uint8_t>& pattern,
        const std::vector<uint8_t>& mask,
        std::vector<uintptr_t>& results,
        size_t scanSize,
        size_t patternSize) {

        for (size_t i = 0; i <= scanSize; i++) {
            bool match = true;

            for (size_t j = 0; j < patternSize; j++) {
                if (mask[j] != 0 && cacheData[i + j] != pattern[j]) {
                    match = false;
                    break;
                }
            }

            if (match) {
                results.push_back(i);
            }
        }
    }

    bool CheckRemainingBytes(size_t offset,
        const std::vector<uint8_t>& pattern,
        const std::vector<uint8_t>& mask,
        size_t patternSize) {
        size_t remaining = patternSize - 32;

        for (size_t i = 0; i < remaining; i++) {
            size_t idx = 32 + i;
            if (mask[idx] != 0 &&
                cacheData[offset + idx] != pattern[idx]) {
                return false;
            }
        }

        return true;
    }
};