rpcsx/kernel/orbis/src/dmem.cpp

#include "dmem.hpp"
#include "KernelAllocator.hpp"
#include "KernelContext.hpp"
#include "KernelObject.hpp"
#include "error.hpp"
#include "kernel/KernelObject.hpp"
#include "pmem.hpp"
#include "rx/AddressRange.hpp"
#include "rx/FileLock.hpp"
#include "rx/Serializer.hpp"
#include "rx/die.hpp"
#include "rx/print.hpp"
#include "utils/Logs.hpp"
#include "vmem.hpp"
#include <array>
#include <rx/format.hpp>
#include <string_view>
#include <thread/Process.hpp>
#include <utility>

struct DirectMemoryAllocation {
  static constexpr std::uint32_t kAllocatedBit = 1 << 31;
  static constexpr std::uint32_t kPooledBit = 1 << 30;
  std::uint32_t type = 0;

  struct Mapping {
    orbis::Process *process;
    rx::AddressRange vmRange;

    void serialize(rx::Serializer &s) const {
      s.serialize(process->pid);
      s.serialize(vmRange);
    }

    void deserialize(rx::Deserializer &d) {
      auto pid = d.deserialize<orbis::pid_t>();
      if (d.failure()) {
        return;
      }
      auto foundProcess = orbis::findProcessById(pid);
      if (foundProcess == nullptr) {
        d.setFailure();
        return;
      }

      process = foundProcess;
      d.deserialize(vmRange);
    }

    bool operator==(const Mapping &) const = default;
  };

  orbis::kvector<Mapping> mappings;

  [[nodiscard]] bool isPooled() const { return type & kPooledBit; }
  void markAsPooled() { type |= kPooledBit | kAllocatedBit; }

  void setMemoryType(orbis::MemoryType memoryType) {
    type = (std::to_underlying(memoryType) & ~(kAllocatedBit | kPooledBit)) |
           (type & kPooledBit) | kAllocatedBit;
  }

  [[nodiscard]] orbis::MemoryType getMemoryType() const {
    return static_cast<orbis::MemoryType>(type & ~(kAllocatedBit | kPooledBit));
  }

  [[nodiscard]] bool isAllocated() const { return (type & kAllocatedBit) != 0; }
  [[nodiscard]] bool isRelated(const DirectMemoryAllocation &other,
                               rx::AddressRange, rx::AddressRange) const {
    return type == other.type;
  }

  [[nodiscard]] DirectMemoryAllocation
  merge(const DirectMemoryAllocation &other, rx::AddressRange,
        rx::AddressRange) const {
    auto result = *this;
    result.mappings.insert(result.mappings.end(), other.mappings.begin(),
                           other.mappings.end());
    return result;
  }

  bool operator==(const DirectMemoryAllocation &) const = default;
};

struct DirectMemoryResourceState {
  std::uint64_t pmemOffset{};
  std::uint64_t dmemTotalSize{};
  std::uint64_t dmemReservedSize{};
};

