Merge pull request #64204 from eeckstein/fix-dce

SILOptimizer: fix a stack overflow in DCE

Author: eeckstein

include/swift/SIL/SILBasicBlock.h

@@ -427,11 +427,6 @@ public SwiftObjectHeader {
     return getTerminator()->getSingleSuccessorBlock();
   }
 
-  /// Returns true if \p BB is a successor of this block.
-  bool isSuccessorBlock(SILBasicBlock *Block) const {
-    return getTerminator()->isSuccessorBlock(Block);
-  }
-
   using SuccessorBlockListTy = TermInst::SuccessorBlockListTy;
   using ConstSuccessorBlockListTy = TermInst::ConstSuccessorBlockListTy;
 
@@ -459,12 +454,6 @@ public SwiftObjectHeader {
     return {pred_begin(), pred_end()};
   }
 
-  bool isPredecessorBlock(SILBasicBlock *BB) const {
-    return any_of(
-        getPredecessorBlocks(),
-        [&BB](const SILBasicBlock *PredBB) -> bool { return BB == PredBB; });
-  }
-
   SILBasicBlock *getSinglePredecessorBlock() {
     if (pred_empty() || std::next(pred_begin()) != pred_end())
       return nullptr;

include/swift/SIL/SILInstruction.h

@@ -8940,14 +8940,6 @@ class TermInst : public NonValueInstruction {
     return const_cast<TermInst *>(this)->getSingleSuccessorBlock();
   }
 
-  /// Returns true if \p BB is a successor of this block.
-  bool isSuccessorBlock(SILBasicBlock *BB) const {
-    auto Range = getSuccessorBlocks();
-    return any_of(Range, [&BB](const SILBasicBlock *SuccBB) -> bool {
-      return BB == SuccBB;
-    });
-  }
-
   using SuccessorBlockArgumentListTy =
       TransformRange<ConstSuccessorListTy, function_ref<ArrayRef<SILArgument *>(
                                                const SILSuccessor &)>>;

lib/SIL/Verifier/SILVerifier.cpp

@@ -6289,18 +6289,6 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
   }
 
   void visitSILBasicBlock(SILBasicBlock *BB) {
-    // Make sure that each of the successors/predecessors of this basic block
-    // have this basic block in its predecessor/successor list.
-    for (const auto *SuccBB : BB->getSuccessorBlocks()) {
-      require(SuccBB->isPredecessorBlock(BB),
-              "Must be a predecessor of each successor.");
-    }
-
-    for (const SILBasicBlock *PredBB : BB->getPredecessorBlocks()) {
-      require(PredBB->isSuccessorBlock(BB),
-              "Must be a successor of each predecessor.");
-    }
-    
     SILInstructionVisitor::visitSILBasicBlock(BB);
     verifyDebugScopeHoles(BB);
   }
@@ -6331,6 +6319,33 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
         continue;
       visitSILBasicBlock(&BB);
     }
+
+    verifyPredecessorSucessorStructure(F);
+  }
+
+  // Make sure that each of the successors/predecessors of a basic block
+  // have this basic block in its predecessor/successor list.
+  void verifyPredecessorSucessorStructure(SILFunction *f) {
+    using PredSuccPair = std::pair<SILBasicBlock *, SILBasicBlock *>;
+    llvm::DenseSet<PredSuccPair> foundSuccessors;
+    llvm::DenseSet<PredSuccPair> foundPredecessors;
+
+    for (auto &block : *f) {
+      for (SILBasicBlock *succ : block.getSuccessorBlocks()) {
+        foundSuccessors.insert({&block, succ});
+      }
+      for (SILBasicBlock *pred : block.getPredecessorBlocks()) {
+        foundPredecessors.insert({pred, &block});
+      }
+    }
+    for (PredSuccPair predSucc : foundSuccessors) {
+      require(foundPredecessors.contains(predSucc),
+              "block is not predecessor of its successor");
+    }
+    for (PredSuccPair predSucc : foundPredecessors) {
+      require(foundSuccessors.contains(predSucc),
+              "block is not successor of its predecessor");
+    }
   }
 
   void visitSILFunction(SILFunction *F) {

lib/SILOptimizer/Analysis/LoopRegionAnalysis.cpp

@@ -218,10 +218,12 @@ void LoopRegionFunctionInfo::verify() {
 
       // If R and OtherR are blocks, then OtherR should be a successor of the
       // real block.
-      if (R->isBlock() && OtherR->isBlock())
-        assert(R->getBlock()->isSuccessorBlock(OtherR->getBlock()) &&
+      if (R->isBlock() && OtherR->isBlock()) {
+        auto succs = R->getBlock()->getSuccessors();
+        assert(std::find(succs.begin(), succs.end(), OtherR->getBlock()) != succs.end() &&
                "Expected either R was not a block or OtherR was a CFG level "
                "successor of R.");
+      }
     }
   }
 #endif

lib/SILOptimizer/Transforms/DeadCodeElimination.cpp

@@ -215,7 +215,6 @@ void DCE::markInstructionLive(SILInstruction *Inst) {
 
   LLVM_DEBUG(llvm::dbgs() << "Marking as live: " << *Inst);
 
-  markControllingTerminatorsLive(Inst->getParent());
   Worklist.push_back(Inst);
 }
 
@@ -450,6 +449,8 @@ void DCE::propagateLiveBlockArgument(SILArgument *Arg) {
 // Given an instruction which is considered live, propagate that liveness
 // back to the instructions that produce values it consumes.
 void DCE::propagateLiveness(SILInstruction *I) {
+  markControllingTerminatorsLive(I->getParent());
+
   if (!isa<TermInst>(I)) {
     for (auto &O : I->getAllOperands())
       markValueLive(O.get());

validation-test/SILOptimizer/large_switch.gyb

@@ -0,0 +1,19 @@
+// RUN: %empty-directory(%t)
+// RUN: %gyb %s > %t/main.swift
+// RUN: %target-swift-frontend -O -parse-as-library -sil-verify-none -emit-sil %t/main.swift -o /dev/null
+
+// REQUIRES: swift_stdlib_no_asserts,optimized_stdlib
+// REQUIRES: long_test
+
+// Check that the compiler does not crash in DCE due to a stack overflow.
+// rdar://106198943
+
+public func test(_ a: Int) -> Int {
+  switch (a) {
+% for i in range(32000):
+    case ${i}: return ${i % 7}
+% end
+    default: return -1
+  }
+}
+