Use OpReturn* in wrap-opkill

The warp-opkill pass is generating incorrect code.  It is placing an
OpUnreachable at the end of a basic block, when the block can be
reached.  We can't reach the end of the block, but we can reach the end.
Instead we will add a return instruction.

Fixes #2875.

source/opt/ir_builder.h

@@ -59,27 +59,44 @@ class InstructionBuilder {
                            preserved_analyses) {}
 
   Instruction* AddNullaryOp(uint32_t type_id, SpvOp opcode) {
-    // TODO(1841): Handle id overflow.
-    std::unique_ptr<Instruction> newUnOp(new Instruction(
-        GetContext(), opcode, type_id,
-        opcode == SpvOpReturn ? 0 : GetContext()->TakeNextId(), {}));
-    return AddInstruction(std::move(newUnOp));
+    uint32_t result_id = 0;
+    if (type_id != 0) {
+      result_id = GetContext()->TakeNextId();
+      if (result_id == 0) {
+        return nullptr;
+      }
+    }
+    std::unique_ptr<Instruction> new_inst(
+        new Instruction(GetContext(), opcode, type_id, result_id, {}));
+    return AddInstruction(std::move(new_inst));
   }
 
   Instruction* AddUnaryOp(uint32_t type_id, SpvOp opcode, uint32_t operand1) {
-    // TODO(1841): Handle id overflow.
+    uint32_t result_id = 0;
+    if (type_id != 0) {
+      result_id = GetContext()->TakeNextId();
+      if (result_id == 0) {
+        return nullptr;
+      }
+    }
     std::unique_ptr<Instruction> newUnOp(new Instruction(
-        GetContext(), opcode, type_id, GetContext()->TakeNextId(),
+        GetContext(), opcode, type_id, result_id,
         {{spv_operand_type_t::SPV_OPERAND_TYPE_ID, {operand1}}}));
     return AddInstruction(std::move(newUnOp));
   }
 
   Instruction* AddBinaryOp(uint32_t type_id, SpvOp opcode, uint32_t operand1,
                            uint32_t operand2) {
-    // TODO(1841): Handle id overflow.
+    uint32_t result_id = 0;
+    if (type_id != 0) {
+      result_id = GetContext()->TakeNextId();
+      if (result_id == 0) {
+        return nullptr;
+      }
+    }
     std::unique_ptr<Instruction> newBinOp(new Instruction(
         GetContext(), opcode, type_id,
-        opcode == SpvOpStore ? 0 : GetContext()->TakeNextId(),
+        opcode == SpvOpStore ? 0 : result_id,
         {{spv_operand_type_t::SPV_OPERAND_TYPE_ID, {operand1}},
          {spv_operand_type_t::SPV_OPERAND_TYPE_ID, {operand2}}}));
     return AddInstruction(std::move(newBinOp));
@@ -87,9 +104,15 @@ class InstructionBuilder {
 
   Instruction* AddTernaryOp(uint32_t type_id, SpvOp opcode, uint32_t operand1,
                             uint32_t operand2, uint32_t operand3) {
-    // TODO(1841): Handle id overflow.
+    uint32_t result_id = 0;
+    if (type_id != 0) {
+      result_id = GetContext()->TakeNextId();
+      if (result_id == 0) {
+        return nullptr;
+      }
+    }
     std::unique_ptr<Instruction> newTernOp(new Instruction(
-        GetContext(), opcode, type_id, GetContext()->TakeNextId(),
+        GetContext(), opcode, type_id, result_id,
         {{spv_operand_type_t::SPV_OPERAND_TYPE_ID, {operand1}},
          {spv_operand_type_t::SPV_OPERAND_TYPE_ID, {operand2}},
          {spv_operand_type_t::SPV_OPERAND_TYPE_ID, {operand3}}}));
@@ -99,9 +122,15 @@ class InstructionBuilder {
   Instruction* AddQuadOp(uint32_t type_id, SpvOp opcode, uint32_t operand1,
                          uint32_t operand2, uint32_t operand3,
                          uint32_t operand4) {
-    // TODO(1841): Handle id overflow.
+    uint32_t result_id = 0;
+    if (type_id != 0) {
+      result_id = GetContext()->TakeNextId();
+      if (result_id == 0) {
+        return nullptr;
+      }
+    }
     std::unique_ptr<Instruction> newQuadOp(new Instruction(
-        GetContext(), opcode, type_id, GetContext()->TakeNextId(),
+        GetContext(), opcode, type_id, result_id,
         {{spv_operand_type_t::SPV_OPERAND_TYPE_ID, {operand1}},
          {spv_operand_type_t::SPV_OPERAND_TYPE_ID, {operand2}},
          {spv_operand_type_t::SPV_OPERAND_TYPE_ID, {operand3}},
@@ -111,9 +140,15 @@ class InstructionBuilder {
 
   Instruction* AddIdLiteralOp(uint32_t type_id, SpvOp opcode, uint32_t id,
                               uint32_t uliteral) {
-    // TODO(1841): Handle id overflow.
+    uint32_t result_id = 0;
+    if (type_id != 0) {
+      result_id = GetContext()->TakeNextId();
+      if (result_id == 0) {
+        return nullptr;
+      }
+    }
     std::unique_ptr<Instruction> newBinOp(new Instruction(
-        GetContext(), opcode, type_id, GetContext()->TakeNextId(),
+        GetContext(), opcode, type_id, result_id,
         {{spv_operand_type_t::SPV_OPERAND_TYPE_ID, {id}},
          {spv_operand_type_t::SPV_OPERAND_TYPE_LITERAL_INTEGER, {uliteral}}}));
     return AddInstruction(std::move(newBinOp));

source/opt/wrap_opkill.cpp

@@ -59,7 +59,14 @@ bool WrapOpKill::ReplaceWithFunctionCall(Instruction* inst) {
   if (ir_builder.AddFunctionCall(GetVoidTypeId(), func_id, {}) == nullptr) {
     return false;
   }
-  ir_builder.AddUnreachable();
+
+  uint32_t return_type_id = GetOwningFunctionsReturnType(inst);
+  if (return_type_id != GetVoidTypeId()) {
+    Instruction* undef = ir_builder.AddNullaryOp(return_type_id, SpvOpUndef);
+    ir_builder.AddUnaryOp(0, SpvOpReturnValue, undef->result_id());
+  } else {
+    ir_builder.AddNullaryOp(0, SpvOpReturn);
+  }
   context()->KillInst(inst);
   return true;
 }
@@ -147,5 +154,15 @@ uint32_t WrapOpKill::GetOpKillFuncId() {
   return opkill_function_->result_id();
 }
 
+uint32_t WrapOpKill::GetOwningFunctionsReturnType(Instruction* inst) {
+  BasicBlock* bb = context()->get_instr_block(inst);
+  if (bb == nullptr) {
+    return 0;
+  }
+
+  Function* func = bb->GetParent();
+  return func->type_id();
+}
+
 }  // namespace opt
 }  // namespace spvtools

source/opt/wrap_opkill.h

@@ -56,6 +56,10 @@ class WrapOpKill : public Pass {
   // function could not be generated.
   uint32_t GetOpKillFuncId();
 
+  // Returns the id of the return type for the function that contains |inst|.
+  // Returns 0 if |inst| is not in a function.
+  uint32_t GetOwningFunctionsReturnType(Instruction* inst);
+
   // The id of the void type.  If its value is 0, then the void type has not
   // been found or created yet.
   uint32_t void_type_id_;

test/opt/wrap_opkill_test.cpp

@@ -31,7 +31,7 @@ TEST_F(WrapOpKillTest, SingleOpKill) {
 ; CHECK: [[orig_kill]] = OpFunction
 ; CHECK-NEXT: OpLabel
 ; CHECK-NEXT: OpFunctionCall %void [[new_kill:%\w+]]
-; CHECK-NEXT: OpUnreachable
+; CHECK-NEXT: OpReturn
 ; CHECK: [[new_kill]] = OpFunction
 ; CHECK-NEXT: OpLabel
 ; CHECK-NEXT: OpKill
@@ -83,10 +83,10 @@ TEST_F(WrapOpKillTest, MultipleOpKillInSameFunc) {
 ; CHECK-NEXT: OpBranchConditional
 ; CHECK-NEXT: OpLabel
 ; CHECK-NEXT: OpFunctionCall %void [[new_kill:%\w+]]
-; CHECK-NEXT: OpUnreachable
+; CHECK-NEXT: OpReturn
 ; CHECK-NEXT: OpLabel
 ; CHECK-NEXT: OpFunctionCall %void [[new_kill]]
-; CHECK-NEXT: OpUnreachable
+; CHECK-NEXT: OpReturn
 ; CHECK: [[new_kill]] = OpFunction
 ; CHECK-NEXT: OpLabel
 ; CHECK-NEXT: OpKill
@@ -143,11 +143,11 @@ TEST_F(WrapOpKillTest, MultipleOpKillInDifferentFunc) {
 ; CHECK: [[orig_kill1]] = OpFunction
 ; CHECK-NEXT: OpLabel
 ; CHECK-NEXT: OpFunctionCall %void [[new_kill:%\w+]]
-; CHECK-NEXT: OpUnreachable
+; CHECK-NEXT: OpReturn
 ; CHECK: [[orig_kill2]] = OpFunction
 ; CHECK-NEXT: OpLabel
 ; CHECK-NEXT: OpFunctionCall %void [[new_kill]]
-; CHECK-NEXT: OpUnreachable
+; CHECK-NEXT: OpReturn
 ; CHECK: [[new_kill]] = OpFunction
 ; CHECK-NEXT: OpLabel
 ; CHECK-NEXT: OpKill
@@ -193,6 +193,58 @@ TEST_F(WrapOpKillTest, MultipleOpKillInDifferentFunc) {
   SinglePassRunAndMatch<WrapOpKill>(text, true);
 }
 
+TEST_F(WrapOpKillTest, FuncWithReturnValue) {
+  const std::string text = R"(
+; CHECK: OpEntryPoint Fragment [[main:%\w+]]
+; CHECK: [[main]] = OpFunction
+; CHECK: OpFunctionCall %int [[orig_kill:%\w+]]
+; CHECK: [[orig_kill]] = OpFunction
+; CHECK-NEXT: OpLabel
+; CHECK-NEXT: OpFunctionCall %void [[new_kill:%\w+]]
+; CHECK-NEXT: [[undef:%\w+]] = OpUndef %int
+; CHECK-NEXT: OpReturnValue [[undef]]
+; CHECK: [[new_kill]] = OpFunction
+; CHECK-NEXT: OpLabel
+; CHECK-NEXT: OpKill
+; CHECK-NEXT: OpFunctionEnd
+               OpCapability Shader
+          %1 = OpExtInstImport "GLSL.std.450"
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint Fragment %main "main"
+               OpExecutionMode %main OriginUpperLeft
+               OpSource GLSL 330
+               OpName %main "main"
+       %void = OpTypeVoid
+          %5 = OpTypeFunction %void
+        %int = OpTypeInt 32 1
+  %func_type = OpTypeFunction %int
+       %bool = OpTypeBool
+       %true = OpConstantTrue %bool
+       %main = OpFunction %void None %5
+          %8 = OpLabel
+               OpBranch %9
+          %9 = OpLabel
+               OpLoopMerge %10 %11 None
+               OpBranch %12
+         %12 = OpLabel
+               OpBranchConditional %true %13 %10
+         %13 = OpLabel
+               OpBranch %11
+         %11 = OpLabel
+         %14 = OpFunctionCall %int %kill_
+               OpBranch %9
+         %10 = OpLabel
+               OpReturn
+               OpFunctionEnd
+      %kill_ = OpFunction %int None %func_type
+         %15 = OpLabel
+               OpKill
+               OpFunctionEnd
+  )";
+
+  SinglePassRunAndMatch<WrapOpKill>(text, true);
+}
+
 TEST_F(WrapOpKillTest, IdBoundOverflow1) {
   const std::string text = R"(
 OpCapability GeometryStreams