spirv-fuzz: Wrap early terminators before merging returns (#3925)

Incorporates some other fixes for issues that were brought to light by
adding this functionality.

Fixes #3717.
Fixes #3924.
diff --git a/source/fuzz/fact_manager/data_synonym_and_id_equation_facts.cpp b/source/fuzz/fact_manager/data_synonym_and_id_equation_facts.cpp
index 44a0851..04726a0 100644
--- a/source/fuzz/fact_manager/data_synonym_and_id_equation_facts.cpp
+++ b/source/fuzz/fact_manager/data_synonym_and_id_equation_facts.cpp
@@ -581,10 +581,18 @@
               synonymous_.IsEquivalent(dd1_prefix, dd2_prefix)) {
             continue;
           }
-
+          opt::Instruction* dd1_object =
+              ir_context_->get_def_use_mgr()->GetDef(dd1->object());
+          opt::Instruction* dd2_object =
+              ir_context_->get_def_use_mgr()->GetDef(dd2->object());
+          if (dd1_object == nullptr || dd2_object == nullptr) {
+            // The objects are not both available in the module, so we cannot
+            // investigate the types of the associated data descriptors; we need
+            // to move on.
+            continue;
+          }
           // Get the type of obj_1
-          auto dd1_root_type_id =
-              ir_context_->get_def_use_mgr()->GetDef(dd1->object())->type_id();
+          auto dd1_root_type_id = dd1_object->type_id();
           // Use this type, together with a_1, ..., a_m, to get the type of
           // obj_1[a_1, ..., a_m].
           auto dd1_prefix_type = fuzzerutil::WalkCompositeTypeIndices(
@@ -592,8 +600,7 @@
 
           // Similarly, get the type of obj_2 and use it to get the type of
           // obj_2[b_1, ..., b_n].
-          auto dd2_root_type_id =
-              ir_context_->get_def_use_mgr()->GetDef(dd2->object())->type_id();
+          auto dd2_root_type_id = dd2_object->type_id();
           auto dd2_prefix_type = fuzzerutil::WalkCompositeTypeIndices(
               ir_context_, dd2_root_type_id, dd2_prefix.index());
 
diff --git a/source/fuzz/fuzzer_pass.cpp b/source/fuzz/fuzzer_pass.cpp
index da6b9a5..486f18f 100644
--- a/source/fuzz/fuzzer_pass.cpp
+++ b/source/fuzz/fuzzer_pass.cpp
@@ -101,6 +101,72 @@
 }
 
 void FuzzerPass::ForEachInstructionWithInstructionDescriptor(
+    opt::Function* function,
+    std::function<
+        void(opt::BasicBlock* block, opt::BasicBlock::iterator inst_it,
+             const protobufs::InstructionDescriptor& instruction_descriptor)>
+        action) {
+  // Consider only reachable blocks. We do this in a separate loop to avoid
+  // recomputing the dominator analysis every time |action| changes the
+  // module.
+  std::vector<opt::BasicBlock*> reachable_blocks;
+
+  const auto* dominator_analysis =
+      GetIRContext()->GetDominatorAnalysis(function);
+  for (auto& block : *function) {
+    if (dominator_analysis->IsReachable(&block)) {
+      reachable_blocks.push_back(&block);
+    }
+  }
+
+  for (auto* block : reachable_blocks) {
+    // We now consider every instruction in the block, randomly deciding
+    // whether to apply a transformation before it.
+
+    // In order for transformations to insert new instructions, they need to
+    // be able to identify the instruction to insert before.  We describe an
+    // instruction via its opcode, 'opc', a base instruction 'base' that has a
+    // result id, and the number of instructions with opcode 'opc' that we
+    // should skip when searching from 'base' for the desired instruction.
+    // (An instruction that has a result id is represented by its own opcode,
+    // itself as 'base', and a skip-count of 0.)
+    std::vector<std::tuple<uint32_t, SpvOp, uint32_t>> base_opcode_skip_triples;
+
+    // The initial base instruction is the block label.
+    uint32_t base = block->id();
+
+    // Counts the number of times we have seen each opcode since we reset the
+    // base instruction.
+    std::map<SpvOp, uint32_t> skip_count;
+
+    // Consider every instruction in the block.  The label is excluded: it is
+    // only necessary to consider it as a base in case the first instruction
+    // in the block does not have a result id.
+    for (auto inst_it = block->begin(); inst_it != block->end(); ++inst_it) {
+      if (inst_it->HasResultId()) {
+        // In the case that the instruction has a result id, we use the
+        // instruction as its own base, and clear the skip counts we have
+        // collected.
+        base = inst_it->result_id();
+        skip_count.clear();
+      }
+      const SpvOp opcode = inst_it->opcode();
+
+      // Invoke the provided function, which might apply a transformation.
+      action(block, inst_it,
+             MakeInstructionDescriptor(
+                 base, opcode,
+                 skip_count.count(opcode) ? skip_count.at(opcode) : 0));
+
+      if (!inst_it->HasResultId()) {
+        skip_count[opcode] =
+            skip_count.count(opcode) ? skip_count.at(opcode) + 1 : 1;
+      }
+    }
+  }
+}
+
+void FuzzerPass::ForEachInstructionWithInstructionDescriptor(
     std::function<
         void(opt::Function* function, opt::BasicBlock* block,
              opt::BasicBlock::iterator inst_it,
@@ -108,65 +174,13 @@
         action) {
   // Consider every block in every function.
   for (auto& function : *GetIRContext()->module()) {
-    // Consider only reachable blocks. We do this in a separate loop to avoid
-    // recomputing the dominator analysis every time |action| changes the
-    // module.
-    std::vector<opt::BasicBlock*> reachable_blocks;
-
-    const auto* dominator_analysis =
-        GetIRContext()->GetDominatorAnalysis(&function);
-    for (auto& block : function) {
-      if (dominator_analysis->IsReachable(&block)) {
-        reachable_blocks.push_back(&block);
-      }
-    }
-
-    for (auto* block : reachable_blocks) {
-      // We now consider every instruction in the block, randomly deciding
-      // whether to apply a transformation before it.
-
-      // In order for transformations to insert new instructions, they need to
-      // be able to identify the instruction to insert before.  We describe an
-      // instruction via its opcode, 'opc', a base instruction 'base' that has a
-      // result id, and the number of instructions with opcode 'opc' that we
-      // should skip when searching from 'base' for the desired instruction.
-      // (An instruction that has a result id is represented by its own opcode,
-      // itself as 'base', and a skip-count of 0.)
-      std::vector<std::tuple<uint32_t, SpvOp, uint32_t>>
-          base_opcode_skip_triples;
-
-      // The initial base instruction is the block label.
-      uint32_t base = block->id();
-
-      // Counts the number of times we have seen each opcode since we reset the
-      // base instruction.
-      std::map<SpvOp, uint32_t> skip_count;
-
-      // Consider every instruction in the block.  The label is excluded: it is
-      // only necessary to consider it as a base in case the first instruction
-      // in the block does not have a result id.
-      for (auto inst_it = block->begin(); inst_it != block->end(); ++inst_it) {
-        if (inst_it->HasResultId()) {
-          // In the case that the instruction has a result id, we use the
-          // instruction as its own base, and clear the skip counts we have
-          // collected.
-          base = inst_it->result_id();
-          skip_count.clear();
-        }
-        const SpvOp opcode = inst_it->opcode();
-
-        // Invoke the provided function, which might apply a transformation.
-        action(&function, block, inst_it,
-               MakeInstructionDescriptor(
-                   base, opcode,
-                   skip_count.count(opcode) ? skip_count.at(opcode) : 0));
-
-        if (!inst_it->HasResultId()) {
-          skip_count[opcode] =
-              skip_count.count(opcode) ? skip_count.at(opcode) + 1 : 1;
-        }
-      }
-    }
+    ForEachInstructionWithInstructionDescriptor(
+        &function,
+        [&action, &function](
+            opt::BasicBlock* block, opt::BasicBlock::iterator inst_it,
+            const protobufs::InstructionDescriptor& instruction_descriptor) {
+          action(&function, block, inst_it, instruction_descriptor);
+        });
   }
 }
 
diff --git a/source/fuzz/fuzzer_pass.h b/source/fuzz/fuzzer_pass.h
index a9aae09..da3f239 100644
--- a/source/fuzz/fuzzer_pass.h
+++ b/source/fuzz/fuzzer_pass.h
@@ -71,8 +71,8 @@
           instruction_is_relevant) const;
 
   // A helper method that iterates through each instruction in each reachable
