[opt] Fix folding rule for extract from copy logical (KhronosGroup#6725)

Fixes KhronosGroup#6712

* Added missing vector cases for composite types

Author: alan-baker

source/opt/folding_rules.cpp

@@ -2180,14 +2180,18 @@ uint32_t GetElementType(uint32_t type_id, Instruction::iterator start,
     const Instruction* type_inst = def_use_manager->GetDef(type_id);
     assert(index.type == SPV_OPERAND_TYPE_LITERAL_INTEGER &&
            index.words.size() == 1);
-    if (type_inst->opcode() == spv::Op::OpTypeArray) {
-      type_id = type_inst->GetSingleWordInOperand(0);
-    } else if (type_inst->opcode() == spv::Op::OpTypeMatrix) {
-      type_id = type_inst->GetSingleWordInOperand(0);
-    } else if (type_inst->opcode() == spv::Op::OpTypeStruct) {
-      type_id = type_inst->GetSingleWordInOperand(index.words[0]);
-    } else {
-      return 0;
+    switch (type_inst->opcode()) {
+      case spv::Op::OpTypeArray:
+      case spv::Op::OpTypeMatrix:
+      case spv::Op::OpTypeVector:
+      case spv::Op::OpTypeVectorIdEXT:
+        type_id = type_inst->GetSingleWordInOperand(0);
+        break;
+      case spv::Op::OpTypeStruct:
+        type_id = type_inst->GetSingleWordInOperand(index.words[0]);
+        break;
+      default:
+        return 0;
     }
   }
   return type_id;

test/opt/fold_test.cpp

@@ -14683,6 +14683,34 @@ ::testing::Values(
                OpReturn
                OpFunctionEnd
         )",
+        13, true),
+    // Test case 31: Fold OpCopyLogical feeding extract into vector.
+    InstructionFoldingCase<bool>(
+        Header() + R"(
+; CHECK: [[uint:%\w+]] = OpTypeInt 32 0
+; CHECK: [[v2uint:%\w+]] = OpTypeVector [[uint]] 2
+; CHECK: [[struct_type1:%\w+]] = OpTypeStruct [[v2uint]]
+; CHECK: [[struct_type2:%\w+]] = OpTypeStruct [[v2uint]]
+; CHECK: [[struct_type3:%\w+]] = OpTypeStruct [[struct_type1]]
+; CHECK: [[struct_type4:%\w+]] = OpTypeStruct [[struct_type2]]
+; CHECK: [[var:%\w+]] = OpVariable
+; CHECK: [[ld:%\w+]] = OpLoad [[struct_type3]] [[var]]
+; CHECK: [[ex:%\w+]] = OpCompositeExtract [[uint]] [[ld]] 0 0 0
+; CHECK: %13 = OpCopyObject [[uint]] [[ex]]
+    %struct1 = OpTypeStruct %v2uint
+    %struct2 = OpTypeStruct %v2uint
+    %struct3 = OpTypeStruct %struct1
+    %struct4 = OpTypeStruct %struct2
+%_ptr_StorageBuffer_struct3 = OpTypePointer StorageBuffer %struct3
+%var1 = OpVariable %_ptr_StorageBuffer_struct3 StorageBuffer
+       %main = OpFunction %void None %void_func
+          %4 = OpLabel
+         %11 = OpLoad %struct3 %var1
+         %12 = OpCopyLogical %struct4 %11
+         %13 = OpCompositeExtract %uint %12 0 0 0
+               OpReturn
+               OpFunctionEnd
+        )",
         13, true)
 ));