1 files changed, 63 insertions, 16 deletions

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