Validate loop merge (#2579)

Fixes #2559

* Validate OpLoopMerge including loop controls
  * add tests
  * fix some bad tests
diff --git a/source/val/validate_cfg.cpp b/source/val/validate_cfg.cpp
index 6e7acbd..17f144c 100644
--- a/source/val/validate_cfg.cpp
+++ b/source/val/validate_cfg.cpp
@@ -230,6 +230,78 @@
   return SPV_SUCCESS;
 }
 
+spv_result_t ValidateLoopMerge(ValidationState_t& _, const Instruction* inst) {
+  const auto merge_id = inst->GetOperandAs<uint32_t>(0);
+  const auto merge = _.FindDef(merge_id);
+  if (!merge || merge->opcode() != SpvOpLabel) {
+    return _.diag(SPV_ERROR_INVALID_ID, inst)
+           << "Merge Block " << _.getIdName(merge_id) << " must be an OpLabel";
+  }
+
+  const auto continue_id = inst->GetOperandAs<uint32_t>(1);
+  const auto continue_target = _.FindDef(continue_id);
+  if (!continue_target || continue_target->opcode() != SpvOpLabel) {
+    return _.diag(SPV_ERROR_INVALID_ID, inst)
+           << "Continue Target " << _.getIdName(continue_id)
+           << " must be an OpLabel";
+  }
+
+  if (merge_id == continue_id) {
+    return _.diag(SPV_ERROR_INVALID_ID, inst)
+           << "Merge Block and Continue Target must be different ids";
+  }
+
+  const auto loop_control = inst->GetOperandAs<uint32_t>(2);
+  if ((loop_control >> SpvLoopControlUnrollShift) & 0x1 &&
+      (loop_control >> SpvLoopControlDontUnrollShift) & 0x1) {
+    return _.diag(SPV_ERROR_INVALID_DATA, inst)
+           << "Unroll and DontUnroll loop controls must not both be specified";
+  }
+  if ((loop_control >> SpvLoopControlDontUnrollShift) & 0x1 &&
+      (loop_control >> SpvLoopControlPeelCountShift) & 0x1) {
+    return _.diag(SPV_ERROR_INVALID_DATA, inst) << "PeelCount and DontUnroll "
+                                                   "loop controls must not "
+                                                   "both be specified";
+  }
+  if ((loop_control >> SpvLoopControlDontUnrollShift) & 0x1 &&
+      (loop_control >> SpvLoopControlPartialCountShift) & 0x1) {
+    return _.diag(SPV_ERROR_INVALID_DATA, inst) << "PartialCount and "
+                                                   "DontUnroll loop controls "
+                                                   "must not both be specified";
+  }
+
+  uint32_t operand = 3;
+  if ((loop_control >> SpvLoopControlDependencyLengthShift) & 0x1) {
+    ++operand;
+  }
+  if ((loop_control >> SpvLoopControlMinIterationsShift) & 0x1) {
+    ++operand;
+  }
+  if ((loop_control >> SpvLoopControlMaxIterationsShift) & 0x1) {
+    ++operand;
+  }
+  if ((loop_control >> SpvLoopControlIterationMultipleShift) & 0x1) {
+    if (inst->operands().size() < operand ||
+        inst->GetOperandAs<uint32_t>(operand) == 0) {
+      return _.diag(SPV_ERROR_INVALID_DATA, inst) << "IterationMultiple loop "
+                                                     "control operand must be "
+                                                     "greater than zero";
+    }
+    ++operand;
+  }
+  if ((loop_control >> SpvLoopControlPeelCountShift) & 0x1) {
+    ++operand;
+  }
+  if ((loop_control >> SpvLoopControlPartialCountShift) & 0x1) {
+    ++operand;
+  }
+
+  // That the right number of operands is present is checked by the parser. The
+  // above code tracks operands for expanded validation checking in the future.
+
+  return SPV_SUCCESS;
+}
+
 }  // namespace
 
 void printDominatorList(const BasicBlock& b) {
@@ -931,6 +1003,9 @@
     case SpvOpSwitch:
       if (auto error = ValidateSwitch(_, inst)) return error;
       break;
+    case SpvOpLoopMerge:
+      if (auto error = ValidateLoopMerge(_, inst)) return error;
+      break;
     default:
       break;
   }
diff --git a/test/val/val_cfg_test.cpp b/test/val/val_cfg_test.cpp
index 72bfa42..00a24ef 100644
--- a/test/val/val_cfg_test.cpp
+++ b/test/val/val_cfg_test.cpp
@@ -1952,7 +1952,7 @@
 %funct   = OpTypeFunction %voidt
 %main    = OpFunction %voidt None %funct
 %loop    = OpLabel