-  // block, at all times tracking an instruction descriptor that allows the
-  // latest instruction to be located even if it has no result id.
+  // block of |function|, at all times tracking an instruction descriptor that
+  // allows the latest instruction to be located even if it has no result id.
   //
   // The code to manipulate the instruction descriptor is a bit fiddly.  The
   // point of this method is to avoiding having to duplicate it in multiple
@@ -87,6 +87,17 @@
   // whether to try to apply some transformation, and then - if selected - to
   // attempt to apply it.
   void ForEachInstructionWithInstructionDescriptor(
+      opt::Function* function,
+      std::function<
+          void(opt::BasicBlock* block, opt::BasicBlock::iterator inst_it,
+               const protobufs::InstructionDescriptor& instruction_descriptor)>
+          action);
+
+  // Applies the above overload of ForEachInstructionWithInstructionDescriptor
+  // to every function in the module, so that |action| is applied to an
+  // |instruction_descriptor| for every instruction, |inst_it|, of every |block|
+  // in every |function|.
+  void ForEachInstructionWithInstructionDescriptor(
       std::function<
           void(opt::Function* function, opt::BasicBlock* block,
                opt::BasicBlock::iterator inst_it,
@@ -100,7 +111,12 @@
                                        *GetTransformationContext()) &&
            "Transformation should be applicable by construction.");
     transformation.Apply(GetIRContext(), GetTransformationContext());
-    *GetTransformations()->add_transformation() = transformation.ToMessage();
+    protobufs::Transformation transformation_message =
+        transformation.ToMessage();
+    assert(transformation_message.transformation_case() !=
+               protobufs::Transformation::TRANSFORMATION_NOT_SET &&
+           "Bad transformation.");
+    *GetTransformations()->add_transformation() = transformation_message;
   }
 
   // A generic helper for applying a transformation only if it is applicable.
@@ -111,7 +127,12 @@
     if (transformation.IsApplicable(GetIRContext(),
                                     *GetTransformationContext())) {
       transformation.Apply(GetIRContext(), GetTransformationContext());
-      *GetTransformations()->add_transformation() = transformation.ToMessage();
+      protobufs::Transformation transformation_message =
+          transformation.ToMessage();
+      assert(transformation_message.transformation_case() !=
+                 protobufs::Transformation::TRANSFORMATION_NOT_SET &&
+             "Bad transformation.");
+      *GetTransformations()->add_transformation() = transformation_message;
       return true;
     }
     return false;
diff --git a/source/fuzz/fuzzer_pass_merge_function_returns.cpp b/source/fuzz/fuzzer_pass_merge_function_returns.cpp
index 7130e43..4bcedb8 100644
--- a/source/fuzz/fuzzer_pass_merge_function_returns.cpp
+++ b/source/fuzz/fuzzer_pass_merge_function_returns.cpp
@@ -16,7 +16,9 @@
 
 #include "source/fuzz/fuzzer_util.h"
 #include "source/fuzz/instruction_descriptor.h"
+#include "source/fuzz/transformation_add_early_terminator_wrapper.h"
 #include "source/fuzz/transformation_merge_function_returns.h"
