Fix detection of blocks bypassed by new edge (#2874)

Fixes an issue where the blocks that would be bypassed by a new break
or continue control flow edge were not properly detected.

Fixes #2871.

2 files changed, 205 insertions, 46 deletions

diff --git a/source/fuzz/fuzzer_util.cpp b/source/fuzz/fuzzer_util.cpp
index 4cf3f205..8173d0f2 100644
--- a/source/fuzz/fuzzer_util.cpp
+++ b/source/fuzz/fuzzer_util.cpp
@@ -205,6 +205,7 @@ opt::BasicBlock::iterator GetIteratorForBaseInstructionAndOffset(
   return nullptr;
 }
 
+// Returns the ids of all successors of |block|
 std::vector<uint32_t> GetSuccessors(opt::BasicBlock* block) {
   std::vector<uint32_t> result;
   switch (block->terminator()->opcode()) {
@@ -226,6 +227,41 @@ std::vector<uint32_t> GetSuccessors(opt::BasicBlock* block) {
   return result;
 }
 
+// The FindBypassedBlocks method and its helpers perform a depth-first search;
+// this struct represents an element of the stack used during depth-first
+// search.
+struct FindBypassedBlocksDfsStackNode {
+  opt::BasicBlock* block;  // The block that is being explored
+  bool handled_merge;  // We visit merge blocks before successors; this field
+                       // tracks whether we have yet processed the merge block
+                       // (if any) associated with the block
+  uint32_t next_successor;  // The next as-yet unexplored successor of this
+                            // block; exploration of a block is complete when
+                            // this field's value reaches the successor count
+};
+
+// Helper method for the depth-first-search routine that collects blocks that a
+// new break or continue control flow graph edge will bypass.
+void HandleSuccessorDuringSearchForBypassedBlocks(
+    opt::BasicBlock* successor, bool new_blocks_will_be_bypassed,
+    std::set<uint32_t>* already_visited,
+    std::set<opt::BasicBlock*>* bypassed_blocks,
+    std::vector<FindBypassedBlocksDfsStackNode>* dfs_stack) {
+  if (already_visited->count(successor->id()) == 0) {
+    // This is a new block; mark it as visited so that we don't regard it as new
+    // in the future, and push it on to the stack for exploration.
+    already_visited->insert(successor->id());
+    dfs_stack->push_back({successor, false, 0});
+    if (new_blocks_will_be_bypassed) {
+      // We are in the region of the control-flow graph consisting of blocks
+      // that the new edge will bypass, so grab this block.
+      bypassed_blocks->insert(successor);
+    }
+  }
+}
+
+// Determines those block that will be bypassed by a break or continue edge from
+// |bb_from| to |bb_to|.
 void FindBypassedBlocks(opt::IRContext* context, opt::BasicBlock* bb_from,
                         opt::BasicBlock* bb_to,
                         std::set<opt::BasicBlock*>* bypassed_blocks) {
@@ -240,51 +276,40 @@ void FindBypassedBlocks(opt::IRContext* context, opt::BasicBlock* bb_from,
   // visited in the sub-search rooted at |bb_from|. (As an optimization, the
   // search terminates as soon as exploration of |bb_from| has completed.)
 
-  // This represents a basic block in a partial state of exploration.  As we
-  // wish to visit merge blocks in advance of regular successors, we track them
-  // separately.
-  struct StackNode {
-    opt::BasicBlock* block;
-    bool handled_merge;
-    std::vector<uint32_t> successors;
-    uint32_t next_successor;
-  };
-
   auto enclosing_function = bb_from->GetParent();
 
   // The set of block ids already visited during search.  We put |bb_to| in
   // there initially so that search automatically backtracks when this block is
   // reached.
-  std::set<uint32_t> visited;
-  visited.insert(bb_to->id());
+  std::set<uint32_t> already_visited;
+  already_visited.insert(bb_to->id());
 
-  // Tracks when we are in the region of blocks that are to be grabbed; we flip
-  // this to 'true' once we reach |bb_from| and have finished searching its
-  // merge block (in the case that it happens to be a header.
-  bool interested = false;
+  // Tracks when we are in the region of blocks that the new edge would bypass;
+  // we flip this to 'true' once we reach |bb_from| and have finished searching
+  // its merge block (in the case that it happens to be a header.
+  bool new_blocks_will_be_bypassed = false;
 
-  std::vector<StackNode> dfs_stack;
+  std::vector<FindBypassedBlocksDfsStackNode> dfs_stack;
   opt::BasicBlock* entry_block = enclosing_function->entry().get();
-  dfs_stack.push_back({entry_block, false, GetSuccessors(entry_block), 0});
+  dfs_stack.push_back({entry_block, false, 0});
   while (!dfs_stack.empty()) {
-    StackNode* node = &dfs_stack.back();
+    auto node_index = dfs_stack.size() - 1;
 
     // First make sure we search the merge block associated ith this block, if
     // there is one.
-    if (!node->handled_merge) {
-      node->handled_merge = true;
-      if (node->block->MergeBlockIdIfAny()) {
-        opt::BasicBlock* merge_block =
-            context->cfg()->block(node->block->MergeBlockIdIfAny());
+    if (!dfs_stack[node_index].handled_merge) {
+      dfs_stack[node_index].handled_merge = true;
+      if (dfs_stack[node_index].block->MergeBlockIdIfAny()) {
+        opt::BasicBlock* merge_block = context->cfg()->block(
+            dfs_stack[node_index].block->MergeBlockIdIfAny());
         // A block can only be the merge block for one header, so this block
         // should only be in |visited| if it is |bb_to|, which we put into
         // |visited| in advance.
-        assert(visited.count(merge_block->id()) == 0 || merge_block == bb_to);
-        if (visited.count(merge_block->id()) == 0) {
-          visited.insert(merge_block->id());
-          dfs_stack.push_back(
-              {merge_block, false, GetSuccessors(merge_block), 0});
-        }
+        assert(already_visited.count(merge_block->id()) == 0 ||
+               merge_block == bb_to);
+        HandleSuccessorDuringSearchForBypassedBlocks(
+            merge_block, new_blocks_will_be_bypassed, &already_visited,
+            bypassed_blocks, &dfs_stack);
       }
       continue;
     }
@@ -292,28 +317,23 @@ void FindBypassedBlocks(opt::IRContext* context, opt::BasicBlock* bb_from,
     // If we find |bb_from|, we are interested in grabbing previously unseen
     // successor blocks (by this point we will have already searched the merge
     // block associated with |bb_from|, if there is one.
-    if (node->block == bb_from) {
-      interested = true;
+    if (dfs_stack[node_index].block == bb_from) {
+      new_blocks_will_be_bypassed = true;
     }
 
     // Consider the next unexplored successor.
-    if (node->next_successor < node->successors.size()) {
-      uint32_t successor_id = node->successors[node->next_successor];
-      if (visited.count(successor_id) == 0) {
-        visited.insert(successor_id);
-        opt::BasicBlock* successor_block = context->cfg()->block(successor_id);
-        if (interested) {
-          // If we're in the region of interest, grab this block.
-          bypassed_blocks->insert(successor_block);
-        }
-        dfs_stack.push_back(
-            {successor_block, false, GetSuccessors(successor_block), 0});
-      }
-      node->next_successor++;
+    auto successors = GetSuccessors(dfs_stack[node_index].block);
+    if (dfs_stack[node_index].next_successor < successors.size()) {
+      HandleSuccessorDuringSearchForBypassedBlocks(
+          context->cfg()->block(
+              successors[dfs_stack[node_index].next_successor]),
+          new_blocks_will_be_bypassed, &already_visited, bypassed_blocks,
+          &dfs_stack);
+      dfs_stack[node_index].next_successor++;
     } else {
       // We have finished exploring |node|.  If it is |bb_from|, we can
       // terminate search -- we have grabbed all the relevant blocks.
-      if (node->block == bb_from) {
+      if (dfs_stack[node_index].block == bb_from) {
         break;
       }
       dfs_stack.pop_back();
diff --git a/test/fuzz/transformation_add_dead_continue_test.cpp b/test/fuzz/transformation_add_dead_continue_test.cpp
index e28af1dc..9f22679a 100644
--- a/test/fuzz/transformation_add_dead_continue_test.cpp
+++ b/test/fuzz/transformation_add_dead_continue_test.cpp
@@ -1249,6 +1249,145 @@ TEST(TransformationAddDeadContinueTest, RespectDominanceRules3) {
   ASSERT_FALSE(bad_transformation.IsApplicable(context.get(), fact_manager));
 }
 
+TEST(TransformationAddDeadContinueTest, Miscellaneous1) {
+  // A miscellaneous test that exposed a bug in spirv-fuzz.
+
+  std::string shader = R"(
+               OpCapability Shader
+          %1 = OpExtInstImport "GLSL.std.450"
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint Fragment %4 "main" %586 %623
+               OpExecutionMode %4 OriginUpperLeft
+               OpSource ESSL 310
+               OpMemberDecorate %34 0 Offset 0
+               OpDecorate %34 Block
+               OpDecorate %36 DescriptorSet 0
+               OpDecorate %36 Binding 0
+               OpDecorate %586 BuiltIn FragCoord
+               OpMemberDecorate %591 0 Offset 0
+               OpDecorate %591 Block
+               OpDecorate %593 DescriptorSet 0
+               OpDecorate %593 Binding 1
+               OpDecorate %623 Location 0
+          %2 = OpTypeVoid
+          %3 = OpTypeFunction %2
+          %6 = OpTypeInt 32 1
+          %7 = OpTypePointer Function %6
+          %9 = OpConstant %6 0
+         %16 = OpConstant %6 2
+         %17 = OpTypeBool
+         %27 = OpTypeFloat 32
+         %28 = OpTypeVector %27 2
+         %29 = OpTypeMatrix %28 2
+         %30 = OpTypePointer Private %29
+         %31 = OpVariable %30 Private
+         %34 = OpTypeStruct %27
+         %35 = OpTypePointer Uniform %34
+         %36 = OpVariable %35 Uniform
+         %37 = OpTypePointer Uniform %27
+         %40 = OpTypePointer Private %27
+         %43 = OpConstant %6 1
+         %62 = OpConstant %6 3
+         %64 = OpTypeVector %27 3
+         %65 = OpTypeMatrix %64 2
+         %66 = OpTypePointer Private %65
+         %67 = OpVariable %66 Private
+         %92 = OpConstant %6 4
+         %94 = OpTypeVector %27 4
+         %95 = OpTypeMatrix %94 2
+         %96 = OpTypePointer Private %95
+         %97 = OpVariable %96 Private
+        %123 = OpTypeMatrix %28 3
+        %124 = OpTypePointer Private %123
+        %125 = OpVariable %124 Private
+        %151 = OpTypeMatrix %64 3
+        %152 = OpTypePointer Private %151
+        %153 = OpVariable %152 Private
+        %179 = OpTypeMatrix %94 3
+        %180 = OpTypePointer Private %179
+        %181 = OpVariable %180 Private
+        %207 = OpTypeMatrix %28 4
+        %208 = OpTypePointer Private %207
+        %209 = OpVariable %208 Private
+        %235 = OpTypeMatrix %64 4
+        %236 = OpTypePointer Private %235
+        %237 = OpVariable %236 Private
+        %263 = OpTypeMatrix %94 4
+        %264 = OpTypePointer Private %263
+        %265 = OpVariable %264 Private
+        %275 = OpTypeInt 32 0
+        %276 = OpConstant %275 9
+        %277 = OpTypeArray %27 %276
+        %278 = OpTypePointer Function %277
+        %280 = OpConstant %27 0
+        %281 = OpTypePointer Function %27
+        %311 = OpConstant %27 16
+        %448 = OpConstant %6 5
+        %482 = OpConstant %6 6
+        %516 = OpConstant %6 7
+        %550 = OpConstant %6 8
+        %585 = OpTypePointer Input %94
+        %586 = OpVariable %585 Input
+        %587 = OpConstant %275 0
+        %588 = OpTypePointer Input %27
+        %591 = OpTypeStruct %28
+        %592 = OpTypePointer Uniform %591
+        %593 = OpVariable %592 Uniform
+        %596 = OpConstant %27 3
+        %601 = OpConstant %275 1
+        %617 = OpConstant %6 9
+        %622 = OpTypePointer Output %94
+        %623 = OpVariable %622 Output
+        %628 = OpConstant %27 1
+        %634 = OpConstantComposite %94 %280 %280 %280 %628
+        %635 = OpUndef %6
+        %636 = OpUndef %17
+        %637 = OpUndef %27
+        %638 = OpUndef %64
+        %639 = OpUndef %94
+        %640 = OpConstantTrue %17
+        %736 = OpConstantFalse %17
+        %642 = OpVariable %37 Uniform
+        %643 = OpVariable %40 Private
+          %4 = OpFunction %2 None %3
+          %5 = OpLabel
+               OpBranch %164
+        %164 = OpLabel
+               OpLoopMerge %166 %167 None
+               OpBranch %165
+        %165 = OpLabel
+               OpBranch %172
+        %172 = OpLabel
+               OpSelectionMerge %174 None
+               OpBranchConditional %640 %174 %174
+        %174 = OpLabel
+        %785 = OpCopyObject %6 %43
+               OpBranch %167
+        %167 = OpLabel
+        %190 = OpIAdd %6 %9 %785
+               OpBranchConditional %640 %164 %166
+        %166 = OpLabel
+               OpBranch %196
+        %196 = OpLabel
+               OpBranch %194
+        %194 = OpLabel
+               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()));
+
+  FactManager fact_manager;
+
+  // This transformation would shortcut the part of the loop body that defines
+  // an id used in the continue target.
+  auto bad_transformation = TransformationAddDeadContinue(165, false, {});
+  ASSERT_FALSE(bad_transformation.IsApplicable(context.get(), fact_manager));
+}
+
 }  // namespace
 }  // namespace fuzz
 }  // namespace spvtools