12 files changed, 183 insertions, 13 deletions

diff --git a/src/common/CMakeLists.txt b/src/common/CMakeLists.txt
index 5b51fcafa..9c6f1c07c 100644
--- a/src/common/CMakeLists.txt
+++ b/src/common/CMakeLists.txt
@@ -74,10 +74,12 @@ add_custom_command(OUTPUT scm_rev.cpp
       "${VIDEO_CORE}/shader/decode/xmad.cpp"
       "${VIDEO_CORE}/shader/ast.cpp"
       "${VIDEO_CORE}/shader/ast.h"
-      "${VIDEO_CORE}/shader/control_flow.cpp"
-      "${VIDEO_CORE}/shader/control_flow.h"
       "${VIDEO_CORE}/shader/compiler_settings.cpp"
       "${VIDEO_CORE}/shader/compiler_settings.h"
+      "${VIDEO_CORE}/shader/const_buffer_locker.cpp"
+      "${VIDEO_CORE}/shader/const_buffer_locker.h"
+      "${VIDEO_CORE}/shader/control_flow.cpp"
+      "${VIDEO_CORE}/shader/control_flow.h"
       "${VIDEO_CORE}/shader/decode.cpp"
       "${VIDEO_CORE}/shader/expr.cpp"
       "${VIDEO_CORE}/shader/expr.h"
diff --git a/src/common/hash.h b/src/common/hash.h
index 40194d1ee..c939709bc 100644
--- a/src/common/hash.h
+++ b/src/common/hash.h
@@ -6,6 +6,8 @@
 
 #include <cstddef>
 #include <cstring>
+#include <utility>
+#include <boost/functional/hash.hpp>
 #include "common/cityhash.h"
 #include "common/common_types.h"
 
@@ -68,4 +70,13 @@ struct HashableStruct {
     }
 };
 
+struct PairHash {
+    template <class T1, class T2>
+    std::size_t operator()(const std::pair<T1, T2>& pair) const {
+        std::size_t seed = std::hash<T1>()(pair.first);
+        boost::hash_combine(seed, std::hash<T2>()(pair.second));
+        return seed;
+    }
+};
+
 } // namespace Common
