4 files changed, 227 insertions, 93 deletions

diff --git a/src/core/hle/service/nvdrv/devices/nvhost_as_gpu.cpp b/src/core/hle/service/nvdrv/devices/nvhost_as_gpu.cpp
index 195421cc0..d4ba88147 100644
--- a/src/core/hle/service/nvdrv/devices/nvhost_as_gpu.cpp
+++ b/src/core/hle/service/nvdrv/devices/nvhost_as_gpu.cpp
@@ -16,11 +16,12 @@
 #include "video_core/renderer_base.h"
 
 namespace Service::Nvidia::Devices {
+
 namespace NvErrCodes {
-enum {
-    InvalidNmapHandle = -22,
-};
-}
+constexpr u32 Success{};
+constexpr u32 OutOfMemory{static_cast<u32>(-12)};
+constexpr u32 InvalidInput{static_cast<u32>(-22)};
+} // namespace NvErrCodes
 
 nvhost_as_gpu::nvhost_as_gpu(Core::System& system, std::shared_ptr<nvmap> nvmap_dev)
     : nvdevice(system), nvmap_dev(std::move(nvmap_dev)) {}
@@ -49,8 +50,9 @@ u32 nvhost_as_gpu::ioctl(Ioctl command, const std::vector<u8>& input, const std:
         break;
     }
 
-    if (static_cast<IoctlCommand>(command.cmd.Value()) == IoctlCommand::IocRemapCommand)
+    if (static_cast<IoctlCommand>(command.cmd.Value()) == IoctlCommand::IocRemapCommand) {
         return Remap(input, output);
+    }
 
     UNIMPLEMENTED_MSG("Unimplemented ioctl command");
     return 0;
