Merge pull request #62039 from eeckstein/fix-simplifycfg

SimplifyCFG: fix an argument use-after-erase in CheckedCastBrJumpThreading

Author: eeckstein

include/swift/SIL/SILArgument.h

@@ -101,6 +101,14 @@ class SILArgument : public ValueBase {
 
   SILBasicBlock *getParent() const { return parentBlock; }
 
+  /// Returns true if this argument is erased from a basic block.
+  ///
+  /// Note that SILArguments which are erased from a SILBasicBlock are not
+  /// destroyed and freed, but are kept in memory. So it's safe to keep a
+  /// pointer to an erased argument and then at a later time check if its
+  /// erased.
+  bool isErased() const { return !parentBlock; }
+
   SILFunction *getFunction();
   const SILFunction *getFunction() const;
 

include/swift/SIL/SILBasicBlock.h

@@ -308,7 +308,7 @@ public SwiftObjectHeader {
                                     const ValueDecl *D = nullptr);
 
   /// Remove all block arguments.
-  void dropAllArguments() { ArgumentList.clear(); }
+  void dropAllArguments();
 
   //===--------------------------------------------------------------------===//
   // Successors

lib/SIL/IR/SILBasicBlock.cpp

@@ -180,6 +180,7 @@ SILFunctionArgument *SILBasicBlock::replaceFunctionArgument(
 
   SILFunctionArgument *NewArg = new (M) SILFunctionArgument(Ty, Kind, D);
   NewArg->setParent(this);
+  ArgumentList[i]->parentBlock = nullptr;
 
   // TODO: When we switch to malloc/free allocation we'll be leaking memory
   // here.
@@ -203,6 +204,7 @@ SILPhiArgument *SILBasicBlock::replacePhiArgument(unsigned i, SILType Ty,
 
   SILPhiArgument *NewArg = new (M) SILPhiArgument(Ty, Kind, D);
   NewArg->setParent(this);
+  ArgumentList[i]->parentBlock = nullptr;
 
   // TODO: When we switch to malloc/free allocation we'll be leaking memory
   // here.
@@ -257,9 +259,17 @@ SILPhiArgument *SILBasicBlock::insertPhiArgument(unsigned AtArgPos, SILType Ty,
   return arg;
 }
 
+void SILBasicBlock::dropAllArguments() {
+  for (SILArgument *arg : ArgumentList) {
+    arg->parentBlock = nullptr;
+  }
+  ArgumentList.clear();
+}
+
 void SILBasicBlock::eraseArgument(int Index) {
   assert(getArgument(Index)->use_empty() &&
          "Erasing block argument that has uses!");
+  ArgumentList[Index]->parentBlock = nullptr;
   ArgumentList.erase(ArgumentList.begin() + Index);
 }
 

lib/SILOptimizer/Utils/CheckedCastBrJumpThreading.cpp

@@ -771,6 +771,9 @@ void CheckedCastBrJumpThreading::optimizeFunction() {
     Fn->verifyCriticalEdges();
 
   for (Edit *edit : Edits) {
+    if (edit->SuccessArg->isErased())
+      continue;
+
     BasicBlockCloner Cloner(edit->CCBBlock, deBlocks);
     if (!Cloner.canCloneBlock())
       continue;

test/SILOptimizer/simplify_cfg.sil

@@ -3233,3 +3233,89 @@ bb6:
   %r = tuple()
   return %r : $()
 }
+
+// CHECK-LABEL: sil @dont_crash_in_cast_jump_threading
+// CHECK:         checked_cast_br
+// CHECK:       } // end sil function 'dont_crash_in_cast_jump_threading'
+sil @dont_crash_in_cast_jump_threading : $@convention(thin) (Optional<Base>) -> () {
+bb0(%0 : $Optional<Base>):
+  br bb1
+
+bb1:
+  switch_enum %0 : $Optional<Base>, case #Optional.some!enumelt: bb2, case #Optional.none!enumelt: bb21
+
+
+bb2(%18 : $Base):
+  checked_cast_br %18 : $Base to Derived, bb3, bb4
+
+bb3(%21 : $Derived):
+  %22 = integer_literal $Builtin.Int1, -1
+  br bb5(%22 : $Builtin.Int1)
+
+bb4:
+  %24 = integer_literal $Builtin.Int1, 0
+  br bb5(%24 : $Builtin.Int1)
+
+
+bb5(%26 : $Builtin.Int1):
+  cond_br %26, bb6, bb19
+
+bb6:
+  checked_cast_br %18 : $Base to Derived, bb7, bb8
+
+
+bb7(%29 : $Derived):
+  %30 = enum $Optional<Derived>, #Optional.some!enumelt, %29 : $Derived
+  br bb9(%30 : $Optional<Derived>)
+
+bb8:
+  %32 = enum $Optional<Derived>, #Optional.none!enumelt
+  br bb9(%32 : $Optional<Derived>)
+
+
+bb9(%34 : $Optional<Derived>):
+  switch_enum %34 : $Optional<Derived>, case #Optional.some!enumelt: bb10, case #Optional.none!enumelt: bb18
+
+
+bb10(%36 : $Derived):
+  br bb11
+
+bb11:
+  checked_cast_br %18 : $Base to Derived, bb12, bb13
+
+
+bb12(%43 : $Derived):
+  %44 = enum $Optional<Derived>, #Optional.some!enumelt, %43 : $Derived
+  br bb14(%44 : $Optional<Derived>)
+
+bb13:
+  %46 = enum $Optional<Derived>, #Optional.none!enumelt
+  br bb14(%46 : $Optional<Derived>)
+
+
+bb14(%48 : $Optional<Derived>):
+  switch_enum %48 : $Optional<Derived>, case #Optional.some!enumelt: bb15, case #Optional.none!enumelt: bb17
+
+
+bb15(%50 : $Derived):
+  br bb16
+
+bb16:
+  br bb20
+
+bb17:
+  br bb16
+
+bb18:
+  br bb11
+
+bb19:
+  br bb20
+
+bb20:
+  br bb1
+
+bb21:
+  %75 = tuple ()
+  return %75 : $()
+}