+#include "source/fuzz/transformation_wrap_early_terminator_in_function.h"
 
 namespace spvtools {
 namespace fuzz {
@@ -31,7 +33,15 @@
 FuzzerPassMergeFunctionReturns::~FuzzerPassMergeFunctionReturns() = default;
 
 void FuzzerPassMergeFunctionReturns::Apply() {
+  // The pass might add new functions to the module (due to wrapping early
+  // terminator instructions in function calls), so we record the functions that
+  // are currently present and then iterate over them.
+  std::vector<opt::Function*> functions;
   for (auto& function : *GetIRContext()->module()) {
+    functions.emplace_back(&function);
+  }
+
+  for (auto* function : functions) {
     // Randomly decide whether to consider this function.
     if (GetFuzzerContext()->ChoosePercentage(
             GetFuzzerContext()->GetChanceOfMergingFunctionReturns())) {
@@ -39,13 +49,66 @@
     }
 
     // Only consider functions that have early returns.
-    if (!function.HasEarlyReturn()) {
+    if (!function->HasEarlyReturn()) {
       continue;
     }
 
+    // Wrap early terminators in function calls.
+    ForEachInstructionWithInstructionDescriptor(
+        function,
+        [this, function](
+            opt::BasicBlock* /*unused*/, opt::BasicBlock::iterator inst_it,
+            const protobufs::InstructionDescriptor& instruction_descriptor) {
+          const SpvOp opcode = inst_it->opcode();
+          switch (opcode) {
+            case SpvOpKill:
+            case SpvOpUnreachable:
+            case SpvOpTerminateInvocation: {
+              // This is an early termination instruction - we need to wrap it
+              // so that it becomes a return.
+              if (TransformationWrapEarlyTerminatorInFunction::
+                      MaybeGetWrapperFunction(GetIRContext(), opcode) ==
+                  nullptr) {
+                // We don't have a suitable wrapper function, so create one.
+                ApplyTransformation(TransformationAddEarlyTerminatorWrapper(
+                    GetFuzzerContext()->GetFreshId(),
+                    GetFuzzerContext()->GetFreshId(), opcode));
+              }
+              // If the function has non-void return type then we need a
+              // suitable value to use in an OpReturnValue instruction.
+              opt::Instruction* function_return_type =
+                  GetIRContext()->get_def_use_mgr()->GetDef(
+                      function->type_id());
+              uint32_t returned_value_id;
+              if (function_return_type->opcode() == SpvOpTypeVoid) {
+                // No value is needed.
+                returned_value_id = 0;
+              } else if (fuzzerutil::CanCreateConstant(
+                             GetIRContext(),
+                             function_return_type->result_id())) {
+                // We favour returning an irrelevant zero.
+                returned_value_id = FindOrCreateZeroConstant(
+                    function_return_type->result_id(), true);
+              } else {
+                // It's not possible to use an irrelevant zero, so we use an
+                // OpUndef instead.
+                returned_value_id =
+                    FindOrCreateGlobalUndef(function_return_type->result_id());
+              }
+              // Wrap the early termination instruction in a function call.
+              ApplyTransformation(TransformationWrapEarlyTerminatorInFunction(
+                  GetFuzzerContext()->GetFreshId(), instruction_descriptor,
+                  returned_value_id));
+              break;
+            }
+            default:
+              break;
+          }
+        });
+
     // Get the return blocks.
     auto return_blocks = fuzzerutil::GetReachableReturnBlocks(
-        GetIRContext(), function.result_id());
+        GetIRContext(), function->result_id());
 
     // Only go ahead if there is more than one reachable return block.
     if (return_blocks.size() <= 1) {
@@ -58,12 +121,12 @@
 
     // Collect the ids available after the entry block of the function.
     auto ids_available_after_entry_block =
-        GetTypesToIdsAvailableAfterEntryBlock(&function);
+        GetTypesToIdsAvailableAfterEntryBlock(function);
 
     // If the entry block does not branch unconditionally to another block,
     // split it.
-    if (function.entry()->terminator()->opcode() != SpvOpBranch) {
-      SplitBlockAfterOpPhiOrOpVariable(function.entry()->id());
+    if (function->entry()->terminator()->opcode() != SpvOpBranch) {
+      SplitBlockAfterOpPhiOrOpVariable(function->entry()->id());
     }
 
     // Collect the merge blocks of the function whose corresponding loops
@@ -109,7 +172,7 @@
     uint32_t outer_return_id = GetFuzzerContext()->GetFreshId();
 
     bool function_is_void =
-        GetIRContext()->get_type_mgr()->GetType(function.type_id())->AsVoid();
+        GetIRContext()->get_type_mgr()->GetType(function->type_id())->AsVoid();
 
     // We only need a return value if the function is not void.
     uint32_t return_val_id =
@@ -120,17 +183,17 @@
     uint32_t returnable_val_id = 0;
     if (!function_is_void && !actual_merge_blocks.empty()) {
       // If there is an id of the suitable type, choose one at random.
-      if (ids_available_after_entry_block.count(function.type_id())) {
+      if (ids_available_after_entry_block.count(function->type_id())) {
         const auto& candidates =
-            ids_available_after_entry_block[function.type_id()];
+            ids_available_after_entry_block[function->type_id()];
         returnable_val_id =
             candidates[GetFuzzerContext()->RandomIndex(candidates)];
       } else {
         // If there is no id, add a global OpUndef.
-        uint32_t suitable_id = FindOrCreateGlobalUndef(function.type_id());
+        uint32_t suitable_id = FindOrCreateGlobalUndef(function->type_id());
         // Add the new id to the map of available ids.
         ids_available_after_entry_block.emplace(
-            function.type_id(), std::vector<uint32_t>({suitable_id}));
+            function->type_id(), std::vector<uint32_t>({suitable_id}));
         returnable_val_id = suitable_id;
       }
     }
@@ -142,7 +205,7 @@
     // Apply the transformation if it is applicable (it could be inapplicable if
     // adding new predecessors to merge blocks breaks dominance rules).
     MaybeApplyTransformation(TransformationMergeFunctionReturns(
-        function.result_id(), outer_header_id, outer_return_id, return_val_id,
+        function->result_id(), outer_header_id, outer_return_id, return_val_id,
         returnable_val_id, merge_blocks_info));
   }
 }
diff --git a/source/fuzz/fuzzer_pass_merge_function_returns.h b/source/fuzz/fuzzer_pass_merge_function_returns.h
index 0fa505c..2d79905 100644
--- a/source/fuzz/fuzzer_pass_merge_function_returns.h
+++ b/source/fuzz/fuzzer_pass_merge_function_returns.h
@@ -21,7 +21,11 @@
 namespace fuzz {
 
 // A fuzzer pass for changing functions in the module so that they don't have an
-// early return.
+// early return.  When handling a function the pass first eliminates early
+// terminator instructions, such as OpKill, by wrapping them in functions and
+// replacing them with a function call followed by a return.  The return
+// instructions that arise are then modified so that the function does not have
+// early returns.
 class FuzzerPassMergeFunctionReturns : public FuzzerPass {
  public:
   FuzzerPassMergeFunctionReturns(
diff --git a/source/fuzz/fuzzer_pass_propagate_instructions_up.cpp b/source/fuzz/fuzzer_pass_propagate_instructions_up.cpp
index 2042d7c..16ec680 100644
--- a/source/fuzz/fuzzer_pass_propagate_instructions_up.cpp
+++ b/source/fuzz/fuzzer_pass_propagate_instructions_up.cpp
@@ -52,7 +52,6 @@
             fresh_id = GetFuzzerContext()->GetFreshId();
           }
         }
-
         ApplyTransformation(
             TransformationPropagateInstructionUp(block.id(), fresh_ids));
       }
diff --git a/source/fuzz/protobufs/spvtoolsfuzz.proto b/source/fuzz/protobufs/spvtoolsfuzz.proto
index 35e0b9a..7778165 100644
--- a/source/fuzz/protobufs/spvtoolsfuzz.proto
+++ b/source/fuzz/protobufs/spvtoolsfuzz.proto
@@ -555,6 +555,7 @@
     TransformationPropagateInstructionDown propagate_instruction_down = 81;
     TransformationReplaceBranchFromDeadBlockWithExit replace_branch_from_dead_block_with_exit = 82;
     TransformationWrapEarlyTerminatorInFunction wrap_early_terminator_in_function = 83;
+    TransformationMergeFunctionReturns merge_function_returns = 84;
     // Add additional option using the next available number.
   }
 }
diff --git a/source/fuzz/transformation.cpp b/source/fuzz/transformation.cpp
index 357d59f..c8bf879 100644
--- a/source/fuzz/transformation.cpp
+++ b/source/fuzz/transformation.cpp
@@ -64,6 +64,7 @@
 #include "source/fuzz/transformation_load.h"
 #include "source/fuzz/transformation_make_vector_operation_dynamic.h"
 #include "source/fuzz/transformation_merge_blocks.h"
+#include "source/fuzz/transformation_merge_function_returns.h"
 #include "source/fuzz/transformation_move_block_down.h"
 #include "source/fuzz/transformation_move_instruction_down.h"
 #include "source/fuzz/transformation_mutate_pointer.h"
@@ -245,6 +246,9 @@
           message.make_vector_operation_dynamic());
     case protobufs::Transformation::TransformationCase::kMergeBlocks:
       return MakeUnique<TransformationMergeBlocks>(message.merge_blocks());
+    case protobufs::Transformation::TransformationCase::kMergeFunctionReturns:
+      return MakeUnique<TransformationMergeFunctionReturns>(
+          message.merge_function_returns());
     case protobufs::Transformation::TransformationCase::kMoveBlockDown:
       return MakeUnique<TransformationMoveBlockDown>(message.move_block_down());
     case protobufs::Transformation::TransformationCase::kMoveInstructionDown:
diff --git a/source/fuzz/transformation_add_loop_preheader.cpp b/source/fuzz/transformation_add_loop_preheader.cpp
index ef5bef2..3d50fa9 100644
--- a/source/fuzz/transformation_add_loop_preheader.cpp
+++ b/source/fuzz/transformation_add_loop_preheader.cpp
@@ -189,11 +189,10 @@
       // Update the OpPhi instruction in the header so that it refers to the
       // back edge block and the preheader as the predecessors, and it uses the
       // newly-defined OpPhi in the preheader for the corresponding value.
-      phi_inst->SetInOperands(
-          {{SPV_OPERAND_TYPE_RESULT_ID, {fresh_phi_id}},
-           {SPV_OPERAND_TYPE_RESULT_ID, {preheader->id()}},
-           {SPV_OPERAND_TYPE_RESULT_ID, {back_edge_val}},
-           {SPV_OPERAND_TYPE_RESULT_ID, {back_edge_block_id}}});
+      phi_inst->SetInOperands({{SPV_OPERAND_TYPE_ID, {fresh_phi_id}},
+                               {SPV_OPERAND_TYPE_ID, {preheader->id()}},
+                               {SPV_OPERAND_TYPE_ID, {back_edge_val}},
+                               {SPV_OPERAND_TYPE_ID, {back_edge_block_id}}});
     }
   });
 
