spirv-fuzz: Wrap OpKill and similar in function calls (#3884)

Part of #3717.
diff --git a/source/fuzz/CMakeLists.txt b/source/fuzz/CMakeLists.txt
index 6843eb8..c936ca8 100644
--- a/source/fuzz/CMakeLists.txt
+++ b/source/fuzz/CMakeLists.txt
@@ -220,6 +220,7 @@
         transformation_swap_conditional_branch_operands.h
         transformation_toggle_access_chain_instruction.h
         transformation_vector_shuffle.h
+        transformation_wrap_early_terminator_in_function.h
         transformation_wrap_region_in_selection.h
         uniform_buffer_element_descriptor.h
         ${CMAKE_CURRENT_BINARY_DIR}/protobufs/spvtoolsfuzz.pb.h
@@ -404,6 +405,7 @@
         transformation_swap_conditional_branch_operands.cpp
         transformation_toggle_access_chain_instruction.cpp
         transformation_vector_shuffle.cpp
+        transformation_wrap_early_terminator_in_function.cpp
         transformation_wrap_region_in_selection.cpp
         uniform_buffer_element_descriptor.cpp
         ${CMAKE_CURRENT_BINARY_DIR}/protobufs/spvtoolsfuzz.pb.cc
diff --git a/source/fuzz/protobufs/spvtoolsfuzz.proto b/source/fuzz/protobufs/spvtoolsfuzz.proto
index e99d5f0..06c585c 100644
--- a/source/fuzz/protobufs/spvtoolsfuzz.proto
+++ b/source/fuzz/protobufs/spvtoolsfuzz.proto
@@ -554,6 +554,7 @@
     TransformationAddEarlyTerminatorWrapper add_early_terminator_wrapper = 80;
     TransformationPropagateInstructionDown propagate_instruction_down = 81;
     TransformationReplaceBranchFromDeadBlockWithExit replace_branch_from_dead_block_with_exit = 82;
+    TransformationWrapEarlyTerminatorInFunction wrap_early_terminator_in_function = 83;
     // Add additional option using the next available number.
   }
 }
@@ -2260,6 +2261,26 @@
 
 }
 
+message TransformationWrapEarlyTerminatorInFunction {
+
+  // Replaces an early terminator - OpKill, OpReachable or OpTerminateInvocation
+  // - with a call to a wrapper function for the terminator.
+
+  // A fresh id for a new OpFunctionCall instruction.
+  uint32 fresh_id = 1;
+
+  // A descriptor for an OpKill, OpUnreachable or OpTerminateInvocation
+  // instruction.
+  InstructionDescriptor early_terminator_instruction = 2;
+
+  // An id with the same type as the enclosing function's return type that is
+  // available at the early terminator.  This is used to change the terminator
+  // to OpReturnValue.  Ignored if the enclosing function has void return type,
+  // in which case OpReturn can be used as the new terminator.
+  uint32 returned_value_id = 3;
+
+}
+
 message TransformationWrapRegionInSelection {
 
   // Transforms a single-entry-single-exit region R into
diff --git a/source/fuzz/transformation.cpp b/source/fuzz/transformation.cpp
index 28903a2..357d59f 100644
--- a/source/fuzz/transformation.cpp
+++ b/source/fuzz/transformation.cpp
@@ -98,6 +98,7 @@
 #include "source/fuzz/transformation_swap_conditional_branch_operands.h"
 #include "source/fuzz/transformation_toggle_access_chain_instruction.h"
 #include "source/fuzz/transformation_vector_shuffle.h"
+#include "source/fuzz/transformation_wrap_early_terminator_in_function.h"
 #include "source/fuzz/transformation_wrap_region_in_selection.h"
 #include "source/util/make_unique.h"
 
@@ -358,6 +359,10 @@
           message.toggle_access_chain_instruction());
     case protobufs::Transformation::TransformationCase::kVectorShuffle:
       return MakeUnique<TransformationVectorShuffle>(message.vector_shuffle());
+    case protobufs::Transformation::TransformationCase::
+        kWrapEarlyTerminatorInFunction:
+      return MakeUnique<TransformationWrapEarlyTerminatorInFunction>(
+          message.wrap_early_terminator_in_function());
     case protobufs::Transformation::TransformationCase::kWrapRegionInSelection:
       return MakeUnique<TransformationWrapRegionInSelection>(
           message.wrap_region_in_selection());
