From bda177ef409242c7ce200f55c1d8eea0e5aba6d3 Mon Sep 17 00:00:00 2001
From: ReinUsesLisp <reinuseslisp@airmail.cc>
Date: Fri, 22 Jan 2021 18:31:08 -0300
Subject: [PATCH] video_core/memory_manager: Add BytesToMapEnd

Track map address sizes in a flat ordered map and add a method to query
the number of bytes until the end of a map in a given address.
---
 src/video_core/memory_manager.cpp | 23 +++++++++++++++++++++--
 src/video_core/memory_manager.h   |  6 ++++++
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/src/video_core/memory_manager.cpp b/src/video_core/memory_manager.cpp
index 65feff588d..f99fa828e5 100644
--- a/src/video_core/memory_manager.cpp
+++ b/src/video_core/memory_manager.cpp
@@ -4,6 +4,7 @@
 
 #include "common/alignment.h"
 #include "common/assert.h"
+#include "common/logging/log.h"
 #include "core/core.h"
 #include "core/hle/kernel/memory/page_table.h"
 #include "core/hle/kernel/process.h"
@@ -38,6 +39,12 @@ GPUVAddr MemoryManager::UpdateRange(GPUVAddr gpu_addr, PageEntry page_entry, std
 }
 
 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});
+    }
     return UpdateRange(gpu_addr, cpu_addr, size);
 }
 
@@ -52,10 +59,16 @@ GPUVAddr MemoryManager::MapAllocate32(VAddr cpu_addr, std::size_t size) {
 }
 
 void MemoryManager::Unmap(GPUVAddr gpu_addr, std::size_t size) {
-    if (!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);
+    }
     // Flush and invalidate through the GPU interface, to be asynchronous if possible.
     const std::optional<VAddr> cpu_addr = GpuToCpuAddress(gpu_addr);
     ASSERT(cpu_addr);
@@ -218,6 +231,12 @@ 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);
+    --it;
+    return it->second - (gpu_addr - it->first);
+}
+
 void MemoryManager::ReadBlock(GPUVAddr gpu_src_addr, void* dest_buffer, std::size_t size) const {
     std::size_t remaining_size{size};
     std::size_t page_index{gpu_src_addr >> page_bits};
diff --git a/src/video_core/memory_manager.h b/src/video_core/memory_manager.h
index c35e576894..ca8e8328e4 100644
--- a/src/video_core/memory_manager.h
+++ b/src/video_core/memory_manager.h
@@ -85,6 +85,9 @@ public:
     [[nodiscard]] u8* GetPointer(GPUVAddr addr);
     [[nodiscard]] const u8* GetPointer(GPUVAddr addr) const;
 
+    /// Returns the number of bytes until the end of the memory map containing the given GPU address
+    [[nodiscard]] size_t BytesToMapEnd(GPUVAddr gpu_addr) const noexcept;
+
     /**
      * ReadBlock and WriteBlock are full read and write operations over virtual
      * GPU Memory. It's important to use these when GPU memory may not be continuous
@@ -150,6 +153,9 @@ private:
     VideoCore::RasterizerInterface* rasterizer = nullptr;
 
     std::vector<PageEntry> page_table;
+
+    using MapRange = std::pair<GPUVAddr, size_t>;
+    std::vector<MapRange> map_ranges;
 };
 
 } // namespace Tegra