@@ -201,11 +200,11 @@
   fuzzerutil::UpdateModuleIdBound(ir_context, message_.fresh_id());
 
   // Add an unconditional branch from the preheader to the header.
-  preheader->AddInstruction(std::unique_ptr<opt::Instruction>(
-      new opt::Instruction(ir_context, SpvOpBranch, 0, 0,
-                           std::initializer_list<opt::Operand>{opt::Operand(
-                               spv_operand_type_t::SPV_OPERAND_TYPE_RESULT_ID,
-                               {loop_header->id()})})));
+  preheader->AddInstruction(
+      std::unique_ptr<opt::Instruction>(new opt::Instruction(
+          ir_context, SpvOpBranch, 0, 0,
+          std::initializer_list<opt::Operand>{opt::Operand(
+              spv_operand_type_t::SPV_OPERAND_TYPE_ID, {loop_header->id()})})));
 
   // Insert the preheader in the module.
   loop_header->GetParent()->InsertBasicBlockBefore(std::move(preheader),
diff --git a/source/fuzz/transformation_merge_function_returns.cpp b/source/fuzz/transformation_merge_function_returns.cpp
index d215de1..bdaec04 100644
--- a/source/fuzz/transformation_merge_function_returns.cpp
+++ b/source/fuzz/transformation_merge_function_returns.cpp
@@ -284,7 +284,7 @@
     // Replace the return instruction with an unconditional branch.
     ret_block->terminator()->SetOpcode(SpvOpBranch);
     ret_block->terminator()->SetInOperands(
-        {{SPV_OPERAND_TYPE_RESULT_ID, {merge_block_id}}});
+        {{SPV_OPERAND_TYPE_ID, {merge_block_id}}});
   }
 
   // Get a list of all the relevant merge blocks.