diff --git a/src/video_core/CMakeLists.txt b/src/video_core/CMakeLists.txt
index eaa694ff8..cb6eda1b8 100644
--- a/src/video_core/CMakeLists.txt
+++ b/src/video_core/CMakeLists.txt
@@ -6,6 +6,7 @@ add_library(video_core STATIC
     dma_pusher.h
     debug_utils/debug_utils.cpp
     debug_utils/debug_utils.h
+    engines/const_buffer_engine_interface.h
     engines/const_buffer_info.h
     engines/engine_upload.cpp
     engines/engine_upload.h
@@ -107,10 +108,12 @@ add_library(video_core STATIC
     shader/decode/other.cpp
     shader/ast.cpp
     shader/ast.h
-    shader/control_flow.cpp
-    shader/control_flow.h
     shader/compiler_settings.cpp
     shader/compiler_settings.h
+    shader/const_buffer_locker.cpp
+    shader/const_buffer_locker.h
+    shader/control_flow.cpp
+    shader/control_flow.h
     shader/decode.cpp
     shader/expr.cpp
     shader/expr.h
diff --git a/src/video_core/engines/const_buffer_engine_interface.h b/src/video_core/engines/const_buffer_engine_interface.h
new file mode 100644
index 000000000..cc41a9cac
--- /dev/null
+++ b/src/video_core/engines/const_buffer_engine_interface.h
@@ -0,0 +1,26 @@
+// Copyright 2019 yuzu Emulator Project
+// Licensed under GPLv2 or any later version
+// Refer to the license.txt file included.
+
+#pragma once
+
+#include "common/common_types.h"
+
+namespace Tegra::Engines {
+
+enum class ShaderType : u32 {
+    Vertex = 0,
+    TesselationControl = 1,
+    TesselationEval = 2,
+    Geometry = 3,
+    Fragment = 4,
+    Compute = 5,
+};
+
+class ConstBufferEngineInterface {
+public:
+    virtual ~ConstBufferEngineInterface() {}
+    virtual u32 AccessConstBuffer32(ShaderType stage, u64 const_buffer, u64 offset) const = 0;
+};
+
+}
diff --git a/src/video_core/engines/kepler_compute.cpp b/src/video_core/engines/kepler_compute.cpp
index 63d449135..ba97c2894 100644
--- a/src/video_core/engines/kepler_compute.cpp
+++ b/src/video_core/engines/kepler_compute.cpp
@@ -70,7 +70,8 @@ Texture::FullTextureInfo KeplerCompute::GetTextureInfo(const Texture::TextureHan
                                     GetTSCEntry(tex_handle.tsc_id)};
 }
 
-u32 KeplerCompute::AccessConstBuffer32(u64 const_buffer, u64 offset) const {
+u32 KeplerCompute::AccessConstBuffer32(ShaderType stage, u64 const_buffer, u64 offset) const {
+    ASSERT(stage == ShaderType::Compute);
     const auto& buffer = launch_description.const_buffer_config[const_buffer];
     u32 result;
     std::memcpy(&result, memory_manager.GetPointer(buffer.Address() + offset), sizeof(u32));
diff --git a/src/video_core/engines/kepler_compute.h b/src/video_core/engines/kepler_compute.h
index 90cf650d2..d7e0dfcd6 100644
--- a/src/video_core/engines/kepler_compute.h
+++ b/src/video_core/engines/kepler_compute.h
@@ -11,6 +11,7 @@
 #include "common/common_funcs.h"
 #include "common/common_types.h"
 #include "video_core/engines/engine_upload.h"
+#include "video_core/engines/const_buffer_engine_interface.h"
 #include "video_core/gpu.h"
 #include "video_core/textures/texture.h"
 
@@ -37,7 +38,7 @@ namespace Tegra::Engines {
 #define KEPLER_COMPUTE_REG_INDEX(field_name)                                                       \
     (offsetof(Tegra::Engines::KeplerCompute::Regs, field_name) / sizeof(u32))
 
-class KeplerCompute final {
+class KeplerCompute final : public ConstBufferEngineInterface  {
 public:
     explicit KeplerCompute(Core::System& system, VideoCore::RasterizerInterface& rasterizer,
                            MemoryManager& memory_manager);
@@ -201,7 +202,7 @@ public:
     Texture::FullTextureInfo GetTextureInfo(const Texture::TextureHandle tex_handle,
                                             std::size_t offset) const;
 
-    u32 AccessConstBuffer32(u64 const_buffer, u64 offset) const;
+    u32 AccessConstBuffer32(ShaderType stage, u64 const_buffer, u64 offset) const override;
 
 private:
     Core::System& system;
diff --git a/src/video_core/engines/maxwell_3d.cpp b/src/video_core/engines/maxwell_3d.cpp
index 59976943a..92e38b071 100644
--- a/src/video_core/engines/maxwell_3d.cpp
+++ b/src/video_core/engines/maxwell_3d.cpp
@@ -847,7 +847,8 @@ void Maxwell3D::ProcessClearBuffers() {
     rasterizer.Clear();
 }
 
-u32 Maxwell3D::AccessConstBuffer32(Regs::ShaderStage stage, u64 const_buffer, u64 offset) const {
+u32 Maxwell3D::AccessConstBuffer32(ShaderType stage, u64 const_buffer, u64 offset) const {
+    ASSERT(stage != ShaderType::Compute);
     const auto& shader_stage = state.shader_stages[static_cast<std::size_t>(stage)];
     const auto& buffer = shader_stage.const_buffers[const_buffer];
     u32 result;
diff --git a/src/video_core/engines/maxwell_3d.h b/src/video_core/engines/maxwell_3d.h
index e3f1047d5..04d02d208 100644
--- a/src/video_core/engines/maxwell_3d.h
+++ b/src/video_core/engines/maxwell_3d.h
@@ -16,6 +16,7 @@
 #include "common/common_types.h"
 #include "common/math_util.h"
 #include "video_core/engines/const_buffer_info.h"
+#include "video_core/engines/const_buffer_engine_interface.h"
 #include "video_core/engines/engine_upload.h"
 #include "video_core/gpu.h"
 #include "video_core/macro_interpreter.h"
@@ -44,7 +45,7 @@ namespace Tegra::Engines {
 #define MAXWELL3D_REG_INDEX(field_name)                                                            \
     (offsetof(Tegra::Engines::Maxwell3D::Regs, field_name) / sizeof(u32))
 
-class Maxwell3D final {
+class Maxwell3D final : public ConstBufferEngineInterface {
 public:
     explicit Maxwell3D(Core::System& system, VideoCore::RasterizerInterface& rasterizer,
                        MemoryManager& memory_manager);
@@ -1257,7 +1258,7 @@ public:
     /// Returns the texture information for a specific texture in a specific shader stage.
     Texture::FullTextureInfo GetStageTexture(Regs::ShaderStage stage, std::size_t offset) const;
 
-    u32 AccessConstBuffer32(Regs::ShaderStage stage, u64 const_buffer, u64 offset) const;
+    u32 AccessConstBuffer32(ShaderType stage, u64 const_buffer, u64 offset) const override;
 
     /// Memory for macro code - it's undetermined how big this is, however 1MB is much larger than
     /// we've seen used.
diff --git a/src/video_core/renderer_opengl/gl_rasterizer.cpp b/src/video_core/renderer_opengl/gl_rasterizer.cpp
index cbcf81414..10114909b 100644
--- a/src/video_core/renderer_opengl/gl_rasterizer.cpp
+++ b/src/video_core/renderer_opengl/gl_rasterizer.cpp
@@ -975,7 +975,8 @@ TextureBufferUsage RasterizerOpenGL::SetupDrawTextures(Maxwell::ShaderStage stag
             }
             const auto cbuf = entry.GetBindlessCBuf();
             Tegra::Texture::TextureHandle tex_handle;
-            tex_handle.raw = maxwell3d.AccessConstBuffer32(stage, cbuf.first, cbuf.second);
+            Tegra::Engines::ShaderType shader_type = static_cast<Tegra::Engines::ShaderType>(stage);
+            tex_handle.raw = maxwell3d.AccessConstBuffer32(shader_type, cbuf.first, cbuf.second);
             return maxwell3d.GetTextureInfo(tex_handle, entry.GetOffset());
         }();
 
@@ -1005,7 +1006,7 @@ TextureBufferUsage RasterizerOpenGL::SetupComputeTextures(const Shader& kernel)
             }
             const auto cbuf = entry.GetBindlessCBuf();
             Tegra::Texture::TextureHandle tex_handle;
-            tex_handle.raw = compute.AccessConstBuffer32(cbuf.first, cbuf.second);
+            tex_handle.raw = compute.AccessConstBuffer32(Tegra::Engines::ShaderType::Compute, cbuf.first, cbuf.second);
             return compute.GetTextureInfo(tex_handle, entry.GetOffset());
         }();
 
@@ -1050,7 +1051,7 @@ void RasterizerOpenGL::SetupComputeImages(const Shader& shader) {
             }
             const auto cbuf = entry.GetBindlessCBuf();
             Tegra::Texture::TextureHandle tex_handle;
-            tex_handle.raw = compute.AccessConstBuffer32(cbuf.first, cbuf.second);
+            tex_handle.raw = compute.AccessConstBuffer32(Tegra::Engines::ShaderType::Compute, cbuf.first, cbuf.second);
             return compute.GetTextureInfo(tex_handle, entry.GetOffset()).tic;
         }();
         SetupImage(bindpoint, tic, entry);
diff --git a/src/video_core/shader/const_buffer_locker.cpp b/src/video_core/shader/const_buffer_locker.cpp
new file mode 100644
index 000000000..6a9e0ed5e
--- /dev/null
+++ b/src/video_core/shader/const_buffer_locker.cpp
@@ -0,0 +1,72 @@
+// Copyright 2019 yuzu Emulator Project
+// Licensed under GPLv2 or any later version
+// Refer to the license.txt file included.
+
+#pragma once
+
+#include "common/assert.h"
+#include "common/common_types.h"
+#include "video_core/engines/maxwell_3d.h"
+#include "video_core/shader/const_buffer_locker.h"
+
+namespace VideoCommon::Shader {
+
+ConstBufferLocker::ConstBufferLocker(Tegra::Engines::ShaderType shader_stage)
+    : engine{nullptr}, shader_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_;
+}
+
+std::optional<u32> ConstBufferLocker::ObtainKey(u32 buffer, u32 offset) {
+    const std::pair<u32, u32> key = {buffer, offset};
+    const auto iter = keys.find(key);
+    if (iter != keys.end()) {
+        return {iter->second};
+    }
+    if (!IsEngineSet()) {
+        return {};
+    }
+    const u32 value = engine->AccessConstBuffer32(shader_stage, buffer, offset);
+    keys.emplace(key, value);
+    return {value};
+}
+
+void ConstBufferLocker::InsertKey(u32 buffer, u32 offset, u32 value) {
+    const std::pair<u32, u32> key = {buffer, offset};
+    keys[key] = value;
+}
+
+u32 ConstBufferLocker::NumKeys() const {
+    return keys.size();
+}
+
+const std::unordered_map<std::pair<u32, u32>, u32, Common::PairHash>&
+ConstBufferLocker::AccessKeys() const {
+    return keys;
+}
+
+bool ConstBufferLocker::AreKeysConsistant() const {
+    if (!IsEngineSet()) {
+        return false;
+    }
+    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;
+        }
+    }
+    return true;
+}
+
+} // namespace VideoCommon::Shader
diff --git a/src/video_core/shader/const_buffer_locker.h b/src/video_core/shader/const_buffer_locker.h
new file mode 100644
index 000000000..39e62584d
--- /dev/null
+++ b/src/video_core/shader/const_buffer_locker.h
@@ -0,0 +1,50 @@
+// Copyright 2019 yuzu Emulator Project
+// Licensed under GPLv2 or any later version
+// Refer to the license.txt file included.
+
+#pragma once
+
+#include <unordered_map>
+#include "common/common_types.h"
+#include "common/hash.h"
+#include "video_core/engines/const_buffer_engine_interface.h"
+
+namespace VideoCommon::Shader {
+
+class ConstBufferLocker {
+public:
+    explicit ConstBufferLocker(Tegra::Engines::ShaderType shader_stage);
+
+    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.
+    std::optional<u32> ObtainKey(u32 buffer, u32 offset);
+
+    // Manually inserts a key.
+    void InsertKey(u32 buffer, u32 offset, u32 value);
+
+    // Retrieves the number of keys registered.
+    u32 NumKeys() const;
+
+    // Gives an accessor to the key's database.
+    const std::unordered_map<std::pair<u32, u32>, u32, Common::PairHash>& AccessKeys() const;
+
+    // Checks keys against maxwell3d's current const buffers. Returns true if they
+    // are the same value, false otherwise;
+    bool AreKeysConsistant() const;
+
+private:
+    Tegra::Engines::ConstBufferEngineInterface* engine;
+    Tegra::Engines::ShaderType shader_stage;
+    std::unordered_map<std::pair<u32, u32>, u32, Common::PairHash> keys{};
+};
+} // namespace VideoCommon::Shader
diff --git a/src/video_core/shader/shader_ir.h b/src/video_core/shader/shader_ir.h
index 91cd0a534..68818643c 100644
--- a/src/video_core/shader/shader_ir.h
+++ b/src/video_core/shader/shader_ir.h
@@ -17,6 +17,7 @@
 #include "video_core/engines/shader_header.h"
 #include "video_core/shader/ast.h"
 #include "video_core/shader/compiler_settings.h"
+#include "video_core/shader/const_buffer_locker.h"
 #include "video_core/shader/node.h"
 
 namespace VideoCommon::Shader {