diff --git a/source/fuzz/transformation_wrap_early_terminator_in_function.cpp b/source/fuzz/transformation_wrap_early_terminator_in_function.cpp
new file mode 100644
index 0000000..b8370c8
--- /dev/null
+++ b/source/fuzz/transformation_wrap_early_terminator_in_function.cpp
@@ -0,0 +1,183 @@
+// Copyright (c) 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "source/fuzz/transformation_wrap_early_terminator_in_function.h"
+
+#include "source/fuzz/fuzzer_util.h"
+#include "source/fuzz/instruction_descriptor.h"
+#include "source/util/make_unique.h"
+
+namespace spvtools {
+namespace fuzz {
+
+TransformationWrapEarlyTerminatorInFunction::
+    TransformationWrapEarlyTerminatorInFunction(
+        const spvtools::fuzz::protobufs::
+            TransformationWrapEarlyTerminatorInFunction& message)
+    : message_(message) {}
+
+TransformationWrapEarlyTerminatorInFunction::
+    TransformationWrapEarlyTerminatorInFunction(
+        uint32_t fresh_id,
+        const protobufs::InstructionDescriptor& early_terminator_instruction,
+        uint32_t returned_value_id) {
+  message_.set_fresh_id(fresh_id);
+  *message_.mutable_early_terminator_instruction() =
+      early_terminator_instruction;
+  message_.set_returned_value_id(returned_value_id);
+}
+
+bool TransformationWrapEarlyTerminatorInFunction::IsApplicable(
+    opt::IRContext* ir_context, const TransformationContext& /*unused*/) const {
+  // The given id must be fresh.
+  if (!fuzzerutil::IsFreshId(ir_context, message_.fresh_id())) {
+    return false;
+  }
+
+  // |message_.early_terminator_instruction| must identify an instruciton, and
+  // the instruction must indeed be an early terminator.
+  auto early_terminator =
+      FindInstruction(message_.early_terminator_instruction(), ir_context);
+  if (!early_terminator) {
+    return false;
+  }
+  switch (early_terminator->opcode()) {
+    case SpvOpKill:
+    case SpvOpUnreachable:
+    case SpvOpTerminateInvocation:
+      break;
+    default:
+      return false;
+  }
+  // A wrapper function for the early terminator must exist.
+  auto wrapper_function =
+      MaybeGetWrapperFunction(ir_context, early_terminator->opcode());
+  if (wrapper_function == nullptr) {
+    return false;
+  }
+  auto enclosing_function =
+      ir_context->get_instr_block(early_terminator)->GetParent();
+  // The wrapper function cannot be the function containing the instruction we
+  // would like to wrap.
+  if (wrapper_function->result_id() == enclosing_function->result_id()) {
+    return false;
+  }
+  if (!ir_context->get_type_mgr()
+           ->GetType(enclosing_function->type_id())
+           ->AsVoid()) {
+    // The enclosing function has non-void return type.  We thus need to make
+    // sure that |message_.returned_value_instruction| provides a suitable
+    // result id to use in an OpReturnValue instruction.
+    auto returned_value_instruction =
+        ir_context->get_def_use_mgr()->GetDef(message_.returned_value_id());
+    if (!returned_value_instruction || !returned_value_instruction->type_id() ||
+        returned_value_instruction->type_id() !=
+            enclosing_function->type_id()) {
+      return false;
+    }
+    if (!fuzzerutil::IdIsAvailableBeforeInstruction(
+            ir_context, early_terminator, message_.returned_value_id())) {
+      return false;
+    }
+  }
+  return true;
+}
+
+void TransformationWrapEarlyTerminatorInFunction::Apply(
+    opt::IRContext* ir_context, TransformationContext* /*unused*/) const {
+  fuzzerutil::UpdateModuleIdBound(ir_context, message_.fresh_id());
+  auto early_terminator =
+      FindInstruction(message_.early_terminator_instruction(), ir_context);
+  auto enclosing_block = ir_context->get_instr_block(early_terminator);
+  auto enclosing_function = enclosing_block->GetParent();
+
+  // We would like to add an OpFunctionCall before the block's terminator
+  // instruction, and then change the block's terminator to OpReturn or
+  // OpReturnValue.
+
+  // We get an iterator to the instruction we would like to insert the function
+  // call before.  It will be an iterator to the final instruction in the block
+  // unless the block is a merge block in which case it will be to the
+  // penultimate instruction (because we cannot insert an OpFunctionCall after
+  // a merge instruction).
+  auto iterator = enclosing_block->tail();
+  if (enclosing_block->MergeBlockIdIfAny()) {
+    --iterator;
+  }
+
+  auto wrapper_function =
+      MaybeGetWrapperFunction(ir_context, early_terminator->opcode());
+
+  iterator->InsertBefore(MakeUnique<opt::Instruction>(
+      ir_context, SpvOpFunctionCall, wrapper_function->type_id(),
+      message_.fresh_id(),
+      opt::Instruction::OperandList(
+          {{SPV_OPERAND_TYPE_ID, {wrapper_function->result_id()}}})));
+
+  opt::Instruction::OperandList new_in_operands;
+  if (!ir_context->get_type_mgr()
+           ->GetType(enclosing_function->type_id())
+           ->AsVoid()) {
+    new_in_operands.push_back(
+        {SPV_OPERAND_TYPE_ID, {message_.returned_value_id()}});
+    early_terminator->SetOpcode(SpvOpReturnValue);
+  } else {
+    early_terminator->SetOpcode(SpvOpReturn);
+  }
+  early_terminator->SetInOperands(std::move(new_in_operands));
+
+  ir_context->InvalidateAnalysesExceptFor(opt::IRContext::kAnalysisNone);
+}
+
+std::unordered_set<uint32_t>
+TransformationWrapEarlyTerminatorInFunction::GetFreshIds() const {
+  return std::unordered_set<uint32_t>({message_.fresh_id()});
+}
+
+protobufs::Transformation
+TransformationWrapEarlyTerminatorInFunction::ToMessage() const {
+  protobufs::Transformation result;
+  *result.mutable_wrap_early_terminator_in_function() = message_;
+  return result;
+}
+
+opt::Function*
+TransformationWrapEarlyTerminatorInFunction::MaybeGetWrapperFunction(
+    opt::IRContext* ir_context, SpvOp early_terminator_opcode) {
+  assert((early_terminator_opcode == SpvOpKill ||
+          early_terminator_opcode == SpvOpUnreachable ||
+          early_terminator_opcode == SpvOpTerminateInvocation) &&
+         "Invalid opcode.");
+  auto void_type_id = fuzzerutil::MaybeGetVoidType(ir_context);
+  if (!void_type_id) {
+    return nullptr;
+  }
+  auto void_function_type_id =
+      fuzzerutil::FindFunctionType(ir_context, {void_type_id});
+  if (!void_function_type_id) {
+    return nullptr;
+  }
+  for (auto& function : *ir_context->module()) {
+    if (function.DefInst().GetSingleWordInOperand(1) != void_function_type_id) {
+      continue;
+    }
+    if (function.begin()->begin()->opcode() == early_terminator_opcode) {
+      return &function;
+    }
+  }
+  return nullptr;
+}
+
+}  // namespace fuzz
+}  // namespace spvtools
diff --git a/source/fuzz/transformation_wrap_early_terminator_in_function.h b/source/fuzz/transformation_wrap_early_terminator_in_function.h
new file mode 100644
index 0000000..2629c50
--- /dev/null
+++ b/source/fuzz/transformation_wrap_early_terminator_in_function.h
@@ -0,0 +1,59 @@
+// Copyright (c) 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef SOURCE_FUZZ_TRANSFORMATION_WRAP_EARLY_TERMINATOR_IN_FUNCTION_H_
+#define SOURCE_FUZZ_TRANSFORMATION_WRAP_EARLY_TERMINATOR_IN_FUNCTION_H_
+
+#include "source/fuzz/protobufs/spirvfuzz_protobufs.h"
+#include "source/fuzz/transformation.h"
+#include "source/fuzz/transformation_context.h"
+#include "source/opt/ir_context.h"
+
+namespace spvtools {
+namespace fuzz {
+
+class TransformationWrapEarlyTerminatorInFunction : public Transformation {
+ public:
+  explicit TransformationWrapEarlyTerminatorInFunction(
+      const protobufs::TransformationWrapEarlyTerminatorInFunction& message);
+
+  TransformationWrapEarlyTerminatorInFunction(
+      uint32_t fresh_id,
+      const protobufs::InstructionDescriptor& early_terminator_instruction,
+      uint32_t returned_value_id);
+
+  // TODO comment
+  bool IsApplicable(
+      opt::IRContext* ir_context,
+      const TransformationContext& transformation_context) const override;
+
+  // TODO comment
+  void Apply(opt::IRContext* ir_context,
+             TransformationContext* transformation_context) const override;
+
+  std::unordered_set<uint32_t> GetFreshIds() const override;
+
+  protobufs::Transformation ToMessage() const override;
+
+ private:
+  static opt::Function* MaybeGetWrapperFunction(opt::IRContext* ir_context,
+                                                SpvOp early_terminator_opcode);
+
+  protobufs::TransformationWrapEarlyTerminatorInFunction message_;
+};
+
+}  // namespace fuzz
+}  // namespace spvtools
+
+#endif  // SOURCE_FUZZ_TRANSFORMATION_WRAP_EARLY_TERMINATOR_IN_FUNCTION_H_
diff --git a/test/fuzz/CMakeLists.txt b/test/fuzz/CMakeLists.txt
index d586425..866a924 100644
--- a/test/fuzz/CMakeLists.txt
+++ b/test/fuzz/CMakeLists.txt
@@ -115,6 +115,7 @@
           transformation_toggle_access_chain_instruction_test.cpp
           transformation_record_synonymous_constants_test.cpp
           transformation_vector_shuffle_test.cpp
