virtualmemorymap.cpp

#include <Windows.h>
#include <algorithm>
#include <cstdint>
#include <vector>
#include <optional>

template<class Type>
class Address
{
public:
   constexpr Address() = default;

   constexpr explicit Address(uint32_t address) :
      mAddress(address)
   {
   }

   constexpr bool operator == (const Address &other) const
   {
      return mAddress == other.mAddress;
   }

   constexpr bool operator != (const Address &other) const
   {
      return mAddress != other.mAddress;
   }

   constexpr bool operator >= (const Address &other) const
   {
      return mAddress >= other.mAddress;
   }

   constexpr bool operator <= (const Address &other) const
   {
      return mAddress <= other.mAddress;
   }

   constexpr bool operator > (const Address &other) const
   {
      return mAddress > other.mAddress;
   }

   constexpr bool operator < (const Address &other) const
   {
      return mAddress < other.mAddress;
   }

   constexpr Address &operator += (ptrdiff_t value)
   {
      mAddress = static_cast<uint32_t>(mAddress + value);
      return *this;
   }

   constexpr Address &operator -= (ptrdiff_t value)
   {
      mAddress = static_cast<uint32_t>(mAddress - value);
      return *this;
   }

   constexpr Address operator + (ptrdiff_t value) const
   {
      return Address { static_cast<uint32_t>(mAddress + value) };
   }

   constexpr ptrdiff_t operator -(const Address &other) const
   {
      return mAddress - other.mAddress;
   }

   constexpr ptrdiff_t operator -(ptrdiff_t value) const
   {
      return mAddress - value;
   }

   constexpr uint32_t getAddress() const
   {
      return mAddress;
   }

private:
   uint32_t mAddress = 0;
};

class Physical;
class Virtual;

using PhysicalAddress = Address<Physical>;
using VirtualAddress = Address<Virtual>;

class MemoryMap
{
   enum PhysicalMemoryType
   {
      Invalid,
      MEM1,
      MEM2,
   };

   static constexpr PhysicalAddress MEM1BaseAddress = PhysicalAddress { 0 };
   static constexpr PhysicalAddress MEM1EndAddress = PhysicalAddress { 0x02000000 };

   static constexpr PhysicalAddress MEM2BaseAddress = PhysicalAddress { 0x10000000 };
   static constexpr PhysicalAddress MEM2EndAddress = PhysicalAddress { 0x90000000 };

   struct VirtualReservation
   {
      VirtualAddress start;
      uint64_t size;
   };

   struct VirtualMemoryMap
   {
      PhysicalMemoryType type;
      VirtualAddress virtualAddress;
      PhysicalAddress physicalAddress;
      uint32_t size;
   };

public:
   MemoryMap()
   {
   }

   ~MemoryMap()
   {
      free();
   }

   bool reserve()
   {
      // Reserve 32mb MEM1 and 2gb MEM2 physical memory
      mMem1 = CreateFileMappingW(INVALID_HANDLE_VALUE, NULL, SEC_RESERVE | PAGE_READWRITE, 0,
                                 static_cast<DWORD>(MEM1EndAddress - MEM1BaseAddress), NULL);
      mMem2 = CreateFileMappingW(INVALID_HANDLE_VALUE, NULL, SEC_RESERVE | PAGE_READWRITE, 0,
                                 static_cast<DWORD>(MEM2EndAddress - MEM2BaseAddress), NULL);

      // Reserve 4gb virtual address space
      mVirtualBase = 0;

      for (auto n = 32; n < 64; n++) {
         auto baseAddress = reinterpret_cast<uint8_t *>(1ull << n);
         auto addr = VirtualAlloc(baseAddress, 0x100000000ull, MEM_RESERVE, PAGE_NOACCESS);

         if (addr) {
            mVirtualBase = reinterpret_cast<uintptr_t>(addr);
            break;
         }
      }

      mReservedMemory.push_back({ VirtualAddress { 0 }, 0x100000000ull });
   }

