[PATCH v2 1/4] vkd3d-shader/tpf: Validate input/output register index counts.

[Conor McCarthy]

From: Conor McCarthy

---
 libs/vkd3d-shader/d3dbc.c                |  3 ++
 libs/vkd3d-shader/ir.c                   | 31 ++++++-------
 libs/vkd3d-shader/tpf.c                  | 55 ++++++++++++++++++++++++
 libs/vkd3d-shader/vkd3d_shader_private.h |  2 +
 4 files changed, 73 insertions(+), 18 deletions(-)

diff --git a/libs/vkd3d-shader/d3dbc.c b/libs/vkd3d-shader/d3dbc.c
index 823df6f6..d4f3ef66 100644
--- a/libs/vkd3d-shader/d3dbc.c
+++ b/libs/vkd3d-shader/d3dbc.c
@@ -465,6 +465,7 @@ static void shader_sm1_parse_src_param(uint32_t param, const struct vkd3d_shader
     src->reg.idx[1].rel_addr = NULL;
     src->reg.idx[2].offset = ~0u;
     src->reg.idx[2].rel_addr = NULL;
+    src->reg.idx_count = 1;
     src->swizzle = swizzle_from_sm1((param & VKD3D_SM1_SWIZZLE_MASK) >> VKD3D_SM1_SWIZZLE_SHIFT);
     src->modifiers = (param & VKD3D_SM1_SRC_MODIFIER_MASK) >> VKD3D_SM1_SRC_MODIFIER_SHIFT;
 }