+          transformation_wrap_early_terminator_in_function_test.cpp
           transformation_wrap_region_in_selection_test.cpp
           uniform_buffer_element_descriptor_test.cpp)
 
diff --git a/test/fuzz/transformation_wrap_early_terminator_in_function_test.cpp b/test/fuzz/transformation_wrap_early_terminator_in_function_test.cpp
new file mode 100644
index 0000000..edef6b4
--- /dev/null
+++ b/test/fuzz/transformation_wrap_early_terminator_in_function_test.cpp
@@ -0,0 +1,315 @@
+// Copyright (c) 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "source/fuzz/transformation_wrap_early_terminator_in_function.h"
+
+#include "source/fuzz/instruction_descriptor.h"
+#include "test/fuzz/fuzz_test_util.h"
+
+namespace spvtools {
+namespace fuzz {
+namespace {
+
+TEST(TransformationWrapEarlyTerminatorInFunctionTest, IsApplicable) {
+  std::string shader = R"(
+               OpCapability Shader
+               OpExtension "SPV_KHR_terminate_invocation"
+          %1 = OpExtInstImport "GLSL.std.450"
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint Fragment %4 "main"
+               OpExecutionMode %4 OriginUpperLeft
+               OpSource ESSL 320
+               OpName %4 "main"
+          %2 = OpTypeVoid
+          %3 = OpTypeFunction %2
+          %6 = OpTypeInt 32 1
+          %7 = OpConstant %6 0
+         %90 = OpTypeBool
+         %91 = OpConstantFalse %90
+
+         %20 = OpTypeFunction %2 %6
+         %21 = OpTypeFunction %6
+
+          %4 = OpFunction %2 None %3
+          %5 = OpLabel
+               OpSelectionMerge %11 None
+               OpSwitch %7 %11 0 %8 1 %9 2 %10
+          %8 = OpLabel
+               OpKill
+          %9 = OpLabel
+               OpUnreachable
+         %10 = OpLabel
+               OpTerminateInvocation
+         %11 = OpLabel
+               OpReturn
+               OpFunctionEnd
+
+         %30 = OpFunction %2 None %3
+         %31 = OpLabel
+               OpKill
+               OpFunctionEnd
+
+         %50 = OpFunction %2 None %3
+         %51 = OpLabel
+               OpTerminateInvocation
+               OpFunctionEnd
+
+         %60 = OpFunction %6 None %21
+         %61 = OpLabel
+               OpBranch %62
+         %62 = OpLabel
+               OpKill
+               OpFunctionEnd
+
+         %70 = OpFunction %6 None %21
+         %71 = OpLabel
+               OpUnreachable
+               OpFunctionEnd
+
+         %80 = OpFunction %2 None %20
+         %81 = OpFunctionParameter %6
+         %82 = OpLabel
+               OpTerminateInvocation
+               OpFunctionEnd
+
+  )";
+
+  const auto env = SPV_ENV_UNIVERSAL_1_4;
+  const auto consumer = nullptr;
+  const auto context = BuildModule(env, consumer, shader, kFuzzAssembleOption);
+  ASSERT_TRUE(IsValid(env, context.get()));
+
+  spvtools::ValidatorOptions validator_options;
+  TransformationContext transformation_context(
+      MakeUnique<FactManager>(context.get()), validator_options);
+
+  // Bad: id is not fresh
+  ASSERT_FALSE(TransformationWrapEarlyTerminatorInFunction(
+                   61, MakeInstructionDescriptor(8, SpvOpKill, 0), 0)
+                   .IsApplicable(context.get(), transformation_context));
+
+  // Bad: early terminator instruction descriptor does not exist
+  ASSERT_FALSE(TransformationWrapEarlyTerminatorInFunction(
+                   100, MakeInstructionDescriptor(82, SpvOpKill, 0), 0)
+                   .IsApplicable(context.get(), transformation_context));
+
+  // Bad: early terminator instruction does not identify an early terminator
+  ASSERT_FALSE(TransformationWrapEarlyTerminatorInFunction(
+                   100, MakeInstructionDescriptor(5, SpvOpSelectionMerge, 0), 0)
+                   .IsApplicable(context.get(), transformation_context));
+
+  // Bad: no wrapper function is available
+  ASSERT_FALSE(TransformationWrapEarlyTerminatorInFunction(
+                   100, MakeInstructionDescriptor(9, SpvOpUnreachable, 0), 0)
+                   .IsApplicable(context.get(), transformation_context));
+
+  // Bad: returned value does not exist
+  ASSERT_FALSE(TransformationWrapEarlyTerminatorInFunction(
+                   100, MakeInstructionDescriptor(62, SpvOpKill, 0), 1000)
+                   .IsApplicable(context.get(), transformation_context));
+
+  // Bad: returned value does not have a type
+  ASSERT_FALSE(TransformationWrapEarlyTerminatorInFunction(
+                   100, MakeInstructionDescriptor(62, SpvOpKill, 0), 61)
+                   .IsApplicable(context.get(), transformation_context));
+
+  // Bad: returned value type does not match
+  ASSERT_FALSE(TransformationWrapEarlyTerminatorInFunction(
+                   100, MakeInstructionDescriptor(62, SpvOpKill, 0), 91)
+                   .IsApplicable(context.get(), transformation_context));
+
+  // Bad: returned value is not available
+  ASSERT_FALSE(TransformationWrapEarlyTerminatorInFunction(
+                   100, MakeInstructionDescriptor(62, SpvOpKill, 0), 81)
+                   .IsApplicable(context.get(), transformation_context));
+
+  // Bad: the OpKill being targeted is in the only available wrapper; we cannot
+  // have the wrapper call itself.
+  ASSERT_FALSE(TransformationWrapEarlyTerminatorInFunction(
+                   100, MakeInstructionDescriptor(31, SpvOpKill, 0), 0)
+                   .IsApplicable(context.get(), transformation_context));
+}
+
+TEST(TransformationWrapEarlyTerminatorInFunctionTest, Apply) {
+  std::string shader = R"(
+               OpCapability Shader
+               OpExtension "SPV_KHR_terminate_invocation"
+          %1 = OpExtInstImport "GLSL.std.450"
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint Fragment %4 "main"
+               OpExecutionMode %4 OriginUpperLeft
+               OpSource ESSL 320
+               OpName %4 "main"
+          %2 = OpTypeVoid
+          %3 = OpTypeFunction %2
+          %6 = OpTypeInt 32 1
+          %7 = OpConstant %6 0
+
+         %20 = OpTypeFunction %2 %6
+         %21 = OpTypeFunction %6
+
+          %4 = OpFunction %2 None %3
+          %5 = OpLabel
+               OpSelectionMerge %11 None
+               OpSwitch %7 %11 0 %8 1 %9 2 %10
+          %8 = OpLabel
+               OpKill
+          %9 = OpLabel
+               OpUnreachable
+         %10 = OpLabel
+               OpTerminateInvocation
+         %11 = OpLabel
+               OpReturn
+               OpFunctionEnd
+
+         %30 = OpFunction %2 None %3
+         %31 = OpLabel
+               OpKill
+               OpFunctionEnd
+
+         %40 = OpFunction %2 None %3
+         %41 = OpLabel
+               OpUnreachable
+               OpFunctionEnd
+
+         %50 = OpFunction %2 None %3
+         %51 = OpLabel
+               OpTerminateInvocation
+               OpFunctionEnd
+
+         %60 = OpFunction %2 None %3
+         %61 = OpLabel
+               OpBranch %62
+         %62 = OpLabel
+               OpKill
+               OpFunctionEnd
+
+         %70 = OpFunction %6 None %21
+         %71 = OpLabel
+               OpUnreachable
+               OpFunctionEnd
+
+         %80 = OpFunction %2 None %20
+         %81 = OpFunctionParameter %6
+         %82 = OpLabel
+               OpTerminateInvocation
+               OpFunctionEnd
+  )";
+
+  const auto env = SPV_ENV_UNIVERSAL_1_4;
+  const auto consumer = nullptr;
+  const auto context = BuildModule(env, consumer, shader, kFuzzAssembleOption);
+  ASSERT_TRUE(IsValid(env, context.get()));
+
+  spvtools::ValidatorOptions validator_options;
+  TransformationContext transformation_context(
+      MakeUnique<FactManager>(context.get()), validator_options);
+
+  for (auto& transformation :
+       {TransformationWrapEarlyTerminatorInFunction(
+            100, MakeInstructionDescriptor(8, SpvOpKill, 0), 0),
+        TransformationWrapEarlyTerminatorInFunction(
+            101, MakeInstructionDescriptor(9, SpvOpUnreachable, 0), 0),
+        TransformationWrapEarlyTerminatorInFunction(
+            102, MakeInstructionDescriptor(10, SpvOpTerminateInvocation, 0), 0),
+        TransformationWrapEarlyTerminatorInFunction(
+            103, MakeInstructionDescriptor(62, SpvOpKill, 0), 0),
+        TransformationWrapEarlyTerminatorInFunction(
+            104, MakeInstructionDescriptor(71, SpvOpUnreachable, 0), 7),
+        TransformationWrapEarlyTerminatorInFunction(
+            105, MakeInstructionDescriptor(82, SpvOpTerminateInvocation, 0),
+            0)}) {
+    ASSERT_TRUE(
+        transformation.IsApplicable(context.get(), transformation_context));
+    ApplyAndCheckFreshIds(transformation, context.get(),
+                          &transformation_context);
+  }
+  ASSERT_TRUE(IsValid(env, context.get()));
+
+  std::string after_transformation = R"(
+               OpCapability Shader
+               OpExtension "SPV_KHR_terminate_invocation"
+          %1 = OpExtInstImport "GLSL.std.450"
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint Fragment %4 "main"
+               OpExecutionMode %4 OriginUpperLeft
+               OpSource ESSL 320
+               OpName %4 "main"
+          %2 = OpTypeVoid
+          %3 = OpTypeFunction %2
+          %6 = OpTypeInt 32 1
+          %7 = OpConstant %6 0
+
+         %20 = OpTypeFunction %2 %6
+         %21 = OpTypeFunction %6
+
+          %4 = OpFunction %2 None %3
+          %5 = OpLabel
+               OpSelectionMerge %11 None
+               OpSwitch %7 %11 0 %8 1 %9 2 %10
+          %8 = OpLabel
+        %100 = OpFunctionCall %2 %30
+               OpReturn
+          %9 = OpLabel
+        %101 = OpFunctionCall %2 %40
+               OpReturn
+         %10 = OpLabel
+        %102 = OpFunctionCall %2 %50
+               OpReturn
+         %11 = OpLabel
+               OpReturn
+               OpFunctionEnd
+
+         %30 = OpFunction %2 None %3
+         %31 = OpLabel
+               OpKill
+               OpFunctionEnd
+
+         %40 = OpFunction %2 None %3
+         %41 = OpLabel
+               OpUnreachable
+               OpFunctionEnd
+
+         %50 = OpFunction %2 None %3
+         %51 = OpLabel
+               OpTerminateInvocation
+               OpFunctionEnd
+
+         %60 = OpFunction %2 None %3
+         %61 = OpLabel
+               OpBranch %62
+         %62 = OpLabel
+        %103 = OpFunctionCall %2 %30
+               OpReturn
+               OpFunctionEnd
+
+         %70 = OpFunction %6 None %21
+         %71 = OpLabel
+        %104 = OpFunctionCall %2 %40
+               OpReturnValue %7
+               OpFunctionEnd
+
+         %80 = OpFunction %2 None %20
+         %81 = OpFunctionParameter %6
+         %82 = OpLabel
+        %105 = OpFunctionCall %2 %50
+               OpReturn
+               OpFunctionEnd
+  )";
+  ASSERT_TRUE(IsEqual(env, after_transformation, context.get()));
+}
+
+}  // namespace
+}  // namespace fuzz
+}  // namespace spvtools