   void free()
   {
      // Unmap all views
      while (mMappedMemory.size()) {
         auto &mapping = mMappedMemory[0];
         unmapMemory(mapping.virtualAddress, mapping.size);
      }

      // Release virtual memory
      if (mVirtualBase) {
         VirtualFree(reinterpret_cast<void *>(mVirtualBase), 0, MEM_RELEASE);
         mVirtualBase = 0;
      }

      // Close file mappings
      if (mMem1) {
         CloseHandle(mMem1);
         mMem1 = 0;
      }

      if (mMem2) {
         CloseHandle(mMem2);
         mMem2 = 0;
      }
   }

   PhysicalMemoryType getPhysicalMemoryType(PhysicalAddress physicalAddress)
   {
      if (physicalAddress >= MEM1BaseAddress && physicalAddress <= MEM1EndAddress) {
         return PhysicalMemoryType::MEM1;
      } else if (physicalAddress >= MEM2BaseAddress && physicalAddress <= MEM2EndAddress) {
         return PhysicalMemoryType::MEM2;
      }

      return PhysicalMemoryType::Invalid;
   }

   std::optional<PhysicalAddress> virtualToPhysicalAddress(VirtualAddress virtualAddress)
   {
      for (auto &mapping : mMappedMemory) {
         if (mapping.virtualAddress <= virtualAddress
             && mapping.virtualAddress + mapping.size >= virtualAddress) {
            ;
            return mapping.physicalAddress + (virtualAddress - mapping.virtualAddress);
         }
      }

      return {};
   }

   void *getVirtualPointer(VirtualAddress virtualAddress)
   {
      return reinterpret_cast<void *>(mVirtualBase + virtualAddress.getAddress());
   }

   bool mapMemory(VirtualAddress virtualAddress, PhysicalAddress physicalAddress, uint32_t size)
   {
      // First we have to break a hole in our reserved address space!
      auto foundHole = false;
      auto physicalMemoryType = getPhysicalMemoryType(physicalAddress);

      if (physicalMemoryType == PhysicalMemoryType::Invalid) {
         // Invalid physical memory address
         return false;
      }

      for (auto itr = mReservedMemory.begin(); itr != mReservedMemory.end(); ++itr) {
         auto &reservation = *itr;

         if (virtualAddress >= reservation.start && virtualAddress + size <= reservation.start + reservation.size) {
            auto result = VirtualFree(getVirtualPointer(reservation.start), 0, MEM_RELEASE);

            if (size == reservation.size) {
               // Consumed whole reservation
               mReservedMemory.erase(itr);
            } else if (virtualAddress == reservation.start) {
               // Consumed from start of reservation
               reservation.start = virtualAddress + size;
               reservation.size -= size;
               VirtualAlloc(getVirtualPointer(reservation.start), reservation.size, MEM_RESERVE, PAGE_NOACCESS);
            } else if (virtualAddress + size == reservation.start + reservation.size) {
               // Consumed from end of reservation
               reservation.size -= size;
               VirtualAlloc(getVirtualPointer(reservation.start), reservation.size, MEM_RESERVE, PAGE_NOACCESS);
            } else {
               // Consumed in middle of reservation
               auto newReservation = VirtualReservation {};
               newReservation.start = virtualAddress + size;
               newReservation.size = (reservation.start + reservation.size) - newReservation.start;
               VirtualAlloc(getVirtualPointer(newReservation.start), newReservation.size, MEM_RESERVE, PAGE_NOACCESS);

               reservation.size = virtualAddress - reservation.start;
               VirtualAlloc(getVirtualPointer(reservation.start), reservation.size, MEM_RESERVE, PAGE_NOACCESS);
               mReservedMemory.insert(itr + 1, newReservation);
            }

            foundHole = true;
            break;
         }
      }

      if (!foundHole) {
         // Virtual address space not free!
         return false;
      }

      // Map virtual address to physical memory
      void *view = nullptr;

      if (physicalMemoryType == PhysicalMemoryType::MEM1) {
         view = MapViewOfFileEx(mMem1, FILE_MAP_ALL_ACCESS, 0,
                                static_cast<DWORD>(physicalAddress - MEM1BaseAddress),
                                size,
                                getVirtualPointer(virtualAddress));
      } else {
         view = MapViewOfFileEx(mMem2, FILE_MAP_ALL_ACCESS, 0,
                                static_cast<DWORD>(physicalAddress - MEM2BaseAddress),
                                size,
                                getVirtualPointer(virtualAddress));
      }

      // Add to the memory map
      auto virtualMemoryMap = VirtualMemoryMap {};
      virtualMemoryMap.type = physicalMemoryType;
      virtualMemoryMap.virtualAddress = virtualAddress;
      virtualMemoryMap.physicalAddress = physicalAddress;
      virtualMemoryMap.size = size;

      mMappedMemory.insert(std::upper_bound(mMappedMemory.begin(),
                                            mMappedMemory.end(),
                                            virtualMemoryMap,
                                            [](const auto &m1, const auto &m2) {
         return m1.virtualAddress < m2.virtualAddress;
      }),
                           virtualMemoryMap);

      if (!view) {
         unmapMemory(virtualAddress, size);
         return false;
      }

      return true;
   }