@@ -483,6 +484,7 @@ static void shader_sm1_parse_dst_param(uint32_t param, const struct vkd3d_shader
     dst->reg.idx[1].rel_addr = NULL;
     dst->reg.idx[2].offset = ~0u;
     dst->reg.idx[2].rel_addr = NULL;
+    dst->reg.idx_count = 1;
     dst->write_mask = (param & VKD3D_SM1_WRITEMASK_MASK) >> VKD3D_SM1_WRITEMASK_SHIFT;
     dst->modifiers = (param & VKD3D_SM1_DST_MODIFIER_MASK) >> VKD3D_SM1_DST_MODIFIER_SHIFT;
     dst->shift = (param & VKD3D_SM1_DSTSHIFT_MASK) >> VKD3D_SM1_DSTSHIFT_SHIFT;
@@ -664,6 +666,7 @@ static void shader_sm1_read_immconst(struct vkd3d_shader_sm1_parser *sm1, const
     src_param->reg.idx[1].rel_addr = NULL;
     src_param->reg.idx[2].offset = ~0u;
     src_param->reg.idx[2].rel_addr = NULL;
+    src_param->reg.idx_count = 0;
     src_param->reg.immconst_type = type;
     memcpy(src_param->reg.u.immconst_uint, *ptr, count * sizeof(uint32_t));
     src_param->swizzle = VKD3D_SHADER_NO_SWIZZLE;
diff --git a/libs/vkd3d-shader/ir.c b/libs/vkd3d-shader/ir.c
index a3ab03f2..cce23afd 100644
--- a/libs/vkd3d-shader/ir.c
+++ b/libs/vkd3d-shader/ir.c
@@ -43,7 +43,7 @@ static void shader_register_eliminate_phase_addressing(struct vkd3d_shader_regis
 {
     unsigned int i;

-    for (i = 0; i < ARRAY_SIZE(reg->idx) && reg->idx[i].offset != ~0u; ++i)
+    for (i = 0; i < reg->idx_count; ++i)
     {
         if (reg->idx[i].rel_addr && shader_register_is_phase_instance_id(&reg->idx[i].rel_addr->reg))
         {
@@ -73,6 +73,7 @@ static void shader_instruction_eliminate_phase_instance_id(struct vkd3d_shader_i
             reg->idx[1].rel_addr = NULL;
             reg->idx[2].offset = ~0u;
             reg->idx[2].rel_addr = NULL;
+            reg->idx_count = 0;
             reg->immconst_type = VKD3D_IMMCONST_SCALAR;
             reg->u.immconst_uint[0] = instance_id;
             continue;
@@ -222,8 +223,8 @@ static enum vkd3d_result shader_normaliser_flatten_phases(struct vkd3d_shader_no
     return VKD3D_OK;
 }

-static void shader_register_init(struct vkd3d_shader_register *reg,
-        enum vkd3d_shader_register_type reg_type, enum vkd3d_data_type data_type)
+static void shader_register_init(struct vkd3d_shader_register *reg, enum vkd3d_shader_register_type reg_type,
+        enum vkd3d_data_type data_type, unsigned int idx_count)
 {
     reg->type = reg_type;
     reg->precision = VKD3D_SHADER_REGISTER_PRECISION_DEFAULT;
@@ -235,6 +236,7 @@ static void shader_register_init(struct vkd3d_shader_register *reg,
     reg->idx[1].rel_addr = NULL;
     reg->idx[2].offset = ~0u;
     reg->idx[2].rel_addr = NULL;
+    reg->idx_count = idx_count;
     reg->immconst_type = VKD3D_IMMCONST_SCALAR;
 }

@@ -286,7 +288,7 @@ static struct vkd3d_shader_src_param *shader_normaliser_create_outpointid_param(
     if (!(rel_addr = shader_src_param_allocator_get(&normaliser->instructions.src_params, 1)))
         return NULL;

-    shader_register_init(&rel_addr->reg, VKD3DSPR_OUTPOINTID, VKD3D_DATA_UINT);
+    shader_register_init(&rel_addr->reg, VKD3DSPR_OUTPOINTID, VKD3D_DATA_UINT, 0);
     rel_addr->swizzle = 0;
     rel_addr->modifiers = 0;

@@ -300,32 +302,25 @@ static bool shader_dst_param_normalise_outpointid(struct vkd3d_shader_dst_param

if (normaliser_is_in_control_point_phase(normaliser) && reg->type == VKD3DSPR_OUTPUT)
     {
-        if (reg->idx[2].offset != ~0u)
-        {
-            FIXME("Cannot insert phase id.\n");
-            return false;
-        }
-        if (reg->idx[1].offset != ~0u)
-        {
-            WARN("Unexpected address at index 1.\n");
-            reg->idx[2] = reg->idx[1];
-        }
+        /* The TPF reader validates idx_count. */
+        assert(reg->idx_count == 1);
         reg->idx[1] = reg->idx[0];
         /* The control point id param is implicit here. Avoid later complications by inserting it. */
         reg->idx[0].offset = 0;
         reg->idx[0].rel_addr = normaliser->outpointid_param;
+        ++reg->idx_count;
     }

return true;
 }

-static void shader_dst_param_io_init(struct vkd3d_shader_dst_param *param,
-        const struct signature_element *e, enum vkd3d_shader_register_type reg_type)
+static void shader_dst_param_io_init(struct vkd3d_shader_dst_param *param, const struct signature_element *e,
+        enum vkd3d_shader_register_type reg_type, unsigned int idx_count)
 {
     param->write_mask = e->mask;
     param->modifiers = 0;
     param->shift = 0;
-    shader_register_init(&param->reg, reg_type, vkd3d_data_type_from_component_type(e->component_type));
+    shader_register_init(&param->reg, reg_type, vkd3d_data_type_from_component_type(e->component_type), idx_count);
 }

static enum vkd3d_result shader_normaliser_emit_hs_input(struct vkd3d_shader_normaliser *normaliser,
@@ -369,7 +364,7 @@ static enum vkd3d_result shader_normaliser_emit_hs_input(struct vkd3d_shader_nor
             param = &ins->declaration.dst;
         }

-        shader_dst_param_io_init(param, e, VKD3DSPR_INPUT);
+        shader_dst_param_io_init(param, e, VKD3DSPR_INPUT, 2);
         param->reg.idx[0].offset = input_control_point_count;
         param->reg.idx[1].offset = i;

diff --git a/libs/vkd3d-shader/tpf.c b/libs/vkd3d-shader/tpf.c
index fcd7cf93..0172c77e 100644
--- a/libs/vkd3d-shader/tpf.c
+++ b/libs/vkd3d-shader/tpf.c
@@ -564,6 +564,8 @@ struct vkd3d_shader_sm4_parser

unsigned int output_map[MAX_REG_OUTPUT];

+    enum vkd3d_shader_opcode phase;
+
     struct vkd3d_shader_parser p;
 };

@@ -1669,6 +1671,8 @@ static bool shader_sm4_read_param(struct vkd3d_shader_sm4_parser *priv, const ui
         return false;
     }

+    param->idx_count = order;
+
     if (register_type == VKD3D_SM4_RT_IMMCONST || register_type == VKD3D_SM4_RT_IMMCONST64)
     {
         enum vkd3d_sm4_dimension dimension = (token & VKD3D_SM4_DIMENSION_MASK) >> VKD3D_SM4_DIMENSION_SHIFT;
@@ -1710,6 +1714,7 @@ static bool shader_sm4_read_param(struct vkd3d_shader_sm4_parser *priv, const ui
          * other values up one slot. Normalize to SM5.1. */
         param->idx[2] = param->idx[1];
         param->idx[1] = param->idx[0];
+        ++param->idx_count;
     }

map_register(priv, param);
@@ -1742,6 +1747,47 @@ static uint32_t swizzle_from_sm4(uint32_t s)
     return vkd3d_shader_create_swizzle(s & 0x3, (s >> 2) & 0x3, (s >> 4) & 0x3, (s >> 6) & 0x3);
 }

+static bool register_is_input_output(const struct vkd3d_shader_register *reg)
+{
+    switch (reg->type)
+    {
+        case VKD3DSPR_INPUT:
+        case VKD3DSPR_OUTPUT:
+        case VKD3DSPR_COLOROUT:
+        case VKD3DSPR_INCONTROLPOINT:
+        case VKD3DSPR_OUTCONTROLPOINT:
+        case VKD3DSPR_PATCHCONST:
+            return true;
+
+        default:
+            return false;
+    }
+}
+
+static bool register_is_control_point_input(const struct vkd3d_shader_register *reg,
+        const struct vkd3d_shader_sm4_parser *priv)
+{
+    return reg->type == VKD3DSPR_INCONTROLPOINT || reg->type == VKD3DSPR_OUTCONTROLPOINT
+            || (reg->type == VKD3DSPR_INPUT && (priv->phase == VKD3DSIH_HS_CONTROL_POINT_PHASE
+            || priv->p.shader_version.type == VKD3D_SHADER_TYPE_GEOMETRY));
+}
+
+static bool shader_sm4_validate_input_output_register(struct vkd3d_shader_sm4_parser *priv,
+        const struct vkd3d_shader_register *reg)
+{
+    unsigned int idx_count = 1 + register_is_control_point_input(reg, priv);
+
+    if (reg->idx_count != idx_count)
+    {
+        vkd3d_shader_parser_error(&priv->p, VKD3D_SHADER_ERROR_TPF_INVALID_REGISTER_INDEX_COUNT,
+                "Invalid index count %u for register type %#x; expected count %u.\n",
+                reg->idx_count, reg->type, idx_count);
+        return false;
+    }
+
+    return true;
+}
+
 static bool shader_sm4_read_src_param(struct vkd3d_shader_sm4_parser *priv, const uint32_t **ptr,
         const uint32_t *end, enum vkd3d_data_type data_type, struct vkd3d_shader_src_param *src_param)
 {
@@ -1793,6 +1839,9 @@ static bool shader_sm4_read_src_param(struct vkd3d_shader_sm4_parser *priv, cons
         }
     }

+    if (register_is_input_output(&src_param->reg) && !shader_sm4_validate_input_output_register(priv, &src_param->reg))
+        return false;
+
     return true;
 }

@@ -1830,6 +1879,9 @@ static bool shader_sm4_read_dst_param(struct vkd3d_shader_sm4_parser *priv, cons
     dst_param->modifiers = 0;
     dst_param->shift = 0;

+    if (register_is_input_output(&dst_param->reg) && !shader_sm4_validate_input_output_register(priv, &dst_param->reg))
+        return false;
+
     return true;
 }

@@ -1967,6 +2019,9 @@ static void shader_sm4_read_instruction(struct vkd3d_shader_sm4_parser *sm4, str
     }

ins->handler_idx = opcode_info->handler_idx;
+    if (ins->handler_idx == VKD3DSIH_HS_CONTROL_POINT_PHASE || ins->handler_idx == VKD3DSIH_HS_FORK_PHASE
+            || ins->handler_idx == VKD3DSIH_HS_JOIN_PHASE)
+        sm4->phase = ins->handler_idx;
     ins->flags = 0;
     ins->coissue = false;
     ins->raw = false;
diff --git a/libs/vkd3d-shader/vkd3d_shader_private.h b/libs/vkd3d-shader/vkd3d_shader_private.h
index f1641c1e..2546c0d2 100644
--- a/libs/vkd3d-shader/vkd3d_shader_private.h
+++ b/libs/vkd3d-shader/vkd3d_shader_private.h
@@ -74,6 +74,7 @@ enum vkd3d_shader_error
     VKD3D_SHADER_ERROR_TPF_MISMATCHED_CF                = 1000,
     VKD3D_SHADER_ERROR_TPF_INVALID_REGISTER_RANGE       = 1001,
     VKD3D_SHADER_ERROR_TPF_OUT_OF_MEMORY                = 1002,
+    VKD3D_SHADER_ERROR_TPF_INVALID_REGISTER_INDEX_COUNT = 1003,

VKD3D_SHADER_ERROR_SPV_DESCRIPTOR_BINDING_NOT_FOUND = 2000,
     VKD3D_SHADER_ERROR_SPV_INVALID_REGISTER_TYPE        = 2001,
@@ -676,6 +677,7 @@ struct vkd3d_shader_register
     bool non_uniform;
     enum vkd3d_data_type data_type;
     struct vkd3d_shader_register_index idx[3];
+    unsigned int idx_count;
     enum vkd3d_immconst_type immconst_type;
     union
     {