struct DirectMemoryResource
    : kernel::AllocableResource<DirectMemoryAllocation, orbis::kallocator>,
      DirectMemoryResourceState {
  using BaseResource =
      kernel::AllocableResource<DirectMemoryAllocation, orbis::kallocator>;

  BaseResource systemResource;

  void create(std::size_t size) {
    auto [pmemRange, errc] = orbis::pmem::allocate(0, size, {}, 64 * 1024);

    rx::dieIf(errc != orbis::ErrorCode{},
              "failed to allocate direct memory: size {:x}, error {}", size,
              static_cast<int>(errc));

    dmemTotalSize = pmemRange.size();
    pmemOffset = pmemRange.beginAddress();
    auto dmemResourceRange =
        rx::AddressRange::fromBeginSize(0, pmemRange.size());

    if (auto errc = BaseResource::create(dmemResourceRange);
        errc != std::errc{}) {
      rx::die("failed to create direct memory resource {:x}, error {}",
              pmemRange.size(), errc);
    }
  }

  orbis::ErrorCode clear() {
    auto result = BaseResource::create(
        rx::AddressRange::fromBeginSize(0, dmemTotalSize + dmemReservedSize));

    if (result != std::errc{}) {
      return orbis::toErrorCode(result);
    }

    systemResource.destroy();
    dmemReservedSize = 0;
    return {};
  }

  void serialize(rx::Serializer &s) const {
    s.serialize(static_cast<const DirectMemoryResourceState &>(*this));
    BaseResource::serialize(s);
  }

  void deserialize(rx::Deserializer &d) {
    d.deserialize(static_cast<DirectMemoryResourceState &>(*this));
    BaseResource::deserialize(d);
  }

  std::pair<std::uint64_t, orbis::ErrorCode> reserveSystem(std::uint64_t size) {
    DirectMemoryAllocation alloc;
    alloc.setMemoryType(orbis::MemoryType::WbOnion);
    auto result = map(0, size, alloc, orbis::AllocationFlags::Stack,
                      orbis::dmem::kPageSize);

    if (result.errc != std::errc{}) {
      return {{}, orbis::toErrorCode(result.errc)};
    }

    allocations.unmap(result.range);
    dmemReservedSize += size;
    systemResource.allocations.map(result.range, {});
    return {result.range.beginAddress(), {}};
  }
};

static std::array g_dmemPools = {
    orbis::createGlobalObject<
        kernel::LockableKernelObject<DirectMemoryResource>>(),
    orbis::createGlobalObject<
        kernel::LockableKernelObject<DirectMemoryResource>>(),
    orbis::createGlobalObject<
        kernel::LockableKernelObject<DirectMemoryResource>>(),
};

orbis::ErrorCode orbis::dmem::initialize() {
  g_dmemPools[0]->create(0x180000000);
  g_dmemPools[1]->create(0x1000000);
  g_dmemPools[2]->create(0x1000000);

  g_dmemPools[0]->reserveSystem(0x60000000);

  return {};
}

orbis::ErrorCode orbis::dmem::clear(unsigned dmemIndex) {
  if (dmemIndex >= std::size(g_dmemPools)) {
    return ErrorCode::INVAL;
  }

  {
    auto dmem = g_dmemPools[dmemIndex];
    std::lock_guard lock(*dmem);
    auto result = dmem->clear();
    if (result != orbis::ErrorCode{}) {
      return result;
    }
  }

  return {};
}

static void dmemDump(unsigned dmemIndex, std::string_view message = {}) {
  auto dmem = g_dmemPools[dmemIndex];

  rx::ScopedFileLock lock(stderr);

  rx::println(stderr, "dmem0 {}", message);
  for (auto alloc : *dmem) {
    rx::println(stderr, "  {:012x}-{:012x}: {}{} {}", alloc.beginAddress(),
                alloc.endAddress(), "_A"[alloc->isAllocated()],
                "_P"[alloc->isPooled()], alloc->getMemoryType());
  }
}

