2 files changed, 76 insertions, 152 deletions

diff --git a/src/video_core/shader/const_buffer_locker.cpp b/src/video_core/shader/const_buffer_locker.cpp
index 37a0968a1..ebeba102d 100644
--- a/src/video_core/shader/const_buffer_locker.cpp
+++ b/src/video_core/shader/const_buffer_locker.cpp
@@ -4,6 +4,8 @@
 
 #pragma once
 
+#include <algorithm>
+#include <memory>
 #include "common/assert.h"
 #include "common/common_types.h"
 #include "video_core/engines/maxwell_3d.h"
@@ -11,140 +13,92 @@
 
 namespace VideoCommon::Shader {
 
+using Tegra::Engines::SamplerDescriptor;
+
 ConstBufferLocker::ConstBufferLocker(Tegra::Engines::ShaderType shader_stage)
-    : engine{nullptr}, shader_stage{shader_stage} {}
+    : stage{shader_stage} {}
 
 ConstBufferLocker::ConstBufferLocker(Tegra::Engines::ShaderType shader_stage,
                                      Tegra::Engines::ConstBufferEngineInterface& engine)
-    : engine{&engine}, shader_stage{shader_stage} {}
-
-bool ConstBufferLocker::IsEngineSet() const {
-    return engine != nullptr;
-}
-
-void ConstBufferLocker::SetEngine(Tegra::Engines::ConstBufferEngineInterface& engine_) {
-    engine = &engine_;
-}
+    : stage{shader_stage}, engine{&engine} {}
 
 std::optional<u32> ConstBufferLocker::ObtainKey(u32 buffer, u32 offset) {
-    if (!keys) {
-        keys = std::make_shared<KeyMap>();
-    }
-    auto& key_map = *keys;
     const std::pair<u32, u32> key = {buffer, offset};
-    const auto iter = key_map.find(key);
-    if (iter != key_map.end()) {
-        return {iter->second};
+    const auto iter = keys.find(key);
+    if (iter != keys.end()) {
+        return iter->second;
     }
-    if (!IsEngineSet()) {
+    if (!engine) {
         return {};
     }
-    const u32 value = engine->AccessConstBuffer32(shader_stage, buffer, offset);
-    key_map.emplace(key, value);
-    return {value};
+    const u32 value = engine->AccessConstBuffer32(stage, buffer, offset);
+    keys.emplace(key, value);
+    return value;
 }
 
-std::optional<Tegra::Engines::SamplerDescriptor> ConstBufferLocker::ObtainBoundSampler(u32 offset) {
-    if (!bound_samplers) {
-        bound_samplers = std::make_shared<BoundSamplerMap>();
-    }
-    auto& key_map = *bound_samplers;
+std::optional<SamplerDescriptor> ConstBufferLocker::ObtainBoundSampler(u32 offset) {
     const u32 key = offset;
-    const auto iter = key_map.find(key);
-    if (iter != key_map.end()) {
-        return {iter->second};
+    const auto iter = bound_samplers.find(key);
+    if (iter != bound_samplers.end()) {
+        return iter->second;
     }
-    if (!IsEngineSet()) {
+    if (!engine) {
         return {};
     }
-    const Tegra::Engines::SamplerDescriptor value =
-        engine->AccessBoundSampler(shader_stage, offset);
-    key_map.emplace(key, value);
-    return {value};
+    const SamplerDescriptor value = engine->AccessBoundSampler(stage, offset);
+    bound_samplers.emplace(key, value);
+    return value;
 }
 
 std::optional<Tegra::Engines::SamplerDescriptor> ConstBufferLocker::ObtainBindlessSampler(
     u32 buffer, u32 offset) {
-    if (!bindless_samplers) {
-        bindless_samplers = std::make_shared<BindlessSamplerMap>();
-    }
-    auto& key_map = *bindless_samplers;
-    const std::pair<u32, u32> key = {buffer, offset};
-    const auto iter = key_map.find(key);
-    if (iter != key_map.end()) {
-        return {iter->second};
+    const std::pair key = {buffer, offset};
+    const auto iter = bindless_samplers.find(key);
+    if (iter != bindless_samplers.end()) {
+        return iter->second;
     }
-    if (!IsEngineSet()) {
+    if (!engine) {
         return {};
     }
-    const Tegra::Engines::SamplerDescriptor value =
-        engine->AccessBindlessSampler(shader_stage, buffer, offset);
-    key_map.emplace(key, value);
-    return {value};
+    const SamplerDescriptor value = engine->AccessBindlessSampler(stage, buffer, offset);
+    bindless_samplers.emplace(key, value);
+    return value;
 }
 
 void ConstBufferLocker::InsertKey(u32 buffer, u32 offset, u32 value) {
-    if (!keys) {
-        keys = std::make_shared<KeyMap>();
-    }
-    const std::pair<u32, u32> key = {buffer, offset};
-    (*keys)[key] = value;
+    keys.insert_or_assign({buffer, offset}, value);
 }
 
-void ConstBufferLocker::InsertBoundSampler(u32 offset, Tegra::Engines::SamplerDescriptor sampler) {
-    if (!bound_samplers) {
-        bound_samplers = std::make_shared<BoundSamplerMap>();
-    }
-    (*bound_samplers)[offset] = sampler;
+void ConstBufferLocker::InsertBoundSampler(u32 offset, SamplerDescriptor sampler) {
+    bound_samplers.insert_or_assign(offset, sampler);
 }
 
-void ConstBufferLocker::InsertBindlessSampler(u32 buffer, u32 offset,
-                                              Tegra::Engines::SamplerDescriptor sampler) {
-    if (!bindless_samplers) {
-        bindless_samplers = std::make_shared<BindlessSamplerMap>();
-    }
-    const std::pair<u32, u32> key = {buffer, offset};
-    (*bindless_samplers)[key] = sampler;
+void ConstBufferLocker::InsertBindlessSampler(u32 buffer, u32 offset, SamplerDescriptor sampler) {
+    bindless_samplers.insert_or_assign({buffer, offset}, sampler);
 }
 
 bool ConstBufferLocker::IsConsistent() const {
-    if (!IsEngineSet()) {
+    if (!engine) {
         return false;
     }
-    if (keys) {
-        for (const auto& key_val : *keys) {
-            const std::pair<u32, u32> key = key_val.first;
-            const u32 value = key_val.second;
-            const u32 other_value =
-                engine->AccessConstBuffer32(shader_stage, key.first, key.second);
-            if (other_value != value) {
-                return false;
-            }
-        }
-    }
-    if (bound_samplers) {
-        for (const auto& sampler_val : *bound_samplers) {
-            const u32 key = sampler_val.first;
-            const Tegra::Engines::SamplerDescriptor value = sampler_val.second;
-            const Tegra::Engines::SamplerDescriptor other_value =
-                engine->AccessBoundSampler(shader_stage, key);
-            if (other_value.raw != value.raw) {
-                return false;
-            }
-        }
-    }
-    if (bindless_samplers) {
-        for (const auto& sampler_val : *bindless_samplers) {
-            const std::pair<u32, u32> key = sampler_val.first;
-            const Tegra::Engines::SamplerDescriptor value = sampler_val.second;
-            const Tegra::Engines::SamplerDescriptor other_value =
-                engine->AccessBindlessSampler(shader_stage, key.first, key.second);
-            if (other_value.raw != value.raw) {
-                return false;
-            }
-        }
-    }
-    return true;
+    return std::all_of(keys.begin(), keys.end(),
+                       [](const auto& key) {
+                           const auto [value, other_value] = key.first;
+                           return value == other_value;
+                       }) &&
+           std::all_of(bound_samplers.begin(), bound_samplers.end(),
+                       [this](const auto& sampler) {
+                           const auto [key, value] = sampler;
+                           const auto other_value = engine->AccessBoundSampler(stage, key);
+                           return value.raw == other_value.raw;
+                       }) &&
+           std::all_of(
+               bindless_samplers.begin(), bindless_samplers.end(), [this](const auto& sampler) {
+                   const auto [cbuf, offset] = sampler.first;
+                   const auto value = sampler.second;
+                   const auto other_value = engine->AccessBindlessSampler(stage, cbuf, offset);
+                   return value.raw == other_value.raw;
+               });
 }
 
 } // namespace VideoCommon::Shader
diff --git a/src/video_core/shader/const_buffer_locker.h b/src/video_core/shader/const_buffer_locker.h
index 54459977f..417d5a16f 100644
--- a/src/video_core/shader/const_buffer_locker.h
+++ b/src/video_core/shader/const_buffer_locker.h
@@ -23,78 +23,48 @@ public:
     explicit ConstBufferLocker(Tegra::Engines::ShaderType shader_stage,
                                Tegra::Engines::ConstBufferEngineInterface& engine);
 
-    // Checks if an engine is setup, it may be possible that during disk shader
-    // cache run, the engines have not been created yet.
-    bool IsEngineSet() const;
-
-    // Use this to set/change the engine used for this shader.
-    void SetEngine(Tegra::Engines::ConstBufferEngineInterface& engine);
-
-    // Retrieves a key from the locker, if it's registered, it will give the
-    // registered value, if not it will obtain it from maxwell3d and register it.
+    /// Retrieves a key from the locker, if it's registered, it will give the registered value, if
+    /// not it will obtain it from maxwell3d and register it.
     std::optional<u32> ObtainKey(u32 buffer, u32 offset);
 
     std::optional<Tegra::Engines::SamplerDescriptor> ObtainBoundSampler(u32 offset);
 
     std::optional<Tegra::Engines::SamplerDescriptor> ObtainBindlessSampler(u32 buffer, u32 offset);
 
-    // Manually inserts a key.
+    /// Inserts a key.
     void InsertKey(u32 buffer, u32 offset, u32 value);
 
+    /// Inserts a bound sampler key.
     void InsertBoundSampler(u32 offset, Tegra::Engines::SamplerDescriptor sampler);
 
+    /// Inserts a bindless sampler key.
     void InsertBindlessSampler(u32 buffer, u32 offset, Tegra::Engines::SamplerDescriptor sampler);
 
-    // Retrieves the number of keys registered.
-    std::size_t NumKeys() const {
-        if (!keys) {
-            return 0;
-        }
-        return keys->size();
-    }
-
-    std::size_t NumBoundSamplers() const {
-        if (!bound_samplers) {
-            return 0;
-        }
-        return bound_samplers->size();
-    }
-
-    std::size_t NumBindlessSamplers() const {
-        if (!bindless_samplers) {
-            return 0;
-        }
-        return bindless_samplers->size();
-    }
+    /// Checks keys and samplers against engine's current const buffers. Returns true if they are
+    /// the same value, false otherwise;
+    bool IsConsistent() const;
 
-    // Gives an accessor to the key's database.
-    // Pre: NumKeys > 0
-    const KeyMap& AccessKeys() const {
-        return *keys;
+    /// Gives an getter to the const buffer keys in the database.
+    const KeyMap& GetKeys() const {
+        return keys;
     }
 
-    // Gives an accessor to the sampler's database.
-    // Pre: NumBindlessSamplers > 0
-    const BoundSamplerMap& AccessBoundSamplers() const {
-        return *bound_samplers;
+    /// Gets samplers database.
+    const BoundSamplerMap& GetBoundSamplers() const {
+        return bound_samplers;
     }
 
-    // Gives an accessor to the sampler's database.
-    // Pre: NumBindlessSamplers > 0
-    const BindlessSamplerMap& AccessBindlessSamplers() const {
-        return *bindless_samplers;
+    /// Gets bindless samplers database.
+    const BindlessSamplerMap& GetBindlessSamplers() const {
+        return bindless_samplers;
     }
 
-    // Checks keys & samplers against engine's current const buffers. Returns true if they
-    // are the same value, false otherwise;
-    bool IsConsistent() const;
-
 private:
-    Tegra::Engines::ConstBufferEngineInterface* engine;
-    Tegra::Engines::ShaderType shader_stage;
-    // All containers are lazy initialized as most shaders don't use them.
-    std::shared_ptr<KeyMap> keys{};
-    std::shared_ptr<BoundSamplerMap> bound_samplers{};
-    std::shared_ptr<BindlessSamplerMap> bindless_samplers{};
+    const Tegra::Engines::ShaderType stage;
+    Tegra::Engines::ConstBufferEngineInterface* engine = nullptr;
+    KeyMap keys;
+    BoundSamplerMap bound_samplers;
+    BindlessSamplerMap bindless_samplers;
 };
+
 } // namespace VideoCommon::Shader