-           OpLoopMerge %exit %exit None
+           OpLoopMerge %exit %loop None
            OpBranch %exit
 %exit    = OpLabel
            OpReturn
@@ -3028,6 +3028,240 @@
                         "to the loop header <ID> 7"));
 }
 
+TEST_F(ValidateCFG, LoopMergeMergeBlockNotLabel) {
+  const std::string text = R"(
+OpCapability Shader
+OpCapability Linkage
+OpMemoryModel Logical GLSL450
+OpName %undef "undef"
+%void = OpTypeVoid
+%bool = OpTypeBool
+%undef = OpUndef %bool
+%void_fn = OpTypeFunction %void
+%func = OpFunction %void None %void_fn
+%1 = OpLabel
+OpLoopMerge %undef %2 None
+OpBranchConditional %undef %2 %2
+%2 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  CompileSuccessfully(text);
+  EXPECT_EQ(SPV_ERROR_INVALID_ID, ValidateInstructions());
+  EXPECT_THAT(getDiagnosticString(),
+              HasSubstr("Merge Block 1[%undef] must be an OpLabel"));
+}
+
+TEST_F(ValidateCFG, LoopMergeContinueTargetNotLabel) {
+  const std::string text = R"(
+OpCapability Shader
+OpCapability Linkage
+OpMemoryModel Logical GLSL450
+OpName %undef "undef"
+%void = OpTypeVoid
+%bool = OpTypeBool
+%undef = OpUndef %bool
+%void_fn = OpTypeFunction %void
+%func = OpFunction %void None %void_fn
+%1 = OpLabel
+OpLoopMerge %2 %undef None
+OpBranchConditional %undef %2 %2
+%2 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  CompileSuccessfully(text);
+  EXPECT_EQ(SPV_ERROR_INVALID_ID, ValidateInstructions());
+  EXPECT_THAT(getDiagnosticString(),
+              HasSubstr("Continue Target 1[%undef] must be an OpLabel"));
+}
+
+TEST_F(ValidateCFG, LoopMergeMergeBlockContinueTargetSameLabel) {
+  const std::string text = R"(
+OpCapability Shader
+OpCapability Linkage
+OpMemoryModel Logical GLSL450
+OpName %undef "undef"
+%void = OpTypeVoid
+%bool = OpTypeBool
+%undef = OpUndef %bool
+%void_fn = OpTypeFunction %void
+%func = OpFunction %void None %void_fn
+%1 = OpLabel
+OpLoopMerge %2 %2 None
+OpBranchConditional %undef %2 %2
+%2 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  CompileSuccessfully(text);
+  EXPECT_EQ(SPV_ERROR_INVALID_ID, ValidateInstructions());
+  EXPECT_THAT(
+      getDiagnosticString(),
+      HasSubstr("Merge Block and Continue Target must be different ids"));
+}
+
+TEST_F(ValidateCFG, LoopMergeUnrollAndDontUnroll) {
+  const std::string text = R"(
+OpCapability Shader
+OpCapability Linkage
+OpMemoryModel Logical GLSL450
+OpName %undef "undef"
+%void = OpTypeVoid
+%bool = OpTypeBool
+%undef = OpUndef %bool
+%void_fn = OpTypeFunction %void
+%func = OpFunction %void None %void_fn
+%5 = OpLabel
+OpBranch %1
+%1 = OpLabel
+OpLoopMerge %2 %3 Unroll|DontUnroll
+OpBranchConditional %undef %2 %3
+%3 = OpLabel
+OpBranch %1
+%2 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  CompileSuccessfully(text);
+  EXPECT_EQ(SPV_ERROR_INVALID_DATA, ValidateInstructions());
+  EXPECT_THAT(
+      getDiagnosticString(),
+      HasSubstr(
+          "Unroll and DontUnroll loop controls must not both be specified"));
+}
+
+TEST_F(ValidateCFG, LoopMergePeelCountAndDontUnroll) {
+  const std::string text = R"(
+OpCapability Shader
+OpCapability Linkage
+OpMemoryModel Logical GLSL450
+OpName %undef "undef"
+%void = OpTypeVoid
+%bool = OpTypeBool
+%undef = OpUndef %bool
+%void_fn = OpTypeFunction %void
+%func = OpFunction %void None %void_fn
+%5 = OpLabel
+OpBranch %1
+%1 = OpLabel
+OpLoopMerge %2 %3 DontUnroll|PeelCount 1
+OpBranchConditional %undef %2 %3
+%3 = OpLabel
+OpBranch %1
+%2 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  CompileSuccessfully(text, SPV_ENV_UNIVERSAL_1_4);
+  EXPECT_EQ(SPV_ERROR_INVALID_DATA,
+            ValidateInstructions(SPV_ENV_UNIVERSAL_1_4));
+  EXPECT_THAT(
+      getDiagnosticString(),
+      HasSubstr(
+          "PeelCount and DontUnroll loop controls must not both be specified"));
+}
+
+TEST_F(ValidateCFG, LoopMergePartialCountAndDontUnroll) {
+  const std::string text = R"(
+OpCapability Shader
+OpCapability Linkage
+OpMemoryModel Logical GLSL450
+OpName %undef "undef"
+%void = OpTypeVoid
+%bool = OpTypeBool
+%undef = OpUndef %bool
+%void_fn = OpTypeFunction %void
+%func = OpFunction %void None %void_fn
+%5 = OpLabel
+OpBranch %1
+%1 = OpLabel
+OpLoopMerge %2 %3 DontUnroll|PartialCount 1
+OpBranchConditional %undef %2 %3
+%3 = OpLabel
+OpBranch %1
+%2 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  CompileSuccessfully(text, SPV_ENV_UNIVERSAL_1_4);
+  EXPECT_EQ(SPV_ERROR_INVALID_DATA,
+            ValidateInstructions(SPV_ENV_UNIVERSAL_1_4));
+  EXPECT_THAT(getDiagnosticString(),
+              HasSubstr("PartialCount and DontUnroll loop controls must not "
+                        "both be specified"));
+}
+
+TEST_F(ValidateCFG, LoopMergeIterationMultipleZero) {
+  const std::string text = R"(
+OpCapability Shader
+OpCapability Linkage
+OpMemoryModel Logical GLSL450
+OpName %undef "undef"
+%void = OpTypeVoid
+%bool = OpTypeBool
+%undef = OpUndef %bool
+%void_fn = OpTypeFunction %void
+%func = OpFunction %void None %void_fn
+%5 = OpLabel
+OpBranch %1
+%1 = OpLabel
+OpLoopMerge %2 %3 IterationMultiple 0
+OpBranchConditional %undef %2 %3
+%3 = OpLabel
+OpBranch %1
+%2 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  CompileSuccessfully(text, SPV_ENV_UNIVERSAL_1_4);
+  EXPECT_EQ(SPV_ERROR_INVALID_DATA,
+            ValidateInstructions(SPV_ENV_UNIVERSAL_1_4));
+  EXPECT_THAT(
+      getDiagnosticString(),
+      HasSubstr(
+          "IterationMultiple loop control operand must be greater than zero"));
+}
+
+TEST_F(ValidateCFG, LoopMergeIterationMultipleZeroMoreOperands) {
+  const std::string text = R"(
+OpCapability Shader
+OpCapability Linkage
+OpMemoryModel Logical GLSL450
+OpName %undef "undef"
+%void = OpTypeVoid
+%bool = OpTypeBool
+%undef = OpUndef %bool
+%void_fn = OpTypeFunction %void
+%func = OpFunction %void None %void_fn
+%5 = OpLabel
+OpBranch %1
+%1 = OpLabel
+OpLoopMerge %2 %3 MaxIterations|IterationMultiple 4 0
+OpBranchConditional %undef %2 %3
+%3 = OpLabel
+OpBranch %1
+%2 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  CompileSuccessfully(text, SPV_ENV_UNIVERSAL_1_4);
+  EXPECT_EQ(SPV_ERROR_INVALID_DATA,
+            ValidateInstructions(SPV_ENV_UNIVERSAL_1_4));
+  EXPECT_THAT(
+      getDiagnosticString(),
+      HasSubstr(
+          "IterationMultiple loop control operand must be greater than zero"));
+}
+
 /// TODO(umar): Nested CFG constructs
 
 }  // namespace
diff --git a/test/val/val_limits_test.cpp b/test/val/val_limits_test.cpp
index 791b709..becf7be 100644
--- a/test/val/val_limits_test.cpp
+++ b/test/val/val_limits_test.cpp
@@ -754,13 +754,17 @@
            OpName %loop "loop"
            OpName %exit "exit"
 %voidt   = OpTypeVoid
+%boolt   = OpTypeBool
+%undef   = OpUndef %boolt
 %funct   = OpTypeFunction %voidt
 %main    = OpFunction %voidt None %funct
 %entry   = OpLabel
            OpBranch %exit
 %loop    = OpLabel
-           OpLoopMerge %loop %loop None
-           OpBranch %loop
+           OpLoopMerge %dead %loop None
+           OpBranchConditional %undef %loop %loop
+%dead    = OpLabel
+           OpUnreachable
 %exit    = OpLabel
            OpReturn
            OpFunctionEnd