1 files changed, 49 insertions, 75 deletions

diff --git a/src/video_core/rasterizer_cache.h b/src/video_core/rasterizer_cache.h
index 7a0492a4e..de1eab86b 100644
--- a/src/video_core/rasterizer_cache.h
+++ b/src/video_core/rasterizer_cache.h
@@ -4,113 +4,87 @@
 
 #pragma once
 
-#include <unordered_map>
+#include <set>
+
 #include <boost/icl/interval_map.hpp>
-#include <boost/range/iterator_range.hpp>
 
 #include "common/common_types.h"
+#include "core/core.h"
 #include "core/memory.h"
 #include "video_core/memory_manager.h"
+#include "video_core/rasterizer_interface.h"
+#include "video_core/renderer_base.h"
 
 template <class T>
 class RasterizerCache : NonCopyable {
 public:
     /// Mark the specified region as being invalidated
-    void InvalidateRegion(Tegra::GPUVAddr region_addr, size_t region_size) {
-        for (auto iter = cached_objects.cbegin(); iter != cached_objects.cend();) {
-            const auto& object{iter->second};
+    void InvalidateRegion(VAddr addr, u64 size) {
+        if (size == 0)
+            return;
 
-            ++iter;
+        const ObjectInterval interval{addr, addr + size};
+        for (auto& pair : boost::make_iterator_range(object_cache.equal_range(interval))) {
+            for (auto& cached_object : pair.second) {
+                if (!cached_object)
+                    continue;
 
-            if (object->GetAddr() <= (region_addr + region_size) &&
-                region_addr <= (object->GetAddr() + object->GetSizeInBytes())) {
-                // Regions overlap, so invalidate
-                Unregister(object);
+                remove_objects.emplace(cached_object);
             }
         }
+
+        for (auto& remove_object : remove_objects) {
+            Unregister(remove_object);
+        }
+
+        remove_objects.clear();
+    }
+
+    /// Invalidates everything in the cache
+    void InvalidateAll() {
+        while (object_cache.begin() != object_cache.end()) {
+            Unregister(*object_cache.begin()->second.begin());
+        }
     }
 
 protected:
     /// Tries to get an object from the cache with the specified address
-    T TryGet(Tegra::GPUVAddr addr) const {
-        const auto& search{cached_objects.find(addr)};
-        if (search != cached_objects.end()) {
-            return search->second;
+    T TryGet(VAddr addr) const {
+        const ObjectInterval interval{addr};
+        for (auto& pair : boost::make_iterator_range(object_cache.equal_range(interval))) {
+            for (auto& cached_object : pair.second) {
+                if (cached_object->GetAddr() == addr) {
+                    return cached_object;
+                }
+            }
         }
-
         return nullptr;
     }
 
-    /// Gets a reference to the cache
-    const std::unordered_map<Tegra::GPUVAddr, T>& GetCache() const {
-        return cached_objects;
-    }
-
     /// Register an object into the cache
     void Register(const T& object) {
-        const auto& search{cached_objects.find(object->GetAddr())};
-        if (search != cached_objects.end()) {
-            // Registered already
-            return;
-        }
-
-        cached_objects[object->GetAddr()] = object;
-        UpdatePagesCachedCount(object->GetAddr(), object->GetSizeInBytes(), 1);
+        object_cache.add({GetInterval(object), ObjectSet{object}});
+        auto& rasterizer = Core::System::GetInstance().Renderer().Rasterizer();
+        rasterizer.UpdatePagesCachedCount(object->GetAddr(), object->GetSizeInBytes(), 1);
     }
 
     /// Unregisters an object from the cache
     void Unregister(const T& object) {
-        const auto& search{cached_objects.find(object->GetAddr())};
-        if (search == cached_objects.end()) {
-            // Unregistered already
-            return;
-        }
-
-        UpdatePagesCachedCount(object->GetAddr(), object->GetSizeInBytes(), -1);
-        cached_objects.erase(search);
+        auto& rasterizer = Core::System::GetInstance().Renderer().Rasterizer();
+        rasterizer.UpdatePagesCachedCount(object->GetAddr(), object->GetSizeInBytes(), -1);
+        object_cache.subtract({GetInterval(object), ObjectSet{object}});
     }
 
 private:
-    using PageMap = boost::icl::interval_map<u64, int>;
-
-    template <typename Map, typename Interval>
-    constexpr auto RangeFromInterval(Map& map, const Interval& interval) {
-        return boost::make_iterator_range(map.equal_range(interval));
-    }
-
-    /// Increase/decrease the number of object in pages touching the specified region
-    void UpdatePagesCachedCount(Tegra::GPUVAddr addr, u64 size, int delta) {
-        const u64 page_start{addr >> Tegra::MemoryManager::PAGE_BITS};
-        const u64 page_end{(addr + size) >> Tegra::MemoryManager::PAGE_BITS};
-
-        // Interval maps will erase segments if count reaches 0, so if delta is negative we have to
-        // subtract after iterating
-        const auto pages_interval = PageMap::interval_type::right_open(page_start, page_end);
-        if (delta > 0)
-            cached_pages.add({pages_interval, delta});
-
-        for (const auto& pair : RangeFromInterval(cached_pages, pages_interval)) {
-            const auto interval = pair.first & pages_interval;
-            const int count = pair.second;
-
-            const Tegra::GPUVAddr interval_start_addr = boost::icl::first(interval)
-                                                        << Tegra::MemoryManager::PAGE_BITS;
-            const Tegra::GPUVAddr interval_end_addr = boost::icl::last_next(interval)
-                                                      << Tegra::MemoryManager::PAGE_BITS;
-            const u64 interval_size = interval_end_addr - interval_start_addr;
-
-            if (delta > 0 && count == delta)
-                Memory::RasterizerMarkRegionCached(interval_start_addr, interval_size, true);
-            else if (delta < 0 && count == -delta)
-                Memory::RasterizerMarkRegionCached(interval_start_addr, interval_size, false);
-            else
-                ASSERT(count >= 0);
-        }
+    using ObjectSet = std::set<T>;
+    using ObjectCache = boost::icl::interval_map<VAddr, ObjectSet>;
+    using ObjectInterval = typename ObjectCache::interval_type;
 
-        if (delta < 0)
-            cached_pages.add({pages_interval, delta});
+    static auto GetInterval(const T& object) {
+        return ObjectInterval::right_open(object->GetAddr(),
+                                          object->GetAddr() + object->GetSizeInBytes());
     }
 
-    std::unordered_map<Tegra::GPUVAddr, T> cached_objects;
-    PageMap cached_pages;
+    ObjectCache object_cache;
+    ObjectSet remove_objects;
 };