GPUMemoryManager: Fix mapping ranges.

Credits to ByLaws from Skyline-emu for pointing this out.

src/core/hle/service/nvdrv/devices/nvhost_as_gpu.cpp

@@ -146,6 +146,14 @@ NvResult nvhost_as_gpu::Remap(const std::vector<u8>& input, std::vector<u8>& out
         LOG_DEBUG(Service_NVDRV, "remap entry, offset=0x{:X} handle=0x{:X} pages=0x{:X}",
                   entry.offset, entry.nvmap_handle, entry.pages);
 
+        if (entry.nvmap_handle == 0) {
+            // If nvmap handle is null, we should unmap instead.
+            const auto offset{static_cast<GPUVAddr>(entry.offset) << 0x10};
+            const auto size{static_cast<u64>(entry.pages) << 0x10};
+            system.GPU().MemoryManager().Unmap(offset, size);
+            continue;
+        }
+
         const auto object{nvmap_dev->GetObject(entry.nvmap_handle)};
         if (!object) {
             LOG_CRITICAL(Service_NVDRV, "invalid nvmap_handle={:X}", entry.nvmap_handle);
@@ -179,24 +187,8 @@ NvResult nvhost_as_gpu::MapBufferEx(const std::vector<u8>& input, std::vector<u8
               params.flags, params.nvmap_handle, params.buffer_offset, params.mapping_size,
               params.offset);
 
-    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 NvResult::InvalidState;
-    }
-
-    // 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();
 
-    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)};
@@ -223,6 +215,23 @@ NvResult nvhost_as_gpu::MapBufferEx(const std::vector<u8>& input, std::vector<u8
         }
     }
 
+    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 NvResult::InvalidState;
+    }
+
+    // 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);
+
+    u64 page_size{params.page_size};
+    if (!page_size) {
+        page_size = object->align;
+    }
+
     // We can only map objects that have already been assigned a CPU address.
     ASSERT(object->status == nvmap::Object::Status::Allocated);
 

src/core/hle/service/nvdrv/devices/nvmap.cpp