@@ -59,6 +61,7 @@ u32 nvhost_as_gpu::ioctl(Ioctl command, const std::vector<u8>& input, const std:
 u32 nvhost_as_gpu::InitalizeEx(const std::vector<u8>& input, std::vector<u8>& output) {
     IoctlInitalizeEx params{};
     std::memcpy(&params, input.data(), input.size());
+
     LOG_WARNING(Service_NVDRV, "(STUBBED) called, big_page_size=0x{:X}", params.big_page_size);
 
     return 0;
@@ -67,53 +70,61 @@ u32 nvhost_as_gpu::InitalizeEx(const std::vector<u8>& input, std::vector<u8>& ou
 u32 nvhost_as_gpu::AllocateSpace(const std::vector<u8>& input, std::vector<u8>& output) {
     IoctlAllocSpace params{};
     std::memcpy(&params, input.data(), input.size());
+
     LOG_DEBUG(Service_NVDRV, "called, pages={:X}, page_size={:X}, flags={:X}", params.pages,
               params.page_size, params.flags);
 
-    auto& gpu = system.GPU();
-    const u64 size{static_cast<u64>(params.pages) * static_cast<u64>(params.page_size)};
-    if (params.flags & 1) {
-        params.offset = gpu.MemoryManager().AllocateSpace(params.offset, size, 1);
+    const auto size{static_cast<u64>(params.pages) * static_cast<u64>(params.page_size)};
+    if ((params.flags & AddressSpaceFlags::FixedOffset) != AddressSpaceFlags::None) {
+        params.offset = *system.GPU().MemoryManager().AllocateFixed(params.offset, size);
     } else {
-        params.offset = gpu.MemoryManager().AllocateSpace(size, params.align);
+        params.offset = system.GPU().MemoryManager().Allocate(size, params.align);
+    }
+
+    auto result{NvErrCodes::Success};
+    if (!params.offset) {
+        LOG_CRITICAL(Service_NVDRV, "allocation failed for size {}", size);
+        result = NvErrCodes::OutOfMemory;
     }
 
     std::memcpy(output.data(), &params, output.size());
-    return 0;
+    return result;
 }
 
 u32 nvhost_as_gpu::Remap(const std::vector<u8>& input, std::vector<u8>& output) {
-    std::size_t num_entries = input.size() / sizeof(IoctlRemapEntry);
+    const auto num_entries = input.size() / sizeof(IoctlRemapEntry);
 
-    LOG_WARNING(Service_NVDRV, "(STUBBED) called, num_entries=0x{:X}", num_entries);
+    LOG_DEBUG(Service_NVDRV, "called, num_entries=0x{:X}", num_entries);
 
+    auto result{NvErrCodes::Success};
     std::vector<IoctlRemapEntry> entries(num_entries);
     std::memcpy(entries.data(), input.data(), input.size());
 
-    auto& gpu = system.GPU();
     for (const auto& entry : entries) {
-        LOG_WARNING(Service_NVDRV, "remap entry, offset=0x{:X} handle=0x{:X} pages=0x{:X}",
-                    entry.offset, entry.nvmap_handle, entry.pages);
-        GPUVAddr offset = static_cast<GPUVAddr>(entry.offset) << 0x10;
-        auto object = nvmap_dev->GetObject(entry.nvmap_handle);
+        LOG_DEBUG(Service_NVDRV, "remap entry, offset=0x{:X} handle=0x{:X} pages=0x{:X}",
+                  entry.offset, entry.nvmap_handle, entry.pages);
+
+        const auto object{nvmap_dev->GetObject(entry.nvmap_handle)};
         if (!object) {
-            LOG_CRITICAL(Service_NVDRV, "nvmap {} is an invalid handle!", entry.nvmap_handle);
-            std::memcpy(output.data(), entries.data(), output.size());
-            return static_cast<u32>(NvErrCodes::InvalidNmapHandle);
+            LOG_CRITICAL(Service_NVDRV, "invalid nvmap_handle={:X}", entry.nvmap_handle);
+            result = NvErrCodes::InvalidInput;
+            break;
         }
 
-        ASSERT(object->status == nvmap::Object::Status::Allocated);
+        const auto offset{static_cast<GPUVAddr>(entry.offset) << 0x10};
+        const auto size{static_cast<u64>(entry.pages) << 0x10};
+        const auto map_offset{static_cast<u64>(entry.map_offset) << 0x10};
+        const auto addr{system.GPU().MemoryManager().Map(object->addr + map_offset, offset, size)};
 
-        const u64 size = static_cast<u64>(entry.pages) << 0x10;
-        ASSERT(size <= object->size);
-        const u64 map_offset = static_cast<u64>(entry.map_offset) << 0x10;
-
-        const GPUVAddr returned =
-            gpu.MemoryManager().MapBufferEx(object->addr + map_offset, offset, size);
-        ASSERT(returned == offset);
+        if (!addr) {
+            LOG_CRITICAL(Service_NVDRV, "map returned an invalid address!");
+            result = NvErrCodes::InvalidInput;
+            break;
+        }
     }
+
     std::memcpy(output.data(), entries.data(), output.size());
-    return 0;
+    return result;
 }
 
 u32 nvhost_as_gpu::MapBufferEx(const std::vector<u8>& input, std::vector<u8>& output) {
@@ -126,44 +137,76 @@ u32 nvhost_as_gpu::MapBufferEx(const std::vector<u8>& input, std::vector<u8>& ou
               params.flags, params.nvmap_handle, params.buffer_offset, params.mapping_size,
               params.offset);
 
-    if (!params.nvmap_handle) {
-        return 0;
+    const auto object{nvmap_dev->GetObject(params.nvmap_handle)};
+    if (!object) {
+        LOG_CRITICAL(Service_NVDRV, "invalid nvmap_handle={:X}", params.nvmap_handle);
+        std::memcpy(output.data(), &params, output.size());
+        return NvErrCodes::InvalidInput;
     }
 
-    auto object = nvmap_dev->GetObject(params.nvmap_handle);
-    ASSERT(object);
-
-    // We can only map objects that have already been assigned a CPU address.
-    ASSERT(object->status == nvmap::Object::Status::Allocated);
-
-    ASSERT(params.buffer_offset == 0);
-
     // The real nvservices doesn't make a distinction between handles and ids, and
     // object can only have one handle and it will be the same as its id. Assert that this is the
     // case to prevent unexpected behavior.
     ASSERT(object->id == params.nvmap_handle);
-
     auto& gpu = system.GPU();
 
-    if (params.flags & 1) {
-        params.offset = gpu.MemoryManager().MapBufferEx(object->addr, params.offset, object->size);
-    } else {
-        params.offset = gpu.MemoryManager().MapBufferEx(object->addr, object->size);
+    u64 page_size{params.page_size};
+    if (!page_size) {
+        page_size = object->align;
+    }
+
+    if ((params.flags & AddressSpaceFlags::Remap) != AddressSpaceFlags::None) {
+        if (const auto buffer_map{FindBufferMap(params.offset)}; buffer_map) {
+            const auto cpu_addr{static_cast<VAddr>(buffer_map->CpuAddr() + params.buffer_offset)};
+            const auto gpu_addr{static_cast<GPUVAddr>(params.offset + params.buffer_offset)};
+
+            if (!gpu.MemoryManager().Map(cpu_addr, gpu_addr, params.mapping_size)) {
+                LOG_CRITICAL(Service_NVDRV,
+                             "remap failed, flags={:X}, nvmap_handle={:X}, buffer_offset={}, "
+                             "mapping_size = {}, offset={}",
+                             params.flags, params.nvmap_handle, params.buffer_offset,
+                             params.mapping_size, params.offset);
+
+                std::memcpy(output.data(), &params, output.size());
+                return NvErrCodes::InvalidInput;
+            }
+
+            std::memcpy(output.data(), &params, output.size());
+            return NvErrCodes::Success;
+        } else {
+            LOG_CRITICAL(Service_NVDRV, "address not mapped offset={}", params.offset);
+
+            std::memcpy(output.data(), &params, output.size());
+            return NvErrCodes::InvalidInput;
+        }
     }
 
-    // Create a new mapping entry for this operation.
-    ASSERT_MSG(buffer_mappings.find(params.offset) == buffer_mappings.end(),
-               "Offset is already mapped");
+    // We can only map objects that have already been assigned a CPU address.
+    ASSERT(object->status == nvmap::Object::Status::Allocated);
+
+    const auto physical_address{object->addr + params.buffer_offset};
+    u64 size{params.mapping_size};
+    if (!size) {
+        size = object->size;
+    }
 
-    BufferMapping mapping{};
-    mapping.nvmap_handle = params.nvmap_handle;
-    mapping.offset = params.offset;
-    mapping.size = object->size;
+    const bool is_alloc{(params.flags & AddressSpaceFlags::FixedOffset) == AddressSpaceFlags::None};
+    if (is_alloc) {
+        params.offset = gpu.MemoryManager().MapAllocate(physical_address, size, page_size);
+    } else {
+        params.offset = gpu.MemoryManager().Map(physical_address, params.offset, size);
+    }
 
-    buffer_mappings[params.offset] = mapping;
+    auto result{NvErrCodes::Success};
+    if (!params.offset) {
+        LOG_CRITICAL(Service_NVDRV, "failed to map size={}", size);
+        result = NvErrCodes::InvalidInput;
+    } else {
+        AddBufferMap(params.offset, size, physical_address, is_alloc);
+    }
 
     std::memcpy(output.data(), &params, output.size());
-    return 0;
+    return result;
 }
 
 u32 nvhost_as_gpu::UnmapBuffer(const std::vector<u8>& input, std::vector<u8>& output) {
@@ -172,24 +215,20 @@ u32 nvhost_as_gpu::UnmapBuffer(const std::vector<u8>& input, std::vector<u8>& ou
 
     LOG_DEBUG(Service_NVDRV, "called, offset=0x{:X}", params.offset);
 
-    const auto itr = buffer_mappings.find(params.offset);
-    if (itr == buffer_mappings.end()) {
-        LOG_WARNING(Service_NVDRV, "Tried to unmap an invalid offset 0x{:X}", params.offset);
-        // Hardware tests shows that unmapping an already unmapped buffer always returns successful
-        // and doesn't fail.
-        return 0;
+    if (const auto size{RemoveBufferMap(params.offset)}; size) {
+        system.GPU().MemoryManager().Unmap(params.offset, *size);
+    } else {
+        LOG_ERROR(Service_NVDRV, "invalid offset=0x{:X}", params.offset);
     }
 
-    params.offset = system.GPU().MemoryManager().UnmapBuffer(params.offset, itr->second.size);
-    buffer_mappings.erase(itr->second.offset);
-
     std::memcpy(output.data(), &params, output.size());
-    return 0;
+    return NvErrCodes::Success;
 }
 
 u32 nvhost_as_gpu::BindChannel(const std::vector<u8>& input, std::vector<u8>& output) {
     IoctlBindChannel params{};
     std::memcpy(&params, input.data(), input.size());
+
     LOG_DEBUG(Service_NVDRV, "called, fd={:X}", params.fd);
 
     channel = params.fd;
@@ -199,6 +238,7 @@ u32 nvhost_as_gpu::BindChannel(const std::vector<u8>& input, std::vector<u8>& ou
 u32 nvhost_as_gpu::GetVARegions(const std::vector<u8>& input, std::vector<u8>& output) {
     IoctlGetVaRegions params{};
     std::memcpy(&params, input.data(), input.size());
+
     LOG_WARNING(Service_NVDRV, "(STUBBED) called, buf_addr={:X}, buf_size={:X}", params.buf_addr,
                 params.buf_size);
 
@@ -210,9 +250,43 @@ u32 nvhost_as_gpu::GetVARegions(const std::vector<u8>& input, std::vector<u8>& o
     params.regions[1].offset = 0x04000000;
     params.regions[1].page_size = 0x10000;
     params.regions[1].pages = 0x1bffff;
+
     // TODO(ogniK): This probably can stay stubbed but should add support way way later
+
     std::memcpy(output.data(), &params, output.size());
     return 0;
 }
 
+std::optional<nvhost_as_gpu::BufferMap> nvhost_as_gpu::FindBufferMap(GPUVAddr gpu_addr) const {
+    const auto end{buffer_mappings.upper_bound(gpu_addr)};
+    for (auto iter{buffer_mappings.begin()}; iter != end; ++iter) {
+        if (gpu_addr >= iter->second.StartAddr() && gpu_addr < iter->second.EndAddr()) {
+            return iter->second;
+        }
+    }
+
+    return {};
+}
+
+void nvhost_as_gpu::AddBufferMap(GPUVAddr gpu_addr, std::size_t size, VAddr cpu_addr,
+                                 bool is_allocated) {
+    buffer_mappings[gpu_addr] = {gpu_addr, size, cpu_addr, is_allocated};
+}
+
+std::optional<std::size_t> nvhost_as_gpu::RemoveBufferMap(GPUVAddr gpu_addr) {
+    if (const auto iter{buffer_mappings.find(gpu_addr)}; iter != buffer_mappings.end()) {
+        std::size_t size{};
+
+        if (iter->second.IsAllocated()) {
+            size = iter->second.Size();
+        }
+
+        buffer_mappings.erase(iter);
+
+        return size;
+    }
+
+    return {};
+}
+
 } // namespace Service::Nvidia::Devices
diff --git a/src/core/hle/service/nvdrv/devices/nvhost_as_gpu.h b/src/core/hle/service/nvdrv/devices/nvhost_as_gpu.h
index f79fcc065..9a0cdff0c 100644
--- a/src/core/hle/service/nvdrv/devices/nvhost_as_gpu.h
+++ b/src/core/hle/service/nvdrv/devices/nvhost_as_gpu.h
@@ -4,9 +4,12 @@
 
 #pragma once
 
+#include <map>
 #include <memory>
-#include <unordered_map>
+#include <optional>
 #include <vector>
+
+#include "common/common_funcs.h"
 #include "common/common_types.h"
 #include "common/swap.h"
 #include "core/hle/service/nvdrv/devices/nvdevice.h"
@@ -15,6 +18,13 @@ namespace Service::Nvidia::Devices {
 
 class nvmap;
 
+enum class AddressSpaceFlags : u32 {
+    None = 0x0,
+    FixedOffset = 0x1,
+    Remap = 0x100,
+};
+DECLARE_ENUM_FLAG_OPERATORS(AddressSpaceFlags);
+
 class nvhost_as_gpu final : public nvdevice {
 public:
     explicit nvhost_as_gpu(Core::System& system, std::shared_ptr<nvmap> nvmap_dev);
@@ -25,6 +35,45 @@ public:
               IoctlVersion version) override;
 
 private:
+    class BufferMap final {
+    public:
+        constexpr BufferMap() = default;
+
+        constexpr BufferMap(GPUVAddr start_addr, std::size_t size)
+            : start_addr{start_addr}, end_addr{start_addr + size} {}
+
+        constexpr BufferMap(GPUVAddr start_addr, std::size_t size, VAddr cpu_addr,
+                            bool is_allocated)
+            : start_addr{start_addr}, end_addr{start_addr + size}, cpu_addr{cpu_addr},
+              is_allocated{is_allocated} {}
+
+        constexpr VAddr StartAddr() const {
+            return start_addr;
+        }
+
+        constexpr VAddr EndAddr() const {
+            return end_addr;
+        }
+
+        constexpr std::size_t Size() const {
+            return end_addr - start_addr;
+        }
+
+        constexpr VAddr CpuAddr() const {
+            return cpu_addr;
+        }
+
+        constexpr bool IsAllocated() const {
+            return is_allocated;
+        }
+
+    private:
+        GPUVAddr start_addr{};
+        GPUVAddr end_addr{};
+        VAddr cpu_addr{};
+        bool is_allocated{};
+    };
+
     enum class IoctlCommand : u32_le {
         IocInitalizeExCommand = 0x40284109,
         IocAllocateSpaceCommand = 0xC0184102,
@@ -49,7 +98,7 @@ private:
     struct IoctlAllocSpace {
         u32_le pages;
         u32_le page_size;
-        u32_le flags;
+        AddressSpaceFlags flags;
         INSERT_PADDING_WORDS(1);
         union {
             u64_le offset;
@@ -69,18 +118,18 @@ private:
     static_assert(sizeof(IoctlRemapEntry) == 20, "IoctlRemapEntry is incorrect size");
 
     struct IoctlMapBufferEx {
-        u32_le flags; // bit0: fixed_offset, bit2: cacheable
-        u32_le kind;  // -1 is default
+        AddressSpaceFlags flags; // bit0: fixed_offset, bit2: cacheable
+        u32_le kind;             // -1 is default
         u32_le nvmap_handle;
         u32_le page_size; // 0 means don't care
-        u64_le buffer_offset;
+        s64_le buffer_offset;
         u64_le mapping_size;
-        u64_le offset;
+        s64_le offset;
     };
     static_assert(sizeof(IoctlMapBufferEx) == 40, "IoctlMapBufferEx is incorrect size");
 
     struct IoctlUnmapBuffer {
-        u64_le offset;
+        s64_le offset;
     };
     static_assert(sizeof(IoctlUnmapBuffer) == 8, "IoctlUnmapBuffer is incorrect size");
 
@@ -106,15 +155,6 @@ private:
     static_assert(sizeof(IoctlGetVaRegions) == 16 + sizeof(IoctlVaRegion) * 2,
                   "IoctlGetVaRegions is incorrect size");
 
-    struct BufferMapping {
-        u64 offset;
-        u64 size;
-        u32 nvmap_handle;
-    };
-
-    /// Map containing the nvmap object mappings in GPU memory.
-    std::unordered_map<u64, BufferMapping> buffer_mappings;
-
     u32 channel{};
 
     u32 InitalizeEx(const std::vector<u8>& input, std::vector<u8>& output);
@@ -125,7 +165,14 @@ private:
     u32 BindChannel(const std::vector<u8>& input, std::vector<u8>& output);
     u32 GetVARegions(const std::vector<u8>& input, std::vector<u8>& output);
 
+    std::optional<BufferMap> FindBufferMap(GPUVAddr gpu_addr) const;
+    void AddBufferMap(GPUVAddr gpu_addr, std::size_t size, VAddr cpu_addr, bool is_allocated);
+    std::optional<std::size_t> RemoveBufferMap(GPUVAddr gpu_addr);
+
     std::shared_ptr<nvmap> nvmap_dev;
+
+    // This is expected to be ordered, therefore we must use a map, not unordered_map
+    std::map<GPUVAddr, BufferMap> buffer_mappings;
 };
 
 } // namespace Service::Nvidia::Devices
diff --git a/src/core/hle/service/nvdrv/devices/nvmap.cpp b/src/core/hle/service/nvdrv/devices/nvmap.cpp
index 8c742316c..7228910a0 100644
--- a/src/core/hle/service/nvdrv/devices/nvmap.cpp
+++ b/src/core/hle/service/nvdrv/devices/nvmap.cpp
@@ -18,7 +18,12 @@ enum {
 };
 }
 
-nvmap::nvmap(Core::System& system) : nvdevice(system) {}
+nvmap::nvmap(Core::System& system) : nvdevice(system) {
+    // Handle 0 appears to be used when remapping, so we create a placeholder empty nvmap object to
+    // represent this.
+    CreateObject(0);
+}
+
 nvmap::~nvmap() = default;
 
 VAddr nvmap::GetObjectAddress(u32 handle) const {
@@ -50,6 +55,21 @@ u32 nvmap::ioctl(Ioctl command, const std::vector<u8>& input, const std::vector<
     return 0;
 }
 
+u32 nvmap::CreateObject(u32 size) {
+    // Create a new nvmap object and obtain a handle to it.
+    auto object = std::make_shared<Object>();
+    object->id = next_id++;
+    object->size = size;
+    object->status = Object::Status::Created;
+    object->refcount = 1;
+
+    const u32 handle = next_handle++;
+
+    handles[handle] = std::move(object);
+
+    return handle;
+}
+
 u32 nvmap::IocCreate(const std::vector<u8>& input, std::vector<u8>& output) {
     IocCreateParams params;
     std::memcpy(&params, input.data(), sizeof(params));
@@ -59,17 +79,8 @@ u32 nvmap::IocCreate(const std::vector<u8>& input, std::vector<u8>& output) {
         LOG_ERROR(Service_NVDRV, "Size is 0");
         return static_cast<u32>(NvErrCodes::InvalidValue);
     }
-    // Create a new nvmap object and obtain a handle to it.
-    auto object = std::make_shared<Object>();
-    object->id = next_id++;
-    object->size = params.size;
-    object->status = Object::Status::Created;
-    object->refcount = 1;
-
-    u32 handle = next_handle++;
-    handles[handle] = std::move(object);
 
-    params.handle = handle;
+    params.handle = CreateObject(params.size);
 
     std::memcpy(output.data(), &params, sizeof(params));
     return 0;
diff --git a/src/core/hle/service/nvdrv/devices/nvmap.h b/src/core/hle/service/nvdrv/devices/nvmap.h
index 73c2e8809..84624be00 100644
--- a/src/core/hle/service/nvdrv/devices/nvmap.h
+++ b/src/core/hle/service/nvdrv/devices/nvmap.h
@@ -49,10 +49,10 @@ public:
 
 private:
     /// Id to use for the next handle that is created.
-    u32 next_handle = 1;
+    u32 next_handle = 0;
 
     /// Id to use for the next object that is created.
-    u32 next_id = 1;
+    u32 next_id = 0;
 
     /// Mapping of currently allocated handles to the objects they represent.
     std::unordered_map<u32, std::shared_ptr<Object>> handles;
@@ -119,6 +119,8 @@ private:
     };
     static_assert(sizeof(IocGetIdParams) == 8, "IocGetIdParams has wrong size");
 
+    u32 CreateObject(u32 size);
+
     u32 IocCreate(const std::vector<u8>& input, std::vector<u8>& output);
     u32 IocAlloc(const std::vector<u8>& input, std::vector<u8>& output);
     u32 IocGetId(const std::vector<u8>& input, std::vector<u8>& output);