std::pair<std::uint64_t, orbis::ErrorCode>
orbis::dmem::allocate(unsigned dmemIndex, rx::AddressRange searchRange,
                      std::uint64_t len, MemoryType memoryType,
                      std::uint64_t alignment, bool pooled) {
  if (dmemIndex >= std::size(g_dmemPools)) {
    return {{}, ErrorCode::INVAL};
  }

  if (searchRange.endAddress() != 0 &&
      (!searchRange.isValid() || searchRange.size() < len)) {
    ORBIS_LOG_ERROR(__FUNCTION__, "invalid range", searchRange.beginAddress(),
                    searchRange.endAddress(), len);
    return {{}, orbis::ErrorCode::INVAL};
  }

  auto dmem = g_dmemPools[dmemIndex];
  std::lock_guard lock(*dmem);

  alignment = alignment == 0 ? kPageSize : alignment;
  len = rx::alignUp(len, dmem::kPageSize);

  if (searchRange.endAddress() == 0) {
    searchRange = rx::AddressRange::fromBeginEnd(searchRange.beginAddress(),
                                                 dmem->dmemTotalSize -
                                                     dmem->dmemReservedSize);
  }

  if (!searchRange.isValid() || searchRange.size() < len) {
    ORBIS_LOG_ERROR(__FUNCTION__, "invalid range", searchRange.beginAddress(),
                    searchRange.endAddress(), len);
    return {{}, orbis::ErrorCode::INVAL};
  }

  if (len == 0) {
    ORBIS_LOG_ERROR(__FUNCTION__, "len is 0", searchRange.beginAddress(),
                    searchRange.endAddress(), len);

    return {{}, ErrorCode::INVAL};
  }

  if (searchRange.endAddress() > dmem->dmemTotalSize - dmem->dmemReservedSize) {
    ORBIS_LOG_ERROR(__FUNCTION__, "out of direct memory size",
                    searchRange.beginAddress(), searchRange.endAddress(),
                    dmem->dmemTotalSize);
    return {{}, ErrorCode::INVAL};
  }

  DirectMemoryAllocation allocation;
  allocation.setMemoryType(memoryType);
  if (pooled) {
    allocation.markAsPooled();
  }

  auto allocResult = dmem->map(searchRange.beginAddress(), len, allocation,
                               AllocationFlags::Dry, alignment);
  if (allocResult.errc != std::errc{}) {
    return {{}, ErrorCode::AGAIN};
  }

  auto result = allocResult.range.beginAddress();

  auto commitResult =
      dmem->map(result, len, allocation, AllocationFlags::Fixed, alignment);

  rx::dieIf(commitResult.errc != std::errc{},
            "dmem: failed to commit memory, error {}", commitResult.errc);

  dmemDump(dmemIndex,
           rx::format("allocated {:x}-{:x}", allocResult.range.beginAddress(),
                      allocResult.range.endAddress()));

  return {result, {}};
}

std::pair<std::uint64_t, orbis::ErrorCode>
orbis::dmem::allocateSystem(unsigned dmemIndex, std::uint64_t len,
                            MemoryType memoryType, std::uint64_t alignment) {
  if (dmemIndex >= std::size(g_dmemPools)) {
    return {{}, ErrorCode::INVAL};
  }

  auto dmem = g_dmemPools[dmemIndex];
  std::lock_guard lock(*dmem);

  alignment = alignment == 0 ? kPageSize : alignment;
  len = rx::alignUp(len, dmem::kPageSize);

  DirectMemoryAllocation allocation;
  allocation.setMemoryType(memoryType);

  auto allocResult = dmem->systemResource.map(0, len, allocation,
                                              AllocationFlags::Dry, alignment);
  if (allocResult.errc != std::errc{}) {
    return {{}, ErrorCode::AGAIN};
  }

  auto result = allocResult.range.beginAddress();

  auto commitResult =
      dmem->map(result, len, allocation, AllocationFlags::Fixed, alignment);

  rx::dieIf(commitResult.errc != std::errc{},
            "dmem: failed to commit main memory, error {}", commitResult.errc);

  dmemDump(dmemIndex, rx::format("allocated main {:x}-{:x}",
                                 allocResult.range.beginAddress(),
                                 allocResult.range.endAddress()));

  return {allocResult.range.beginAddress() | 0x4000000000, {}};
}

orbis::ErrorCode orbis::dmem::release(unsigned dmemIndex,
                                      rx::AddressRange range, bool pooled) {
  if (dmemIndex >= std::size(g_dmemPools)) {
    return ErrorCode::INVAL;
  }

  if (range.beginAddress() % kPageSize || range.size() % vmem::kPageSize ||
      !range.isValid()) {
    return ErrorCode::INVAL;
  }

  auto dmem = g_dmemPools[dmemIndex];
  std::lock_guard lock(*dmem);

  constexpr auto razorGpuMemory =
      rx::AddressRange::fromBeginSize(0x3000000000, 0x20000000);

  if (dmemIndex == 0 && razorGpuMemory.contains(range.beginAddress())) {
    return ErrorCode::OPNOTSUPP;
  }

  auto it = dmem->query(range.beginAddress());

  if (it == dmem->end() || !it->isAllocated()) {
    return ErrorCode::NOENT;
  }

  if (it->isPooled() && !pooled) {
    return ErrorCode::NOENT;
  }

  for (auto mapping : it->mappings) {
    mapping.process->invoke(
        [=] { vmem::unmap(mapping.process, mapping.vmRange); });
  }

  it->mappings.clear();

  DirectMemoryAllocation allocation{};
  auto result = dmem->map(range.beginAddress(), range.size(), allocation,
                          AllocationFlags::Fixed, vmem::kPageSize);

  if (result.errc != std::errc{}) {
    return toErrorCode(result.errc);
  }

  dmemDump(dmemIndex, rx::format("released {:x}-{:x}", range.beginAddress(),
                                 range.endAddress()));
  return {};
}