@@ -345,32 +345,33 @@
     auto merge_block = ir_context->get_instr_block(merge_block_id);
 
     // Adjust the existing OpPhi instructions.
-    merge_block->ForEachPhiInst([&preds, &returning_preds, &phi_to_id,
-                                 &types_to_available_ids](
-                                    opt::Instruction* inst) {
-      // We need a placeholder value id. If |phi_to_id| contains a mapping
-      // for this instruction, we use the given id, otherwise a suitable id
-      // for the instruction's type from |types_to_available_ids|.
-      uint32_t placeholder_val_id =
-          phi_to_id.count(inst->result_id())
-              ? phi_to_id[inst->result_id()]
-              : types_to_available_ids[inst->type_id()];
-      assert(placeholder_val_id &&
-             "We should always be able to find a suitable if the "
-             "transformation is applicable.");
+    merge_block->ForEachPhiInst(
+        [&preds, &returning_preds, &phi_to_id,
+         &types_to_available_ids](opt::Instruction* inst) {
+          // We need a placeholder value id. If |phi_to_id| contains a mapping
+          // for this instruction, we use the given id, otherwise a suitable id
+          // for the instruction's type from |types_to_available_ids|.
+          uint32_t placeholder_val_id =
+              phi_to_id.count(inst->result_id())
+                  ? phi_to_id[inst->result_id()]
+                  : types_to_available_ids[inst->type_id()];
+          assert(placeholder_val_id &&
+                 "We should always be able to find a suitable if the "
+                 "transformation is applicable.");
 
-      // Add a pair of operands (placeholder id, new predecessor) for each
-      // new predecessor of the merge block.
-      for (const auto& entry : returning_preds) {
-        // A returning predecessor may already be a predecessor of the
-        // block. In that case, we should not add new operands.
-        // Each entry is in the form (predecessor, {return val, is returning}).
-        if (!preds.count(entry.first)) {
-          inst->AddOperand({SPV_OPERAND_TYPE_RESULT_ID, {placeholder_val_id}});
-          inst->AddOperand({SPV_OPERAND_TYPE_RESULT_ID, {entry.first}});
-        }
-      }
-    });
+          // Add a pair of operands (placeholder id, new predecessor) for each
+          // new predecessor of the merge block.
+          for (const auto& entry : returning_preds) {
+            // A returning predecessor may already be a predecessor of the
+            // block. In that case, we should not add new operands.
+            // Each entry is in the form (predecessor, {return val, is
+            // returning}).
+            if (!preds.count(entry.first)) {
+              inst->AddOperand({SPV_OPERAND_TYPE_ID, {placeholder_val_id}});
+              inst->AddOperand({SPV_OPERAND_TYPE_ID, {entry.first}});
+            }
+          }
+        });
 
     // If the function is not void, add a new OpPhi instructions to collect the
     // return value from the returning predecessors.
@@ -383,9 +384,9 @@
         // Each entry is in the form (predecessor, {return value,
         // is returning}).
         operand_list.emplace_back(
-            opt::Operand{SPV_OPERAND_TYPE_RESULT_ID, {entry.second.first}});
+            opt::Operand{SPV_OPERAND_TYPE_ID, {entry.second.first}});
         operand_list.emplace_back(
-            opt::Operand{SPV_OPERAND_TYPE_RESULT_ID, {entry.first}});
+            opt::Operand{SPV_OPERAND_TYPE_ID, {entry.first}});
       }
 
       // Add two operands for each original predecessor from which the function
@@ -398,9 +399,9 @@
         }
 
         operand_list.emplace_back(
-            opt::Operand{SPV_OPERAND_TYPE_RESULT_ID, {returnable_val_id}});
+            opt::Operand{SPV_OPERAND_TYPE_ID, {returnable_val_id}});
         operand_list.emplace_back(
-            opt::Operand{SPV_OPERAND_TYPE_RESULT_ID, {original_pred}});
+            opt::Operand{SPV_OPERAND_TYPE_ID, {original_pred}});
       }
 
       // Insert the instruction.
@@ -421,9 +422,9 @@
         // Each entry is in the form (predecessor, {return value,
         // is returning}).
         operand_list.emplace_back(
-            opt::Operand{SPV_OPERAND_TYPE_RESULT_ID, {entry.second.second}});
+            opt::Operand{SPV_OPERAND_TYPE_ID, {entry.second.second}});
         operand_list.emplace_back(
-            opt::Operand{SPV_OPERAND_TYPE_RESULT_ID, {entry.first}});
+            opt::Operand{SPV_OPERAND_TYPE_ID, {entry.first}});
       }
 
       // Add two operands for each original predecessor from which the function
@@ -436,9 +437,9 @@
         }
 
         operand_list.emplace_back(
-            opt::Operand{SPV_OPERAND_TYPE_RESULT_ID, {constant_false}});
+            opt::Operand{SPV_OPERAND_TYPE_ID, {constant_false}});
         operand_list.emplace_back(
-            opt::Operand{SPV_OPERAND_TYPE_RESULT_ID, {original_pred}});
+            opt::Operand{SPV_OPERAND_TYPE_ID, {original_pred}});
       }
 
       // Insert the instruction.