@@ -11,11 +11,7 @@
 
 namespace Service::Nvidia::Devices {
 
-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(Core::System& system_) : nvdevice{system_} {}
 
 nvmap::~nvmap() = default;
 

src/core/hle/service/nvdrv/devices/nvmap.h

@@ -56,10 +56,10 @@ public:
 
 private:
     /// Id to use for the next handle that is created.
-    u32 next_handle = 0;
+    u32 next_handle = 1;
 
     /// Id to use for the next object that is created.
-    u32 next_id = 0;
+    u32 next_id = 1;
 
     /// Mapping of currently allocated handles to the objects they represent.
     std::unordered_map<u32, std::shared_ptr<Object>> handles;

src/video_core/memory_manager.cpp

@@ -38,13 +38,110 @@ GPUVAddr MemoryManager::UpdateRange(GPUVAddr gpu_addr, PageEntry page_entry, std
     return gpu_addr;
 }
 
-GPUVAddr MemoryManager::Map(VAddr cpu_addr, GPUVAddr gpu_addr, std::size_t size) {
-    const auto it = std::ranges::lower_bound(map_ranges, gpu_addr, {}, &MapRange::first);
-    if (it != map_ranges.end() && it->first == gpu_addr) {
-        it->second = size;
-    } else {
-        map_ranges.insert(it, MapRange{gpu_addr, size});
+MemoryManager::MemoryRegion::MemoryRegion() = default;
+
+MemoryManager::MemoryRegion::MemoryRegion(GPUVAddr gpu_addr_, size_t size_)
+    : address{gpu_addr_}, size{size_} {}
+
+void MemoryManager::InsertRange(GPUVAddr gpu_addr, size_t size) {
+    if (map_ranges.empty()) {
+        map_ranges.insert(MemoryRegion{gpu_addr, size});
+        return;
     }
+    const GPUVAddr gpu_addr_end = gpu_addr + size;
+    auto clear_and_map = [&](std::set<MemoryRegion>::iterator it_start) {
+        auto it_end = map_ranges.upper_bound(MemoryRegion{gpu_addr_end - 1, 1});
+        GPUVAddr alloc_gpu_addr = gpu_addr;
+        GPUVAddr alloc_gpu_addr_end = gpu_addr_end;
+        while (it_start != it_end) {
+            const GPUVAddr local_addr_end = it_start->address + it_start->size;
+            if (it_start->address < gpu_addr) {
+                if (local_addr_end <= gpu_addr) {
+                    it_start++;
+                    continue;
+                }
+                alloc_gpu_addr = std::min(alloc_gpu_addr, it_start->address);
+            }
+            if (local_addr_end > gpu_addr_end) {
+                alloc_gpu_addr_end = std::max(alloc_gpu_addr_end, local_addr_end);
+            }
+            it_start = map_ranges.erase(it_start);
+            //@Note: do I have to recalculate the upper_bound or are the iterators safe after
+            // delete?
+            it_end = map_ranges.upper_bound(MemoryRegion{gpu_addr_end - 1, 1});
+        }
+        size_t new_size = alloc_gpu_addr_end - alloc_gpu_addr;
+        map_ranges.insert(MemoryRegion{alloc_gpu_addr, new_size});
+    };
+    auto it_first = map_ranges.upper_bound(MemoryRegion{gpu_addr, size});
+    --it_first;
+    if (it_first != map_ranges.end()) {
+        clear_and_map(it_first);
+    } else {
+        auto it = map_ranges.lower_bound(MemoryRegion{gpu_addr, size});
+        if (it == map_ranges.end() || it->address >= gpu_addr_end) {
+            map_ranges.insert(MemoryRegion{gpu_addr, size});
+            return;
+        }
+        clear_and_map(it);
+    }
+}
+
+void MemoryManager::RemoveRange(GPUVAddr gpu_addr, size_t size) {
+    if (map_ranges.empty()) {
+        UNREACHABLE_MSG("Unmapping non-existent GPU address=0x{:x}", gpu_addr);
+        return;
+    }
+    const GPUVAddr gpu_addr_end = gpu_addr + size;
+    auto clear_and_unmap = [&](std::set<MemoryRegion>::iterator it_start) {
+        auto it_end = map_ranges.upper_bound(MemoryRegion{gpu_addr_end - 1, 1});
+        while (it_start != it_end) {
+            const GPUVAddr local_addr_end = it_start->address + it_start->size;
+            if (it_start->address < gpu_addr) {
+                if (local_addr_end <= gpu_addr) {
+                    it_start++;
+                    continue;
+                }
+                MemoryRegion new_region{};
+                new_region.address = it_start->address;
+                new_region.size = local_addr_end - gpu_addr;
+                map_ranges.erase(it_start);
+                auto pair = map_ranges.insert(new_region);
+                it_start = pair.first;
+                it_end = map_ranges.upper_bound(MemoryRegion{gpu_addr_end - 1, 1});
+                it_start++;
+                continue;
+            }
+            if (local_addr_end > gpu_addr_end) {
+                MemoryRegion new_region{};
+                new_region.address = gpu_addr_end;
+                new_region.size = local_addr_end - gpu_addr_end;
+                map_ranges.insert(new_region);
+                break;
+            }
+            it_start = map_ranges.erase(it_start);
+            //@Note: do I have to recalculate the upper_bound or are the iterators safe after
+            // delete?
+            it_end = map_ranges.upper_bound(MemoryRegion{gpu_addr_end - 1, 1});
+        }
+    };
+    auto it_first = map_ranges.upper_bound(MemoryRegion{gpu_addr, size});
+    --it_first;
+    if (it_first != map_ranges.end()) {
+        clear_and_unmap(it_first);
+    } else {
+        auto it = map_ranges.lower_bound(MemoryRegion{gpu_addr, size});
+        if (it == map_ranges.end() || it->address >= gpu_addr_end) {
+            UNREACHABLE_MSG("Unmapping non-existent GPU address=0x{:x}", gpu_addr);
+            return;
+        }
+        clear_and_unmap(it);
+    }
+}
+
+GPUVAddr MemoryManager::Map(VAddr cpu_addr, GPUVAddr gpu_addr, std::size_t size) {
+    InsertRange(gpu_addr, size);
+    ASSERT(IsFullyMapped(gpu_addr, size));
     return UpdateRange(gpu_addr, cpu_addr, size);
 }
 
@@ -62,13 +159,8 @@ void MemoryManager::Unmap(GPUVAddr gpu_addr, std::size_t size) {
     if (size == 0) {
         return;
     }
-    const auto it = std::ranges::lower_bound(map_ranges, gpu_addr, {}, &MapRange::first);
-    if (it != map_ranges.end()) {
-        ASSERT(it->first == gpu_addr);
-        map_ranges.erase(it);
-    } else {
-        UNREACHABLE_MSG("Unmapping non-existent GPU address=0x{:x}", gpu_addr);
-    }
+    RemoveRange(gpu_addr, size);
+    ASSERT(IsFullyUnmapped(gpu_addr, size));
 
     const auto submapped_ranges = GetSubmappedRange(gpu_addr, size);
 
@@ -258,10 +350,39 @@ const u8* MemoryManager::GetPointer(GPUVAddr gpu_addr) const {
     return system.Memory().GetPointer(*address);
 }
 
-size_t MemoryManager::BytesToMapEnd(GPUVAddr gpu_addr) const noexcept {
-    auto it = std::ranges::upper_bound(map_ranges, gpu_addr, {}, &MapRange::first);
+bool MemoryManager::IsFullyMapped(GPUVAddr gpu_addr, size_t size) const {
+    if (map_ranges.empty()) {
+        return false;
+    }
+    auto it = map_ranges.upper_bound(MemoryRegion{gpu_addr, size});
     --it;
-    return it->second - (gpu_addr - it->first);
+    if (it == map_ranges.end()) {
+        return false;
+    }
+    const GPUVAddr gpu_addr_end = gpu_addr + size;
+    return it->address <= gpu_addr && (it->address + it->size) >= gpu_addr_end;
+}
+
+bool MemoryManager::IsFullyUnmapped(GPUVAddr gpu_addr, size_t size) const {
+    if (map_ranges.empty()) {
+        return true;
+    }
+    auto it = map_ranges.upper_bound(MemoryRegion{gpu_addr, size});
+    const GPUVAddr gpu_addr_end = gpu_addr + size;
+    if (it != map_ranges.end() && it->address < gpu_addr_end) {
+        return false;
+    }
+    --it;
+    if (it == map_ranges.end()) {
+        return true;
+    }
+    return (it->address + it->size) <= gpu_addr;
+}
+
+size_t MemoryManager::BytesToMapEnd(GPUVAddr gpu_addr) const noexcept {
+    auto it = map_ranges.upper_bound(MemoryRegion{gpu_addr, 1});
+    --it;
+    return it->size - (gpu_addr - it->address);
 }
 
 void MemoryManager::ReadBlock(GPUVAddr gpu_src_addr, void* dest_buffer, std::size_t size) const {

src/video_core/memory_manager.h

@@ -6,6 +6,7 @@
 
 #include <map>
 #include <optional>
+#include <set>
 #include <vector>
 
 #include "common/common_types.h"
@@ -155,6 +156,13 @@ private:
 
     void FlushRegion(GPUVAddr gpu_addr, size_t size) const;
 
+    void InsertRange(GPUVAddr gpu_addr, size_t size);
+    void RemoveRange(GPUVAddr gpu_addr, size_t size);
+
+    // For debugging only
+    bool IsFullyMapped(GPUVAddr gpu_addr, size_t size) const;
+    bool IsFullyUnmapped(GPUVAddr gpu_addr, size_t size) const;
+
     [[nodiscard]] static constexpr std::size_t PageEntryIndex(GPUVAddr gpu_addr) {
         return (gpu_addr >> page_bits) & page_table_mask;
     }
@@ -175,8 +183,22 @@ private:
 
     std::vector<PageEntry> page_table;
 
-    using MapRange = std::pair<GPUVAddr, size_t>;
-    std::vector<MapRange> map_ranges;
+    struct MemoryRegion {
+        MemoryRegion();
+        MemoryRegion(GPUVAddr gpu_addr, size_t size);
+
+        bool operator<(const MemoryRegion& other) const {
+            return address < other.address;
+        }
+
+        bool operator==(const MemoryRegion&) const = default;
+        bool operator!=(const MemoryRegion&) const = default;
+
+        GPUVAddr address{};
+        size_t size{};
+    };
+
+    std::set<MemoryRegion> map_ranges;
 
     std::vector<std::pair<VAddr, std::size_t>> cache_invalidate_queue;
 };