std::pair<std::uint64_t, orbis::ErrorCode>
orbis::dmem::reserveSystem(unsigned dmemIndex, std::uint64_t size) {
  if (dmemIndex >= std::size(g_dmemPools)) {
    return {{}, ErrorCode::INVAL};
  }

  auto dmem = g_dmemPools[dmemIndex];
  std::lock_guard lock(*dmem);

  if (size == 0 || size % vmem::kPageSize) {
    return {{}, ErrorCode::INVAL};
  }

  if (dmem->dmemReservedSize + size > dmem->dmemTotalSize) {
    return {{}, ErrorCode::NOMEM};
  }

  return dmem->reserveSystem(size);
}

std::pair<std::uint64_t, orbis::ErrorCode>
orbis::dmem::reservePooled(unsigned dmemIndex, rx::AddressRange searchRange,
                           std::uint64_t size, std::uint64_t alignment) {
  if (dmemIndex >= std::size(g_dmemPools)) {
    return {{}, ErrorCode::INVAL};
  }

  alignment = alignment == 0 ? kPageSize : alignment;

  if (alignment % kPageSize) {
    return {{}, ErrorCode::INVAL};
  }

  auto dmem = g_dmemPools[dmemIndex];

  if (searchRange.endAddress() > dmem->dmemTotalSize) {
    ORBIS_LOG_ERROR(__FUNCTION__, "out of direct memory size",
                    searchRange.endAddress(), dmem->dmemTotalSize);
    return {{}, ErrorCode::INVAL};
  }

  if (!searchRange.isValid() &&
      searchRange.beginAddress() != searchRange.endAddress()) {
    ORBIS_LOG_ERROR(__FUNCTION__, "invalid range", searchRange.beginAddress(),
                    searchRange.endAddress());
    return {{}, ErrorCode::INVAL};
  }

  std::lock_guard lock(*dmem);

  if (searchRange.beginAddress() == searchRange.endAddress()) {
    searchRange = rx::AddressRange::fromBeginEnd(searchRange.beginAddress(),
                                                 dmem->dmemTotalSize -
                                                     dmem->dmemReservedSize);

    if (!searchRange.isValid()) {
      ORBIS_LOG_ERROR(__FUNCTION__, "invalid range", searchRange.beginAddress(),
                      searchRange.endAddress());
      return {{}, ErrorCode::INVAL};
    }
  }

  if (searchRange.endAddress() > dmem->dmemTotalSize - dmem->dmemReservedSize) {
    ORBIS_LOG_ERROR(__FUNCTION__, "out of direct memory size",
                    searchRange.endAddress(), dmem->dmemTotalSize,
                    dmem->dmemReservedSize);
    return {{}, ErrorCode::INVAL};
  }

  auto it = dmem->lowerBound(searchRange.beginAddress());
  rx::AddressRange result;

  while (it != dmem->end() && it.beginAddress() < searchRange.endAddress()) {
    if (!it->isAllocated()) {
      auto viewRange = searchRange.intersection(it.range());

      viewRange = rx::AddressRange::fromBeginEnd(
          rx::alignUp(viewRange.beginAddress(), alignment),
          viewRange.endAddress());

      if (viewRange.isValid() && viewRange.size() >= size) {
        result =
            rx::AddressRange::fromBeginSize(viewRange.beginAddress(), size);
        break;
      }
    }

    ++it;
  }

  if (!result.isValid()) {
    return {{}, ErrorCode::NOMEM};
  }

  DirectMemoryAllocation allocation;
  allocation.markAsPooled();

  auto commitResult = dmem->map(result.beginAddress(), result.size(),
                                allocation, AllocationFlags::Fixed, alignment);

  if (commitResult.errc != std::errc{}) {
    return {{}, toErrorCode(commitResult.errc)};
  }

  return {result.beginAddress(), {}};
}