@@ -480,9 +481,9 @@
     // true, to |original_succ| otherwise.
     merge_block->terminator()->SetOpcode(SpvOpBranchConditional);
     merge_block->terminator()->SetInOperands(
-        {{SPV_OPERAND_TYPE_RESULT_ID, {is_returning_id}},
-         {SPV_OPERAND_TYPE_RESULT_ID, {enclosing_merge}},
-         {SPV_OPERAND_TYPE_RESULT_ID, {original_succ}}});
+        {{SPV_OPERAND_TYPE_ID, {is_returning_id}},
+         {SPV_OPERAND_TYPE_ID, {enclosing_merge}},
+         {SPV_OPERAND_TYPE_ID, {original_succ}}});
   }
 
   assert(function->entry()->terminator()->opcode() == SpvOpBranch &&
@@ -503,8 +504,8 @@
   outer_loop_header->AddInstruction(MakeUnique<opt::Instruction>(
       ir_context, SpvOpLoopMerge, 0, 0,
       opt::Instruction::OperandList{
-          {SPV_OPERAND_TYPE_RESULT_ID, {message_.outer_return_id()}},
-          {SPV_OPERAND_TYPE_RESULT_ID, {message_.outer_header_id()}},
+          {SPV_OPERAND_TYPE_ID, {message_.outer_return_id()}},
+          {SPV_OPERAND_TYPE_ID, {message_.outer_header_id()}},
           {SPV_OPERAND_TYPE_LOOP_CONTROL, {SpvLoopControlMaskNone}}}));
 
   // Add conditional branch:
@@ -514,8 +515,8 @@
   outer_loop_header->AddInstruction(MakeUnique<opt::Instruction>(
       ir_context, SpvOpBranchConditional, 0, 0,
       opt::Instruction::OperandList{
-          {SPV_OPERAND_TYPE_RESULT_ID, {constant_true}},
-          {SPV_OPERAND_TYPE_RESULT_ID, {block_after_entry}},
+          {SPV_OPERAND_TYPE_ID, {constant_true}},
+          {SPV_OPERAND_TYPE_ID, {block_after_entry}},
           {SPV_OPERAND_TYPE_LOOP_CONTROL, {message_.outer_header_id()}}}));
 
   // Insert the header right after the entry block.
@@ -524,7 +525,18 @@
 
   // Update the branching instruction of the entry block.
   function->entry()->terminator()->SetInOperands(
-      {{SPV_OPERAND_TYPE_RESULT_ID, {message_.outer_header_id()}}});
+      {{SPV_OPERAND_TYPE_ID, {message_.outer_header_id()}}});
+
+  // If the entry block is referenced in an OpPhi instruction, the header for
+  // the new loop should be referenced instead.
+  ir_context->get_def_use_mgr()->ForEachUse(
+      function->entry()->id(),
+      [this](opt::Instruction* use_instruction, uint32_t use_operand_index) {
+        if (use_instruction->opcode() == SpvOpPhi) {
+          use_instruction->SetOperand(use_operand_index,
+                                      {message_.outer_header_id()});
+        }
+      });
 
   // Create the merge block for the loop (and return block for the function).
   auto outer_return_block =
@@ -543,9 +555,9 @@
     for (auto entry : outer_merge_predecessors) {
       // Each entry is in the form (predecessor, return value).
       operand_list.emplace_back(
-          opt::Operand{SPV_OPERAND_TYPE_RESULT_ID, {entry.second}});
+          opt::Operand{SPV_OPERAND_TYPE_ID, {entry.second}});
       operand_list.emplace_back(
-          opt::Operand{SPV_OPERAND_TYPE_RESULT_ID, {entry.first}});
+          opt::Operand{SPV_OPERAND_TYPE_ID, {entry.first}});
     }
 
     // Insert the OpPhi instruction.
@@ -559,7 +571,7 @@
     outer_return_block->AddInstruction(MakeUnique<opt::Instruction>(
         ir_context, SpvOpReturnValue, 0, 0,
         opt::Instruction::OperandList{
-            {SPV_OPERAND_TYPE_RESULT_ID, {message_.return_val_id()}}}));
+            {SPV_OPERAND_TYPE_ID, {message_.return_val_id()}}}));
   } else {
     // Insert an OpReturn instruction (the function is void).
     outer_return_block->AddInstruction(MakeUnique<opt::Instruction>(
@@ -598,7 +610,9 @@
 
 protobufs::Transformation TransformationMergeFunctionReturns::ToMessage()
     const {
-  return protobufs::Transformation();
+  protobufs::Transformation result;
+  *result.mutable_merge_function_returns() = message_;
+  return result;
 }
 
 std::map<uint32_t, protobufs::ReturnMergingInfo>
diff --git a/source/fuzz/transformation_wrap_early_terminator_in_function.cpp b/source/fuzz/transformation_wrap_early_terminator_in_function.cpp
index b8370c8..4c436f5 100644
--- a/source/fuzz/transformation_wrap_early_terminator_in_function.cpp
+++ b/source/fuzz/transformation_wrap_early_terminator_in_function.cpp
@@ -45,7 +45,7 @@
     return false;
   }
 
-  // |message_.early_terminator_instruction| must identify an instruciton, and
+  // |message_.early_terminator_instruction| must identify an instruction, and
   // the instruction must indeed be an early terminator.
   auto early_terminator =
       FindInstruction(message_.early_terminator_instruction(), ir_context);
diff --git a/source/fuzz/transformation_wrap_early_terminator_in_function.h b/source/fuzz/transformation_wrap_early_terminator_in_function.h
index 2629c50..00151d4 100644
--- a/source/fuzz/transformation_wrap_early_terminator_in_function.h
+++ b/source/fuzz/transformation_wrap_early_terminator_in_function.h
@@ -33,12 +33,26 @@
       const protobufs::InstructionDescriptor& early_terminator_instruction,
       uint32_t returned_value_id);
 
-  // TODO comment
+  // - |message_.fresh_id| must be fresh.
+  // - |message_.early_terminator_instruction| must identify an early terminator
+  //   instruction, i.e. an instruction with opcode OpKill, OpUnreachable or
+  //   OpTerminateInvocation.
+  // - A suitable wrapper function for the early terminator must exist, and it
+  //   must be distinct from the function containing
+  //   |message_.early_terminator_instruction|.
+  // - If the enclosing function has non-void return type then
+  //   |message_.returned_value_instruction| must be the id of an instruction of
+  //   the return type that is available at the point of the early terminator
+  //   instruction.
   bool IsApplicable(
       opt::IRContext* ir_context,
       const TransformationContext& transformation_context) const override;
 
-  // TODO comment
+  // An OpFunctionCall instruction to an appropriate wrapper function is
+  // inserted before |message_.early_terminator_instruction|, and
+  // |message_.early_terminator_instruction| is replaced with either OpReturn
+  // or OpReturnValue |message_.returned_value_instruction| depending on whether
+  // the enclosing function's return type is void.
   void Apply(opt::IRContext* ir_context,
              TransformationContext* transformation_context) const override;
 
@@ -46,10 +60,10 @@
 
   protobufs::Transformation ToMessage() const override;
 
- private:
   static opt::Function* MaybeGetWrapperFunction(opt::IRContext* ir_context,
                                                 SpvOp early_terminator_opcode);
 
+ private:
   protobufs::TransformationWrapEarlyTerminatorInFunction message_;
 };
 
