3 files changed, 24 insertions, 24 deletions

diff --git a/src/video_core/command_processor.cpp b/src/video_core/command_processor.cpp
index fef0b4ceb..4955ff9f9 100644
--- a/src/video_core/command_processor.cpp
+++ b/src/video_core/command_processor.cpp
@@ -151,10 +151,11 @@ static void WritePicaReg(u32 id, u32 value, u32 mask) {
                         g_debug_context->OnEvent(DebugContext::Event::VertexShaderInvocation,
                                                  static_cast<void*>(&immediate_input));
                     Shader::UnitState shader_unit;
+                    Shader::AttributeBuffer output{};
+
                     shader_unit.LoadInput(regs.vs, immediate_input);
                     shader_engine->Run(g_state.vs, shader_unit);
-                    auto output_vertex = Shader::OutputVertex::FromRegisters(
-                        shader_unit.registers.output, regs, regs.vs.output_mask);
+                    shader_unit.WriteOutput(regs.vs, output);
 
                     // Send to renderer
                     using Pica::Shader::OutputVertex;
@@ -163,7 +164,8 @@ static void WritePicaReg(u32 id, u32 value, u32 mask) {
                         VideoCore::g_renderer->Rasterizer()->AddTriangle(v0, v1, v2);
                     };
 
-                    g_state.primitive_assembler.SubmitVertex(output_vertex, AddTriangle);
+                    g_state.primitive_assembler.SubmitVertex(
+                        Shader::OutputVertex::FromAttributeBuffer(regs, output), AddTriangle);
                 }
             }
         }
@@ -281,7 +283,7 @@ static void WritePicaReg(u32 id, u32 value, u32 mask) {
 
             if (!vertex_cache_hit) {
                 // Initialize data for the current vertex
-                Shader::AttributeBuffer input;
+                Shader::AttributeBuffer input, output{};
                 loader.LoadVertex(base_address, index, vertex, input, memory_accesses);
 
                 // Send to vertex shader
@@ -290,10 +292,10 @@ static void WritePicaReg(u32 id, u32 value, u32 mask) {
                                              (void*)&input);
                 shader_unit.LoadInput(regs.vs, input);
                 shader_engine->Run(g_state.vs, shader_unit);
+                shader_unit.WriteOutput(regs.vs, output);
 
                 // Retrieve vertex from register data
-                output_vertex = Shader::OutputVertex::FromRegisters(shader_unit.registers.output,
-                                                                    regs, regs.vs.output_mask);
+                output_vertex = Shader::OutputVertex::FromAttributeBuffer(regs, output);
 
                 if (is_indexed) {
                     vertex_cache[vertex_cache_pos] = output_vertex;
diff --git a/src/video_core/shader/shader.cpp b/src/video_core/shader/shader.cpp
index dbad167e9..99a22c2dd 100644
--- a/src/video_core/shader/shader.cpp
+++ b/src/video_core/shader/shader.cpp
@@ -4,6 +4,7 @@
 
 #include <cmath>
 #include <cstring>
+#include "common/bit_set.h"
 #include "common/logging/log.h"
 #include "common/microprofile.h"
 #include "video_core/pica.h"
@@ -19,22 +20,13 @@ namespace Pica {
 
 namespace Shader {
 
-OutputVertex OutputVertex::FromRegisters(Math::Vec4<float24> output_regs[16], const Regs& regs,
-                                         u32 output_mask) {
+OutputVertex OutputVertex::FromAttributeBuffer(const Regs& regs, AttributeBuffer& input) {
     // Setup output data
     OutputVertex ret;
-    // TODO(neobrain): Under some circumstances, up to 16 attributes may be output. We need to
-    // figure out what those circumstances are and enable the remaining outputs then.
-    unsigned index = 0;
-    for (unsigned i = 0; i < 7; ++i) {
 
-        if (index >= regs.vs_output_total)
-            break;
-
-        if ((output_mask & (1 << i)) == 0)
-            continue;
-
-        const auto& output_register_map = regs.vs_output_attributes[index];
+    unsigned int num_attributes = regs.vs_output_total;
+    for (unsigned int i = 0; i < num_attributes; ++i) {
+        const auto& output_register_map = regs.vs_output_attributes[i];
 
         u32 semantics[4] = {output_register_map.map_x, output_register_map.map_y,
                             output_register_map.map_z, output_register_map.map_w};
@@ -42,15 +34,13 @@ OutputVertex OutputVertex::FromRegisters(Math::Vec4<float24> output_regs[16], co
         for (unsigned comp = 0; comp < 4; ++comp) {
             float24* out = ((float24*)&ret) + semantics[comp];
             if (semantics[comp] != Regs::VSOutputAttributes::INVALID) {
-                *out = output_regs[i][comp];
+                *out = input.attr[i][comp];
             } else {
                 // Zero output so that attributes which aren't output won't have denormals in them,
                 // which would slow us down later.
                 memset(out, 0, sizeof(*out));
             }
         }
-
-        index++;
     }
 
     // The hardware takes the absolute and saturates vertex colors like this, *before* doing
@@ -80,6 +70,13 @@ void UnitState::LoadInput(const Regs::ShaderConfig& config, const AttributeBuffe
     }
 }
 
+void UnitState::WriteOutput(const Regs::ShaderConfig& config, AttributeBuffer& output) {
+    unsigned int output_i = 0;
+    for (unsigned int reg : Common::BitSet<u32>(config.output_mask)) {
+        output.attr[output_i++] = registers.output[reg];
+    }
+}
+
 MICROPROFILE_DEFINE(GPU_Shader, "GPU", "Shader", MP_RGB(50, 50, 240));
 
 #ifdef ARCHITECTURE_x86_64
diff --git a/src/video_core/shader/shader.h b/src/video_core/shader/shader.h
index 43a8b848c..00bd723cf 100644
--- a/src/video_core/shader/shader.h
+++ b/src/video_core/shader/shader.h
@@ -74,8 +74,7 @@ struct OutputVertex {
         return ret;
     }
 
-    static OutputVertex FromRegisters(Math::Vec4<float24> output_regs[16], const Regs& regs,
-                                      u32 output_mask);
+    static OutputVertex FromAttributeBuffer(const Regs& regs, AttributeBuffer& output);
 };
 static_assert(std::is_pod<OutputVertex>::value, "Structure is not POD");
 static_assert(sizeof(OutputVertex) == 32 * sizeof(float), "OutputVertex has invalid size");
@@ -141,6 +140,8 @@ struct UnitState {
      * @param input Attribute buffer to load into the input registers.
      */
     void LoadInput(const Regs::ShaderConfig& config, const AttributeBuffer& input);
+
+    void WriteOutput(const Regs::ShaderConfig& config, AttributeBuffer& output);
 };
 
 struct ShaderSetup {