orbis::ErrorCode orbis::dmem::setType(unsigned dmemIndex,
                                      rx::AddressRange range, MemoryType type,
                                      rx::EnumBitSet<QueryFlags> flags) {
  if (dmemIndex >= std::size(g_dmemPools)) {
    return ErrorCode::INVAL;
  }

  bool expectedPooled = (flags & QueryFlags::Pooled) == QueryFlags::Pooled;

  auto dmem = g_dmemPools[dmemIndex];
  std::lock_guard lock(*dmem);

  auto it = dmem->query(range.beginAddress());

  if (it == dmem->end() || !it->isAllocated() ||
      it->isPooled() != expectedPooled) {
    return ErrorCode::ACCES;
  }

  if (it->getMemoryType() == type && it.range().contains(range)) {
    return {};
  }

  DirectMemoryAllocation allocation;
  allocation.setMemoryType(type);
  auto [_it, errc, _range] =
      dmem->map(range.beginAddress(), range.size(), allocation,
                AllocationFlags::Fixed, vmem::kPageSize);

  return toErrorCode(errc);
}

std::optional<orbis::dmem::QueryResult>
orbis::dmem::query(unsigned dmemIndex, std::uint64_t dmemOffset,
                   rx::EnumBitSet<QueryFlags> flags) {
  if (dmemIndex >= std::size(g_dmemPools)) {
    return {};
  }

  auto dmem = g_dmemPools[dmemIndex];
  std::lock_guard lock(*dmem);

  auto it = dmem->lowerBound(dmemOffset);

  if (flags & QueryFlags::LowerBound) {
    while (it != dmem->end() && !it->isAllocated()) {
      ++it;
    }

    if (it == dmem->end()) {
      return {};
    }
  } else if (it == dmem->end() || !it->isAllocated()) {
    return {};
  }

  if (!(flags & QueryFlags::Pooled)) {
    if (it->isPooled()) {
      return {};
    }
  }

  return QueryResult{.range = it.range(), .memoryType = it->getMemoryType()};
}

std::uint64_t orbis::dmem::getSize(unsigned dmemIndex) {
  if (dmemIndex >= std::size(g_dmemPools)) {
    return 0;
  }

  auto dmem = g_dmemPools[dmemIndex];
  std::lock_guard lock(*dmem);

  return dmem->dmemTotalSize - dmem->dmemReservedSize;
}

