vkd3d-shader/ir: Check that TEMP registers have consistent dimensions.

Approved-by: Giovanni Mascellani (@giomasce)
Approved-by: Henri Verbeet (@hverbeet)
Approved-by: Alexandre Julliard (@julliard)
Merge-Request: https://gitlab.winehq.org/wine/vkd3d/-/merge_requests/550

libs/vkd3d-shader/d3dbc.c

@@ -1360,18 +1360,21 @@ int vkd3d_shader_sm1_parser_create(const struct vkd3d_shader_compile_info *compi
         sm1->p.shader_desc.flat_constant_count[i].external = get_external_constant_count(sm1, i);
 
     if (!sm1->p.failed)
-        vsir_validate(&sm1->p);
+        ret = vsir_validate(&sm1->p);
 
-    if (sm1->p.failed)
+    if (sm1->p.failed && ret >= 0)
+        ret = VKD3D_ERROR_INVALID_SHADER;
+
+    if (ret < 0)
     {
         WARN("Failed to parse shader.\n");
         shader_sm1_destroy(&sm1->p);
-        return VKD3D_ERROR_INVALID_SHADER;
+        return ret;
     }
 
     *parser = &sm1->p;
 
-    return VKD3D_OK;
+    return ret;
 }
 
 bool hlsl_sm1_register_from_semantic(struct hlsl_ctx *ctx, const struct hlsl_semantic *semantic,

libs/vkd3d-shader/dxil.c

@@ -6282,7 +6282,7 @@ int vkd3d_shader_sm6_parser_create(const struct vkd3d_shader_compile_info *compi
     vkd3d_free(byte_code);
 
     if (!sm6->p.failed && ret >= 0)
-        vsir_validate(&sm6->p);
+        ret = vsir_validate(&sm6->p);
 
     if (sm6->p.failed && ret >= 0)
         ret = VKD3D_ERROR_INVALID_SHADER;

libs/vkd3d-shader/ir.c

@@ -1511,7 +1511,7 @@ enum vkd3d_result vkd3d_shader_normalise(struct vkd3d_shader_parser *parser,
         vkd3d_shader_trace(instructions, &parser->shader_version);
 
     if (result >= 0 && !parser->failed)
-        vsir_validate(parser);
+        result = vsir_validate(parser);
 
     if (result >= 0 && parser->failed)
         result = VKD3D_ERROR_INVALID_SHADER;
@@ -1527,6 +1527,12 @@ struct validation_context
     unsigned int temp_count;
     enum vkd3d_shader_opcode phase;
 
+    struct validation_context_temp_data
+    {
+        enum vsir_dimension dimension;
+        size_t first_seen;
+    } *temps;
+
     enum vkd3d_shader_opcode *blocks;
     size_t depth;
     size_t blocks_capacity;
@@ -1592,6 +1598,9 @@ static void vsir_validate_register(struct validation_context *ctx,
     switch (reg->type)
     {
         case VKD3DSPR_TEMP:
+        {
+            struct validation_context_temp_data *data;
+
             if (reg->idx_count != 1)
             {
                 validator_error(ctx, VKD3D_SHADER_ERROR_VSIR_INVALID_INDEX_COUNT, "Invalid index count %u for a TEMP register.",
@@ -1603,9 +1612,41 @@ static void vsir_validate_register(struct validation_context *ctx,
                 validator_error(ctx, VKD3D_SHADER_ERROR_VSIR_INVALID_INDEX, "Non-NULL relative address for a TEMP register.");
 
             if (reg->idx[0].offset >= temp_count)
+            {
                 validator_error(ctx, VKD3D_SHADER_ERROR_VSIR_INVALID_INDEX, "TEMP register index %u exceeds the maximum count %u.",
                         reg->idx[0].offset, temp_count);
+                break;
+            }
+
+            /* parser->shader_desc.temp_count might be smaller then
+             * temp_count if the parser made a mistake; we still don't
+             * want to overflow the array. */
+            if (reg->idx[0].offset >= ctx->parser->shader_desc.temp_count)
+                break;
+            data = &ctx->temps[reg->idx[0].offset];
+
+            if (reg->dimension == VSIR_DIMENSION_NONE)
+            {
+                validator_error(ctx, VKD3D_SHADER_ERROR_VSIR_INVALID_DIMENSION, "Invalid dimension NONE for a TEMP register.");
+                break;
+            }
+
+            /* TEMP registers can be scalar or vec4, provided that
+             * each individual register always appears with the same
+             * dimension. */
+            if (data->dimension == VSIR_DIMENSION_NONE)
+            {
+                data->dimension = reg->dimension;
+                data->first_seen = ctx->instruction_idx;
+            }
+            else if (data->dimension != reg->dimension)
+            {
+                validator_error(ctx, VKD3D_SHADER_ERROR_VSIR_INVALID_DIMENSION, "Invalid dimension %#x for a TEMP register: "
+                        "it has already been seen with dimension %#x at instruction %zu.",
+                        reg->dimension, data->dimension, data->first_seen);
+            }
             break;
+        }
 
         case VKD3DSPR_SSA:
             if (reg->idx_count != 1)
@@ -1883,7 +1924,7 @@ static void vsir_validate_instruction(struct validation_context *ctx)
     }
 }
 
-void vsir_validate(struct vkd3d_shader_parser *parser)
+enum vkd3d_result vsir_validate(struct vkd3d_shader_parser *parser)
 {
     struct validation_context ctx =
     {
@@ -1892,7 +1933,10 @@ void vsir_validate(struct vkd3d_shader_parser *parser)
     };
 
     if (!(parser->config_flags & VKD3D_SHADER_CONFIG_FLAG_FORCE_VALIDATION))
-        return;
+        return VKD3D_OK;
+
+    if (!(ctx.temps = vkd3d_calloc(parser->shader_desc.temp_count, sizeof(*ctx.temps))))
+        goto fail;
 
     for (ctx.instruction_idx = 0; ctx.instruction_idx < parser->instructions.count; ++ctx.instruction_idx)
         vsir_validate_instruction(&ctx);
@@ -1901,4 +1945,13 @@ void vsir_validate(struct vkd3d_shader_parser *parser)
         validator_error(&ctx, VKD3D_SHADER_ERROR_VSIR_INVALID_INSTRUCTION_NESTING, "%zu nested blocks were not closed.", ctx.depth);
 
     vkd3d_free(ctx.blocks);
+    vkd3d_free(ctx.temps);
+
+    return VKD3D_OK;
+
+fail:
+    vkd3d_free(ctx.blocks);
+    vkd3d_free(ctx.temps);
+
+    return VKD3D_ERROR_OUT_OF_MEMORY;
 }

libs/vkd3d-shader/vkd3d_shader_private.h

@@ -1450,7 +1450,7 @@ int preproc_lexer_parse(const struct vkd3d_shader_compile_info *compile_info,
 int hlsl_compile_shader(const struct vkd3d_shader_code *hlsl, const struct vkd3d_shader_compile_info *compile_info,
         struct vkd3d_shader_code *out, struct vkd3d_shader_message_context *message_context);
 
-void vsir_validate(struct vkd3d_shader_parser *parser);
+enum vkd3d_result vsir_validate(struct vkd3d_shader_parser *parser);
 
 static inline enum vkd3d_shader_component_type vkd3d_component_type_from_data_type(
         enum vkd3d_data_type data_type)