diff --git a/test/fuzz/fuzzer_replayer_test.cpp b/test/fuzz/fuzzer_replayer_test.cpp
index 57e8c20..3925e0b 100644
--- a/test/fuzz/fuzzer_replayer_test.cpp
+++ b/test/fuzz/fuzzer_replayer_test.cpp
@@ -1650,9 +1650,10 @@
     // Every 4th time we run the fuzzer, enable all fuzzer passes.
     bool enable_all_passes = (seed % 4) == 0;
     auto fuzzer_result =
-        Fuzzer(env, kSilentConsumer, binary_in, initial_facts, donor_suppliers,
-               MakeUnique<PseudoRandomGenerator>(seed), enable_all_passes,
-               strategies[strategy_index], true, validator_options)
+        Fuzzer(env, kConsoleMessageConsumer, binary_in, initial_facts,
+               donor_suppliers, MakeUnique<PseudoRandomGenerator>(seed),
+               enable_all_passes, strategies[strategy_index], true,
+               validator_options)
             .Run();
 
     // Cycle the repeated pass strategy so that we try a different one next time
diff --git a/test/fuzz/transformation_compute_data_synonym_fact_closure_test.cpp b/test/fuzz/transformation_compute_data_synonym_fact_closure_test.cpp
index df496c5..1fc66e2 100644
--- a/test/fuzz/transformation_compute_data_synonym_fact_closure_test.cpp
+++ b/test/fuzz/transformation_compute_data_synonym_fact_closure_test.cpp
@@ -371,6 +371,104 @@
       MakeDataDescriptor(40, {2, 1, 1}), MakeDataDescriptor(108, {2, 1, 1})));
 }
 
+TEST(TransformationComputeDataSynonymFactClosureTest,
+     ComputeClosureWithMissingIds) {
+  std::string shader = R"(
+               OpCapability Shader
+          %1 = OpExtInstImport "GLSL.std.450"
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint Fragment %12 "main"
+               OpExecutionMode %12 OriginUpperLeft
+               OpSource ESSL 310
+          %2 = OpTypeVoid
+          %3 = OpTypeFunction %2
+          %6 = OpTypeInt 32 1
+          %7 = OpTypeVector %6 4
+         %15 = OpConstant %6 24
+         %16 = OpConstantComposite %7 %15 %15 %15 %15
+         %17 = OpConstantComposite %7 %15 %15 %15 %15
+         %18 = OpTypeStruct %7
+         %19 = OpConstantComposite %18 %16
+         %30 = OpConstantComposite %18 %17
+         %12 = OpFunction %2 None %3
+         %13 = OpLabel
+         %50 = OpCopyObject %7 %16
+         %51 = OpCopyObject %7 %17
+         %20 = OpCopyObject %6 %15
+         %21 = OpCopyObject %6 %15
+         %22 = OpCopyObject %6 %15
+               OpReturn
+               OpFunctionEnd
+  )";
+
+  const auto env = SPV_ENV_UNIVERSAL_1_3;
+  const auto consumer = nullptr;
+  const auto context = BuildModule(env, consumer, shader, kFuzzAssembleOption);
+  ASSERT_TRUE(IsValid(env, context.get()));
+
+  ValidatorOptions validator_options;
+  TransformationContext transformation_context(
+      MakeUnique<FactManager>(context.get()), validator_options);
+
+  transformation_context.GetFactManager()->AddFactDataSynonym(
+      MakeDataDescriptor(20, {}), MakeDataDescriptor(15, {}));
+  transformation_context.GetFactManager()->AddFactDataSynonym(
+      MakeDataDescriptor(21, {}), MakeDataDescriptor(15, {}));
+  transformation_context.GetFactManager()->AddFactDataSynonym(
+      MakeDataDescriptor(22, {}), MakeDataDescriptor(15, {}));
+
+  transformation_context.GetFactManager()->AddFactDataSynonym(
+      MakeDataDescriptor(17, {0}), MakeDataDescriptor(15, {}));
+  transformation_context.GetFactManager()->AddFactDataSynonym(
+      MakeDataDescriptor(17, {1}), MakeDataDescriptor(15, {}));
+  transformation_context.GetFactManager()->AddFactDataSynonym(
+      MakeDataDescriptor(17, {2}), MakeDataDescriptor(15, {}));
+  transformation_context.GetFactManager()->AddFactDataSynonym(
+      MakeDataDescriptor(17, {3}), MakeDataDescriptor(15, {}));
+
+  transformation_context.GetFactManager()->AddFactDataSynonym(
+      MakeDataDescriptor(16, {0}), MakeDataDescriptor(20, {}));
+  transformation_context.GetFactManager()->AddFactDataSynonym(
+      MakeDataDescriptor(16, {1}), MakeDataDescriptor(21, {}));
+  transformation_context.GetFactManager()->AddFactDataSynonym(
+      MakeDataDescriptor(16, {2}), MakeDataDescriptor(22, {}));
+  transformation_context.GetFactManager()->AddFactDataSynonym(
+      MakeDataDescriptor(16, {3}), MakeDataDescriptor(15, {}));
+
+  transformation_context.GetFactManager()->AddFactDataSynonym(
+      MakeDataDescriptor(51, {0}), MakeDataDescriptor(15, {}));
+  transformation_context.GetFactManager()->AddFactDataSynonym(
+      MakeDataDescriptor(51, {1}), MakeDataDescriptor(15, {}));
+  transformation_context.GetFactManager()->AddFactDataSynonym(
+      MakeDataDescriptor(51, {2}), MakeDataDescriptor(15, {}));
+  transformation_context.GetFactManager()->AddFactDataSynonym(
+      MakeDataDescriptor(51, {3}), MakeDataDescriptor(15, {}));
+
+  transformation_context.GetFactManager()->AddFactDataSynonym(
+      MakeDataDescriptor(50, {0}), MakeDataDescriptor(20, {}));
+  transformation_context.GetFactManager()->AddFactDataSynonym(
+      MakeDataDescriptor(50, {1}), MakeDataDescriptor(21, {}));
+  transformation_context.GetFactManager()->AddFactDataSynonym(
+      MakeDataDescriptor(50, {2}), MakeDataDescriptor(22, {}));
+  transformation_context.GetFactManager()->AddFactDataSynonym(
+      MakeDataDescriptor(50, {3}), MakeDataDescriptor(15, {}));
+
+  ASSERT_FALSE(transformation_context.GetFactManager()->IsSynonymous(
+      MakeDataDescriptor(19, {}), MakeDataDescriptor(30, {})));
+
+  context->KillDef(20);
+  context->KillDef(21);
+  context->KillDef(22);
+  context->KillDef(50);
+  context->KillDef(51);
+  context->InvalidateAnalysesExceptFor(opt::IRContext::kAnalysisNone);
+
+  ApplyAndCheckFreshIds(TransformationComputeDataSynonymFactClosure(100),
+                        context.get(), &transformation_context);
+  ASSERT_FALSE(transformation_context.GetFactManager()->IsSynonymous(
+      MakeDataDescriptor(19, {}), MakeDataDescriptor(30, {})));
+}
+
 }  // namespace
 }  // namespace fuzz
 }  // namespace spvtools