std::pair<rx::AddressRange, orbis::ErrorCode>
orbis::dmem::getAvailSize(unsigned dmemIndex, rx::AddressRange searchRange,
                          std::uint64_t alignment) {
  if (dmemIndex >= std::size(g_dmemPools)) {
    return {{}, ErrorCode::INVAL};
  }

  alignment = alignment == 0 ? kPageSize : alignment;

  if (alignment % vmem::kPageSize) {
    return {{}, ErrorCode::INVAL};
  }

  auto dmem = g_dmemPools[dmemIndex];

  if (searchRange.endAddress() > dmem->dmemTotalSize) {
    ORBIS_LOG_ERROR(__FUNCTION__, "out of direct memory size",
                    searchRange.endAddress(), dmem->dmemTotalSize);
    // return {{}, orbis::ErrorCode::INVAL};
    searchRange = rx::AddressRange::fromBeginEnd(searchRange.beginAddress(),
                                                 dmem->dmemTotalSize);
  }

  if (!searchRange.isValid() &&
      searchRange.beginAddress() != searchRange.endAddress()) {
    ORBIS_LOG_ERROR(__FUNCTION__, "invalid range", searchRange.beginAddress(),
                    searchRange.endAddress());
    return {{}, orbis::ErrorCode::INVAL};
  }

  std::lock_guard lock(*dmem);

  if (searchRange.beginAddress() == searchRange.endAddress()) {
    searchRange = rx::AddressRange::fromBeginEnd(searchRange.beginAddress(),
                                                 dmem->dmemTotalSize -
                                                     dmem->dmemReservedSize);

    if (!searchRange.isValid()) {
      ORBIS_LOG_ERROR(__FUNCTION__, "invalid range", searchRange.beginAddress(),
                      searchRange.endAddress());
      return {{}, orbis::ErrorCode::INVAL};
    }
  }

  if (searchRange.endAddress() > dmem->dmemTotalSize - dmem->dmemReservedSize) {
    ORBIS_LOG_ERROR(__FUNCTION__, "out of direct memory size",
                    searchRange.endAddress(), dmem->dmemTotalSize,
                    dmem->dmemReservedSize);
    searchRange = rx::AddressRange::fromBeginEnd(searchRange.beginAddress(),
                                                 dmem->dmemTotalSize -
                                                     dmem->dmemReservedSize);

    // return {{}, orbis::ErrorCode::INVAL};
  }

  auto it = dmem->lowerBound(searchRange.beginAddress());
  rx::AddressRange result = rx::AddressRange::fromBeginEnd(
      searchRange.beginAddress(), searchRange.beginAddress());

  while (it != dmem->end() && it.beginAddress() < searchRange.endAddress()) {
    if (!it->isAllocated()) {
      auto viewRange = searchRange.intersection(it.range());

      viewRange = rx::AddressRange::fromBeginEnd(
          rx::alignUp(viewRange.beginAddress(), alignment),
          viewRange.endAddress());

      if (viewRange.isValid() &&
          (!result.isValid() || viewRange.size() > result.size())) {
        result = viewRange;
      }
    }

    ++it;
  }

  return {result, {}};
}

orbis::ErrorCode orbis::dmem::map(orbis::Process *process, unsigned dmemIndex,
                                  rx::AddressRange range, std::uint64_t offset,
                                  rx::EnumBitSet<vmem::Protection> protection) {
  if (dmemIndex >= std::size(g_dmemPools)) {
    return ErrorCode::INVAL;
  }

  auto dmem = g_dmemPools[dmemIndex];
  std::lock_guard lock(*dmem);

  if (offset < 0 || offset + range.size() < offset ||
      offset + range.size() > dmem->dmemTotalSize) {
    return orbis::ErrorCode::INVAL;
  }

  if (offset + range.size() > dmem->dmemTotalSize - dmem->dmemReservedSize) {
    return orbis::ErrorCode::ACCES;
  }

  auto allocationInfoIt = dmem->query(offset);

  if (allocationInfoIt == dmem->end() || !allocationInfoIt->isAllocated()) {
    if (allocationInfoIt != dmem->end()) {
      dmemDump(
          dmemIndex,
          rx::format("map unallocated {:x}-{:x}, requested range {:x}-{:x}",
                     allocationInfoIt.beginAddress(),
                     allocationInfoIt.endAddress(), range.beginAddress(),
                     range.endAddress()));
    } else {
      dmemDump(dmemIndex, rx::format("map out of memory {:x}-{:x}",
                                     range.beginAddress(), range.endAddress()));
    }
    return orbis::ErrorCode::ACCES;
  }

  if (allocationInfoIt->isPooled()) {
    return orbis::ErrorCode::ACCES;
  }

  if (!vmem::validateMemoryType(allocationInfoIt->getMemoryType(),
                                protection)) {
    return ErrorCode::ACCES;
  }

  // if (!allocationInfoIt->mappings.empty() && !process->allowDmemAliasing) {
  //   return ErrorCode::INVAL;
  // }

  auto directRange = rx::AddressRange::fromBeginSize(offset, range.size())
                         .intersection(allocationInfoIt.range());

  if (range.size() > directRange.size()) {
    return orbis::ErrorCode::INVAL;
  }

  auto physicalRange =
      rx::AddressRange::fromBeginSize(dmem->pmemOffset + offset, range.size());

  auto result = process->invoke([=] {
    return orbis::pmem::map(range.beginAddress(), physicalRange,
                            vmem::toCpuProtection(protection));
  });

  if (result == ErrorCode{}) {
    allocationInfoIt->mappings.push_back({
        .process = process,
        .vmRange = range,
    });
  }

  return result;
}

