6 files changed, 40 insertions, 44 deletions

diff --git a/src/video_core/command_processor.cpp b/src/video_core/command_processor.cpp
index 66d19cba0..c3872d06c 100644
--- a/src/video_core/command_processor.cpp
+++ b/src/video_core/command_processor.cpp
@@ -143,7 +143,7 @@ static void WritePicaReg(u32 id, u32 value, u32 mask) {
                     immediate_attribute_id = 0;
 
                     auto* shader_engine = Shader::GetEngine();
-                    shader_engine->SetupBatch(&g_state.vs);
+                    shader_engine->SetupBatch(g_state.vs);
 
                     // Send to vertex shader
                     if (g_debug_context)
@@ -151,7 +151,7 @@ static void WritePicaReg(u32 id, u32 value, u32 mask) {
                                                  static_cast<void*>(&immediate_input));
                     Shader::UnitState shader_unit;
                     shader_unit.LoadInputVertex(immediate_input, regs.vs.num_input_attributes + 1);
-                    shader_engine->Run(shader_unit, regs.vs.main_offset);
+                    shader_engine->Run(g_state.vs, shader_unit, regs.vs.main_offset);
                     auto output_vertex = Shader::OutputVertex::FromRegisters(
                         shader_unit.registers.output, regs, regs.vs.output_mask);
 
@@ -248,7 +248,7 @@ static void WritePicaReg(u32 id, u32 value, u32 mask) {
         auto* shader_engine = Shader::GetEngine();
         Shader::UnitState shader_unit;
 
-        shader_engine->SetupBatch(&g_state.vs);
+        shader_engine->SetupBatch(g_state.vs);
 
         for (unsigned int index = 0; index < regs.num_vertices; ++index) {
             // Indexed rendering doesn't use the start offset
@@ -288,7 +288,7 @@ static void WritePicaReg(u32 id, u32 value, u32 mask) {
                     g_debug_context->OnEvent(DebugContext::Event::VertexShaderInvocation,
                                              (void*)&input);
                 shader_unit.LoadInputVertex(input, loader.GetNumTotalAttributes());
-                shader_engine->Run(shader_unit, regs.vs.main_offset);
+                shader_engine->Run(g_state.vs, shader_unit, regs.vs.main_offset);
 
                 // Retrieve vertex from register data
                 output_vertex = Shader::OutputVertex::FromRegisters(shader_unit.registers.output,
diff --git a/src/video_core/shader/shader.h b/src/video_core/shader/shader.h
index 7d51d0044..f26d2ba4f 100644
--- a/src/video_core/shader/shader.h
+++ b/src/video_core/shader/shader.h
@@ -167,6 +167,12 @@ struct ShaderSetup {
 
     std::array<u32, 1024> program_code;
     std::array<u32, 1024> swizzle_data;
+
+    /// Data private to ShaderEngines
+    struct EngineData {
+        /// Used by the JIT, points to a compiled shader object.
+        const void* cached_shader = nullptr;
+    } engine_data;
 };
 
 class ShaderEngine {
@@ -177,13 +183,16 @@ public:
      * Performs any shader unit setup that only needs to happen once per shader (as opposed to once
      * per vertex, which would happen within the `Run` function).
      */
-    virtual void SetupBatch(const ShaderSetup* setup) = 0;
+    virtual void SetupBatch(ShaderSetup& setup) = 0;
 
     /**
-     * Runs the currently setup shader
-     * @param state Shader unit state, must be setup per shader and per shader unit
+     * Runs the currently setup shader.
+     *
+     * @param setup Shader engine state, must be setup with SetupBatch on each shader change.
+     * @param state Shader unit state, must be setup with input data before each shader invocation.
      */
-    virtual void Run(UnitState& state, unsigned int entry_point) const = 0;
+    virtual void Run(const ShaderSetup& setup, UnitState& state,
+                     unsigned int entry_point) const = 0;
 };
 
 // TODO(yuriks): Remove and make it non-global state somewhere
diff --git a/src/video_core/shader/shader_interpreter.cpp b/src/video_core/shader/shader_interpreter.cpp
index a6197c10a..e44abbf1d 100644
--- a/src/video_core/shader/shader_interpreter.cpp
+++ b/src/video_core/shader/shader_interpreter.cpp
@@ -652,25 +652,23 @@ static void RunInterpreter(const ShaderSetup& setup, UnitState& state, DebugData
     }
 }
 
-void InterpreterEngine::SetupBatch(const ShaderSetup* setup_) {
-    setup = setup_;
-}
+void InterpreterEngine::SetupBatch(ShaderSetup& setup) {}
 
 MICROPROFILE_DECLARE(GPU_Shader);
 
-void InterpreterEngine::Run(UnitState& state, unsigned int entry_point) const {
-    ASSERT(setup != nullptr);
+void InterpreterEngine::Run(const ShaderSetup& setup, UnitState& state,
+                            unsigned int entry_point) const {
     ASSERT(entry_point < 1024);
 
     MICROPROFILE_SCOPE(GPU_Shader);
 
     DebugData<false> dummy_debug_data;
-    RunInterpreter(*setup, state, dummy_debug_data, entry_point);
+    RunInterpreter(setup, state, dummy_debug_data, entry_point);
 }
 
-DebugData<true> InterpreterEngine::ProduceDebugInfo(const InputVertex& input, int num_attributes,
+DebugData<true> InterpreterEngine::ProduceDebugInfo(const ShaderSetup& setup,
+                                                    const InputVertex& input, int num_attributes,
                                                     unsigned int entry_point) const {
-    ASSERT(setup != nullptr);
     ASSERT(entry_point < 1024);
 
     UnitState state;
@@ -679,7 +677,7 @@ DebugData<true> InterpreterEngine::ProduceDebugInfo(const InputVertex& input, in
     // Setup input register table
     boost::fill(state.registers.input, Math::Vec4<float24>::AssignToAll(float24::Zero()));
     state.LoadInputVertex(input, num_attributes);
-    RunInterpreter(*setup, state, debug_data, entry_point);
+    RunInterpreter(setup, state, debug_data, entry_point);
     return debug_data;
 }
 
diff --git a/src/video_core/shader/shader_interpreter.h b/src/video_core/shader/shader_interpreter.h
index c3691da70..7f94d405f 100644
--- a/src/video_core/shader/shader_interpreter.h
+++ b/src/video_core/shader/shader_interpreter.h
@@ -13,8 +13,8 @@ namespace Shader {
 
 class InterpreterEngine final : public ShaderEngine {
 public:
-    void SetupBatch(const ShaderSetup* setup) override;
-    void Run(UnitState& state, unsigned int entry_point) const override;
+    void SetupBatch(ShaderSetup& setup) override;
+    void Run(const ShaderSetup& setup, UnitState& state, unsigned int entry_point) const override;
 
     /**
      * Produce debug information based on the given shader and input vertex
@@ -23,11 +23,8 @@ public:
      * @param config Configuration object for the shader pipeline
      * @return Debug information for this shader with regards to the given vertex
      */
-    DebugData<true> ProduceDebugInfo(const InputVertex& input, int num_attributes,
-                                     unsigned int entry_point) const;
-
-private:
-    const ShaderSetup* setup = nullptr;
+    DebugData<true> ProduceDebugInfo(const ShaderSetup& setup, const InputVertex& input,
+                                     int num_attributes, unsigned int entry_point) const;
 };
 
 } // namespace
diff --git a/src/video_core/shader/shader_jit_x64.cpp b/src/video_core/shader/shader_jit_x64.cpp
index 755ae119f..15c1d60b5 100644
--- a/src/video_core/shader/shader_jit_x64.cpp
+++ b/src/video_core/shader/shader_jit_x64.cpp
@@ -14,37 +14,32 @@ namespace Shader {
 JitX64Engine::JitX64Engine() = default;
 JitX64Engine::~JitX64Engine() = default;
 
-void JitX64Engine::SetupBatch(const ShaderSetup* setup_) {
-    cached_shader = nullptr;
-    setup = setup_;
-    if (setup == nullptr)
-        return;
-
-    u64 code_hash = Common::ComputeHash64(&setup->program_code, sizeof(setup->program_code));
-    u64 swizzle_hash = Common::ComputeHash64(&setup->swizzle_data, sizeof(setup->swizzle_data));
+void JitX64Engine::SetupBatch(ShaderSetup& setup) {
+    u64 code_hash = Common::ComputeHash64(&setup.program_code, sizeof(setup.program_code));
+    u64 swizzle_hash = Common::ComputeHash64(&setup.swizzle_data, sizeof(setup.swizzle_data));
 
     u64 cache_key = code_hash ^ swizzle_hash;
     auto iter = cache.find(cache_key);
     if (iter != cache.end()) {
-        cached_shader = iter->second.get();
+        setup.engine_data.cached_shader = iter->second.get();
     } else {
         auto shader = std::make_unique<JitShader>();
-        shader->Compile(&setup->program_code, &setup->swizzle_data);
-        cached_shader = shader.get();
+        shader->Compile(&setup.program_code, &setup.swizzle_data);
+        setup.engine_data.cached_shader = shader.get();
         cache.emplace_hint(iter, cache_key, std::move(shader));
     }
 }
 
 MICROPROFILE_DECLARE(GPU_Shader);
 
-void JitX64Engine::Run(UnitState& state, unsigned int entry_point) const {
-    ASSERT(setup != nullptr);
-    ASSERT(cached_shader != nullptr);
+void JitX64Engine::Run(const ShaderSetup& setup, UnitState& state, unsigned int entry_point) const {
+    ASSERT(setup.engine_data.cached_shader != nullptr);
     ASSERT(entry_point < 1024);
 
     MICROPROFILE_SCOPE(GPU_Shader);
 
-    cached_shader->Run(*setup, state, entry_point);
+    const JitShader* shader = static_cast<const JitShader*>(setup.engine_data.cached_shader);
+    shader->Run(setup, state, entry_point);
 }
 
 } // namespace Shader
diff --git a/src/video_core/shader/shader_jit_x64.h b/src/video_core/shader/shader_jit_x64.h
index b26044477..bd30f51e2 100644
--- a/src/video_core/shader/shader_jit_x64.h
+++ b/src/video_core/shader/shader_jit_x64.h
@@ -19,14 +19,11 @@ public:
     JitX64Engine();
     ~JitX64Engine() override;
 
-    void SetupBatch(const ShaderSetup* setup) override;
-    void Run(UnitState& state, unsigned int entry_point) const override;
+    void SetupBatch(ShaderSetup& setup) override;
+    void Run(const ShaderSetup& setup, UnitState& state, unsigned int entry_point) const override;
 
 private:
-    const ShaderSetup* setup = nullptr;
-
     std::unordered_map<u64, std::unique_ptr<JitShader>> cache;
-    const JitShader* cached_shader = nullptr;
 };
 
 } // namespace Shader