Merge pull request #74329 from eeckstein/fix-block-cloning

SIL: update borrowed-from instructions when cloning a basic block

include/swift/SIL/SILCloner.h

@@ -87,6 +87,10 @@ protected:
   /// blocks.
   llvm::DenseMap<SILBasicBlock*, SILBasicBlock*> BBMap;
 
+  /// Blocks, where edge-spitting may have "converted" terminator result
+  /// arguments to phi-arguments.
+  llvm::SmallVector<SILBasicBlock *> blocksWithNewPhiArgs;
+
 private:
   /// MARK: Private state hidden from CRTP extensions.
 
@@ -872,9 +876,15 @@ void SILCloner<ImplClass>::visitBlocksDepthFirst(SILBasicBlock *startBB) {
             && isa<BranchInst>(BB->getTerminator())) {
           continue;
         }
+        TermInst *term = BB->getTerminator();
+
+        // After edge-spitting, terminator result arguments become phi arguments.
+        if (!isa<BranchInst>(term))
+          blocksWithNewPhiArgs.push_back(BB->getSuccessors()[succIdx]);
+
         // This predecessor has multiple successors, so cloning it without
         // cloning its successors would create a critical edge.
-        splitEdge(BB->getTerminator(), succIdx, DomTree);
+        splitEdge(term, succIdx, DomTree);
         assert(!BBMap.count(BB->getSuccessors()[succIdx]));
       }
       // Map the successor to a new BB. Layout the cloned blocks in the order

lib/SILOptimizer/Utils/BasicBlockOptUtils.cpp

@@ -125,41 +125,47 @@ void BasicBlockCloner::updateSSAAfterCloning() {
       break;
     }
   }
-  if (!needsSSAUpdate)
-    return;
-
-  SILSSAUpdater ssaUpdater(&updateSSAPhis);
-  for (auto availValPair : availVals) {
-    auto inst = availValPair.first;
-    if (inst->use_empty())
-      continue;
-
-    SILValue newResult(availValPair.second);
-
-    SmallVector<UseWrapper, 16> useList;
-    // Collect the uses of the value.
-    for (auto *use : inst->getUses())
-      useList.push_back(UseWrapper(use));
-
-    ssaUpdater.initialize(inst->getFunction(), inst->getType(),
-                          inst->getOwnershipKind());
-    ssaUpdater.addAvailableValue(origBB, inst);
-    ssaUpdater.addAvailableValue(getNewBB(), newResult);
-
-    if (useList.empty())
-      continue;
-
-    // Update all the uses.
-    for (auto useWrapper : useList) {
-      Operand *use = useWrapper; // unwrap
-      SILInstruction *user = use->getUser();
-      assert(user && "Missing user");
-
-      // Ignore uses in the same basic block.
-      if (user->getParent() == origBB)
+  if (needsSSAUpdate) {
+    SILSSAUpdater ssaUpdater(&updateSSAPhis);
+    for (auto availValPair : availVals) {
+      auto inst = availValPair.first;
+      if (inst->use_empty())
         continue;
 
-      ssaUpdater.rewriteUse(*use);
+      SILValue newResult(availValPair.second);
+
+      SmallVector<UseWrapper, 16> useList;
+      // Collect the uses of the value.
+      for (auto *use : inst->getUses())
+        useList.push_back(UseWrapper(use));
+
+      ssaUpdater.initialize(inst->getFunction(), inst->getType(),
+                            inst->getOwnershipKind());
+      ssaUpdater.addAvailableValue(origBB, inst);
+      ssaUpdater.addAvailableValue(getNewBB(), newResult);
+
+      if (useList.empty())
+        continue;
+
+      // Update all the uses.
+      for (auto useWrapper : useList) {
+        Operand *use = useWrapper; // unwrap
+        SILInstruction *user = use->getUser();
+        assert(user && "Missing user");
+
+        // Ignore uses in the same basic block.
+        if (user->getParent() == origBB)
+          continue;
+
+        ssaUpdater.rewriteUse(*use);
+      }
+    }
+  }
+  for (SILBasicBlock *b : blocksWithNewPhiArgs) {
+    for (SILArgument *arg : b->getArguments()) {
+      if (arg->getOwnershipKind() == OwnershipKind::Guaranteed) {
+        updateSSAPhis.push_back(cast<SILPhiArgument>(arg));
+      }
     }
   }
   updateBorrowedFromPhis(pm, updateSSAPhis);

test/SILOptimizer/simplify_cfg_ossa.sil

@@ -25,6 +25,7 @@ class Klass {
   var a: Int
    deinit
   init()
+  func foo() -> Int?
 }
 
 class B {}
@@ -38,6 +39,10 @@ struct KlassWrapper {
   var k: Klass
 }
 
+struct OptionalKlassWrapper {
+  @_hasStorage var k: Optional<Klass>
+}
+
 internal enum CompareResult {
   case equal
   case less
@@ -1843,3 +1848,42 @@ bb10(%7 : $()):
 bb11:
   unreachable
 }
+
+// CHECK-LABEL: sil [ossa] @jump_threading_creates_phi :
+// CHECK:         = borrowed %{{[0-9]+}} : $Klass from (%1 : $OptionalKlassWrapper)
+// CHECK:       } // end sil function 'jump_threading_creates_phi'
+sil [ossa] @jump_threading_creates_phi: $@convention(thin) (Optional<A>, @guaranteed OptionalKlassWrapper) -> () {
+bb0(%0 : $Optional<A>, %1 : @guaranteed $OptionalKlassWrapper):
+  %2 = alloc_stack [lexical] $Optional<A>
+  store %0 to [trivial] %2 : $*Optional<A>
+  switch_enum_addr %2 : $*Optional<A>, case #Optional.some!enumelt: bb1, case #Optional.none!enumelt: bb2
+
+bb1:
+  %5 = unchecked_take_enum_data_addr %2 : $*Optional<A>, #Optional.some!enumelt
+  %6 = load [trivial] %5 : $*A
+  br bb3(%6 : $A)
+
+bb2:
+  %8 = enum $A, #A.B!enumelt
+  br bb3(%8 : $A)
+
+bb3(%10 : $A):
+  dealloc_stack %2 : $*Optional<A>
+  %12 = enum $Optional<A>, #Optional.some!enumelt, %10 : $A
+  %13 = struct_extract %1 : $OptionalKlassWrapper, #OptionalKlassWrapper.k
+  switch_enum %13 : $Optional<Klass>, case #Optional.some!enumelt: bb4, case #Optional.none!enumelt: bb5
+
+bb4(%15 : @guaranteed $Klass):
+  %16 = class_method %15 : $Klass, #Klass.foo : (Klass) -> () -> Int?, $@convention(method) (@guaranteed Klass) -> Optional<Int>
+  %17 = apply %16(%15) : $@convention(method) (@guaranteed Klass) -> Optional<Int>
+  br bb6(%17 : $Optional<Int>)
+
+bb5:
+  %19 = enum $Optional<Int>, #Optional.none!enumelt
+  br bb6(%19 : $Optional<Int>)
+
+bb6(%21 : $Optional<Int>):
+  %22 = tuple ()
+  return %22 : $()
+}
+