   bool unmapMemory(VirtualAddress virtualAddress, uint32_t size)
   {
      // Unmap memory
      auto itr = std::find_if(mMappedMemory.begin(), mMappedMemory.end(),
                              [&](const auto &x) {
         return x.virtualAddress == virtualAddress && x.size == size;
      });

      if (itr == mMappedMemory.end()) {
         return false;
      }

      auto virtualMemoryMap = *itr;

      if (!UnmapViewOfFile(getVirtualPointer(virtualAddress))) {
         auto error = GetLastError();
         printf("error: %X", error);
      }

      mMappedMemory.erase(itr);

      // Check if we can combine with a previous reservation
      for (auto itr = mReservedMemory.begin(); itr != mReservedMemory.end(); ++itr) {
         auto &reservation = *itr;
         auto mergeItr = mReservedMemory.end();

         if (reservation.start == virtualAddress + size) {
            // Expand reservation backward
            VirtualFree(getVirtualPointer(reservation.start), reservation.size, MEM_RELEASE);
            reservation.start -= size;
            reservation.size += size;

            if (itr != mReservedMemory.begin()) {
               // Check if we can merge with the previous reservation
               auto prev = itr - 1;

               if (prev->start + prev->size == reservation.start) {
                  VirtualFree(getVirtualPointer(prev->start), prev->size, MEM_RELEASE);
                  reservation.start = prev->start;
                  reservation.size += prev->size;
                  mergeItr = prev;
               }
            }

            VirtualAlloc(getVirtualPointer(reservation.start), reservation.size, MEM_RESERVE, PAGE_NOACCESS);

            if (mergeItr != mReservedMemory.end()) {
               mReservedMemory.erase(mergeItr);
            }

            return true;
         } else if (reservation.start + reservation.size == virtualAddress) {
            // Expand reservation forward
            VirtualFree(getVirtualPointer(reservation.start), reservation.size, MEM_RELEASE);
            reservation.size += size;

            auto next = itr + 1;

            if (next != mReservedMemory.end()) {
               // Check if we can merge with the next reservation
               if (next->start == reservation.start + reservation.size) {
                  VirtualFree(getVirtualPointer(next->start), next->size, MEM_RELEASE);
                  reservation.size += next->size;
                  mergeItr = next;
               }
            }

            VirtualAlloc(getVirtualPointer(reservation.start), reservation.size, MEM_RESERVE, PAGE_NOACCESS);

            if (mergeItr != mReservedMemory.end()) {
               mReservedMemory.erase(mergeItr);
            }

            return true;
         }
      }

      // We cannot combine with a previous reservation so we must make a new one!
      auto reservation = VirtualReservation {};
      reservation.start = virtualAddress;
      reservation.size = size;
      VirtualAlloc(getVirtualPointer(reservation.start), reservation.size, MEM_RESERVE, PAGE_NOACCESS);

      mReservedMemory.insert(std::upper_bound(mReservedMemory.begin(),
                                              mReservedMemory.end(),
                                              reservation,
                                              [](const auto &m1, const auto &m2) {
         return m1.start < m2.start;
      }),
                             reservation);
      return true;
   }

private:
   HANDLE mMem1;
   HANDLE mMem2;
   uintptr_t mVirtualBase;
   std::vector<VirtualReservation> mReservedMemory;
   std::vector<VirtualMemoryMap> mMappedMemory;
};