diff --git a/test/fuzz/transformation_merge_function_returns_test.cpp b/test/fuzz/transformation_merge_function_returns_test.cpp
index 2c8d949..330bc97 100644
--- a/test/fuzz/transformation_merge_function_returns_test.cpp
+++ b/test/fuzz/transformation_merge_function_returns_test.cpp
@@ -25,7 +25,7 @@
 protobufs::ReturnMergingInfo MakeReturnMergingInfo(
     uint32_t merge_block_id, uint32_t is_returning_id,
     uint32_t maybe_return_val_id,
-    std::map<uint32_t, uint32_t> opphi_to_suitable_id) {
+    const std::map<uint32_t, uint32_t>& opphi_to_suitable_id) {
   protobufs::ReturnMergingInfo result;
   result.set_merge_block_id(merge_block_id);
   result.set_is_returning_id(is_returning_id);
@@ -1779,6 +1779,84 @@
 
   ASSERT_TRUE(IsEqual(env, after_transformation, context.get()));
 }
+
+TEST(TransformationMergeFunctionReturnsTest, OpPhiAfterFirstBlock) {
+  std::string shader = R"(
+               OpCapability Shader
+          %1 = OpExtInstImport "GLSL.std.450"
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint Fragment %2 "main"
+               OpExecutionMode %2 OriginUpperLeft
+               OpSource ESSL 310
+          %3 = OpTypeVoid
+          %4 = OpTypeFunction %3
+          %5 = OpTypeBool
+          %6 = OpConstantTrue %5
+          %7 = OpConstantFalse %5
+          %2 = OpFunction %3 None %4
+          %8 = OpLabel
+               OpBranch %9
+          %9 = OpLabel
+         %10 = OpPhi %5 %6 %8
+               OpSelectionMerge %11 None
+               OpBranchConditional %6 %12 %11
+         %12 = OpLabel
+               OpReturn
+         %11 = OpLabel
+               OpReturn
+               OpFunctionEnd
+)";
+
+  const auto env = SPV_ENV_UNIVERSAL_1_5;
+  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);
+
+  auto transformation =
+      TransformationMergeFunctionReturns(2, 100, 101, 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
+          %1 = OpExtInstImport "GLSL.std.450"
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint Fragment %2 "main"
+               OpExecutionMode %2 OriginUpperLeft
+               OpSource ESSL 310
+          %3 = OpTypeVoid
+          %4 = OpTypeFunction %3
+          %5 = OpTypeBool
+          %6 = OpConstantTrue %5
+          %7 = OpConstantFalse %5
+          %2 = OpFunction %3 None %4
+          %8 = OpLabel
+               OpBranch %100
+        %100 = OpLabel
+               OpLoopMerge %101 %100 None
+               OpBranchConditional %6 %9 %100
+          %9 = OpLabel
+         %10 = OpPhi %5 %6 %100
+               OpSelectionMerge %11 None
+               OpBranchConditional %6 %12 %11
+         %12 = OpLabel
+               OpBranch %101
+         %11 = OpLabel
+               OpBranch %101
+        %101 = OpLabel
+               OpReturn
+               OpFunctionEnd
+)";
+
+  ASSERT_TRUE(IsEqual(env, after_transformation, context.get()));
+}
+
 }  // namespace
 }  // namespace fuzz
 }  // namespace spvtools
\ No newline at end of file