orbis::ErrorCode orbis::dmem::notifyUnmap(orbis::Process *process,
                                          unsigned dmemIndex,
                                          std::uint64_t offset,
                                          rx::AddressRange range) {
  if (dmemIndex >= std::size(g_dmemPools)) {
    return ErrorCode::INVAL;
  }

  auto dmem = g_dmemPools[dmemIndex];
  std::lock_guard lock(*dmem);

  auto it = dmem->query(offset);
  if (it == dmem->end()) {
    return ErrorCode::INVAL;
  }

  for (auto mapIt = it->mappings.begin(); mapIt != it->mappings.end();) {
    if (mapIt->process == process && mapIt->vmRange.intersects(range)) {
      auto blockRange = range.intersection(mapIt->vmRange);
      if (mapIt->vmRange == blockRange) {
        mapIt = it->mappings.erase(mapIt);
        break;
      }

      if (mapIt->vmRange.beginAddress() == blockRange.beginAddress()) {
        mapIt->vmRange = rx::AddressRange::fromBeginEnd(
            blockRange.endAddress(), mapIt->vmRange.endAddress());
        break;
      }

      if (mapIt->vmRange.endAddress() == blockRange.endAddress()) {
        mapIt->vmRange = rx::AddressRange::fromBeginEnd(
            mapIt->vmRange.beginAddress(), blockRange.beginAddress());
        break;
      }

      auto leftAllocation = rx::AddressRange::fromBeginEnd(
          mapIt->vmRange.beginAddress(), blockRange.beginAddress());

      auto rightAllocation = rx::AddressRange::fromBeginEnd(
          blockRange.endAddress(), mapIt->vmRange.endAddress());

      mapIt->vmRange = leftAllocation;
      it->mappings.push_back({.process = process, .vmRange = rightAllocation});
      break;
    }

    ++mapIt;
  }

  return {};
}

orbis::ErrorCode orbis::dmem::protect(orbis::Process *process,
                                      unsigned dmemIndex,
                                      rx::AddressRange range,
                                      rx::EnumBitSet<vmem::Protection> prot) {
  auto dmem = g_dmemPools[dmemIndex];
  std::lock_guard lock(*dmem);

  auto it = dmem->query(range.beginAddress());
  if (it == dmem->end()) {
    return ErrorCode::INVAL;
  }

  if (!it.range().contains(range)) {
    return ErrorCode::INVAL;
  }

  for (auto mapping : it->mappings) {
    if (process == nullptr || process == mapping.process) {
      vmem::protect(mapping.process, mapping.vmRange, prot);
    }
  }

  return {};
}

std::pair<std::uint64_t, orbis::ErrorCode>
orbis::dmem::getPmemOffset(unsigned dmemIndex, std::uint64_t dmemOffset) {
  if (dmemIndex > std::size(g_dmemPools)) {
    return {{}, ErrorCode::INVAL};
  }

  auto dmem = g_dmemPools[dmemIndex];
  std::lock_guard lock(*dmem);

  if (dmemOffset >= dmem->dmemTotalSize) {
    return {{}, orbis::ErrorCode::INVAL};
  }

  if (dmemOffset >= dmem->dmemTotalSize - dmem->dmemReservedSize) {
    return {{}, orbis::ErrorCode::ACCES};
  }

  auto allocationInfoIt = dmem->query(dmemOffset);

  if (allocationInfoIt == dmem->end() || !allocationInfoIt->isAllocated()) {
    return {{}, orbis::ErrorCode::ACCES};
  }

  return {dmem->pmemOffset + dmemOffset, {}};
}