Fix SILCombine of existential applies

include/swift/SILOptimizer/Utils/Existential.h

@@ -74,6 +74,9 @@ struct ConcreteExistentialInfo {
   SILValue ConcreteValue;
   // True if the ConcreteValue is copied from another stack location
   bool isConcreteValueCopied = false;
+  // True if the ConcreteValue should replace all uses of the opened
+  // existential.
+  bool ConcreteValueNeedsFixup = false;
   // When ConcreteType is itself an opened existential, record the type
   // definition. May be nullptr for a valid AppliedConcreteType.
   SingleValueInstruction *ConcreteTypeDef = nullptr;

lib/SILOptimizer/SILCombiner/SILCombinerApplyVisitors.cpp

@@ -798,30 +798,13 @@ SILCombiner::buildConcreteOpenedExistentialInfoFromSoleConformingType(
   // Prepare the code by adding UncheckedCast instructions that cast opened
   // existentials to concrete types. Set the ConcreteValue of CEI.
   if (auto *OER = dyn_cast<OpenExistentialRefInst>(OAI.OpenedArchetypeValue)) {
-    // If we have an owned open_existential_ref, we only optimize for now if our
-    // open_existential_ref has a single non-debug consuming use that is a
-    // destroy_value.
-    if (OER->getForwardingOwnershipKind() != OwnershipKind::Owned) {
-      // We use OER as the insertion point so that
-      SILBuilderWithScope b(std::next(OER->getIterator()), Builder);
-      auto loc = RegularLocation::getAutoGeneratedLocation();
-      SoleCEI.ConcreteValue =
-          b.createUncheckedRefCast(loc, OER, concreteSILType);
-      return COAI;
-    }
-
-    auto *consumingUse = OER->getSingleConsumingUse();
-    if (!consumingUse || !isa<DestroyValueInst>(consumingUse->getUser())) {
-      return std::nullopt;
-    }
-
-    // We use std::next(OER) as the insertion point so that we can reuse the
-    // destroy_value of consumingUse.
     SILBuilderWithScope b(std::next(OER->getIterator()), Builder);
     auto loc = RegularLocation::getAutoGeneratedLocation();
-    auto *uri = b.createUncheckedRefCast(loc, OER, concreteSILType);
-    SoleCEI.ConcreteValue = uri;
-    replaceInstUsesWith(*OER, uri);
+    SoleCEI.ConcreteValue = b.createUncheckedRefCast(loc, OER, concreteSILType);
+    if (F->hasOwnership() &&
+        SoleCEI.ConcreteValue->getOwnershipKind() == OwnershipKind::Owned) {
+      SoleCEI.ConcreteValueNeedsFixup = true;
+    }
     return COAI;
   }
 
@@ -852,8 +835,19 @@ SILCombiner::buildConcreteOpenedExistentialInfo(Operand &ArgOperand) {
   // Build a ConcreteOpenedExistentialInfo following the data flow chain of the
   // ArgOperand through the open_existential backward to an init_existential.
   ConcreteOpenedExistentialInfo COEI(ArgOperand);
-  if (COEI.CEI)
+  if (COEI.CEI) {
+    auto ConcreteValue = COEI.CEI->ConcreteValue;
+    if (ConcreteValue->getFunction()->hasOwnership() &&
+        ConcreteValue->getOwnershipKind() == OwnershipKind::Owned) {
+      SILBuilderWithScope b(COEI.OAI.OpenedArchetypeValue->getNextInstruction(),
+                            Builder);
+      auto loc = RegularLocation::getAutoGeneratedLocation();
+      COEI.CEI->ConcreteValue = b.createUncheckedRefCast(
+          loc, COEI.OAI.OpenedArchetypeValue, ConcreteValue->getType());
+      COEI.CEI->ConcreteValueNeedsFixup = true;
+    }
     return COEI;
+  }
 
   // Use SoleConformingType information.
   return buildConcreteOpenedExistentialInfoFromSoleConformingType(ArgOperand);
@@ -1147,6 +1141,7 @@ SILInstruction *SILCombiner::createApplyWithConcreteType(
 
   // Transform the parameter arguments.
   SmallVector<ConcreteArgumentCopy, 4> concreteArgCopies;
+  llvm::SmallMapVector<SILValue, SILValue, 4> valuesToReplace;
   for (unsigned EndIdx = Apply.getNumArguments(); ArgIdx < EndIdx; ++ArgIdx) {
     auto ArgIt = COAIs.find(ArgIdx);
     if (ArgIt == COAIs.end()) {
@@ -1199,6 +1194,9 @@ SILInstruction *SILCombiner::createApplyWithConcreteType(
       concreteArgCopies.push_back(*argSub);
       NewArgs.push_back(argSub->tempArg);
     } else {
+      if (CEI.ConcreteValueNeedsFixup) {
+        valuesToReplace[OAI.OpenedArchetypeValue] = CEI.ConcreteValue;
+      }
       NewArgs.push_back(CEI.ConcreteValue);
     }
 
@@ -1260,9 +1258,8 @@ SILInstruction *SILCombiner::createApplyWithConcreteType(
     // SILCombine Worklist.
     InstructionDeleter deleter;
     for (SILInstruction *inst : *Builder.getTrackingList()) {
-      deleter.trackIfDead(inst);
+      deleter.deleteIfDead(inst, false);
     }
-    deleter.cleanupDeadInstructions();
     Builder.getTrackingList()->clear();
     return nullptr;
   }
@@ -1304,6 +1301,20 @@ SILInstruction *SILCombiner::createApplyWithConcreteType(
       cleanupBuilder.createDeallocStack(cleanupLoc, argCopy.tempArg);
     }
   }
+  // Replace all uses of the opened existential with the unconditional cast to
+  // concrete type.
+  for (auto &valuePair : valuesToReplace) {
+    auto openedExistential = valuePair.first;
+    auto uncheckedCast = valuePair.second;
+    SmallVector<Operand *> uses(openedExistential->getNonTypeDependentUses());
+    for (auto use : uses) {
+      auto *user = use->getUser();
+      if (user == cast<SingleValueInstruction>(uncheckedCast)) {
+        continue;
+      }
+      use->set(uncheckedCast);
+    }
+  }
   return NewApply.getInstruction();
 }
 

test/SILOptimizer/sil_combine_ossa.sil

@@ -2855,7 +2855,21 @@ bb0(%0 : @owned $Klass):
   %2 = open_existential_ref %1 : $AnyObject to $@opened("7CAE06CE-5F10-11E4-AF13-C82A1428F987", AnyObject) Self
   %f = function_ref @use_generic_obj_guaranteed : $@convention(thin) <τ_0_0> (@guaranteed τ_0_0) -> ()
   apply %f<@opened("7CAE06CE-5F10-11E4-AF13-C82A1428F987", AnyObject) Self>(%2) : $@convention(thin) <τ_0_0> (@guaranteed τ_0_0) -> ()
+  %3 = unchecked_ref_cast %2 : $@opened("7CAE06CE-5F10-11E4-AF13-C82A1428F987", AnyObject) Self to $Builtin.NativeObject
+  return %3 : $Builtin.NativeObject
+}
 
+// CHECK-LABEL: sil [ossa] @collapse_existential_pack_unpack_unchecked_ref_cast_owned4 :
+// CHECK:     apply {{.*}}<Klass>
+// CHECK: } // end sil function 'collapse_existential_pack_unpack_unchecked_ref_cast_owned4'
+sil [ossa] @collapse_existential_pack_unpack_unchecked_ref_cast_owned4 : $@convention(thin) (@owned Klass) -> @owned Builtin.NativeObject {
+bb0(%0 : @owned $Klass):
+  %1 = init_existential_ref %0 : $Klass : $Klass, $AnyObject
+  %f1 = function_ref @use_anyobject_guaranteed : $@convention(thin) (@guaranteed AnyObject) -> ()
+  apply %f1(%1) : $@convention(thin) (@guaranteed AnyObject) -> ()
+  %2 = open_existential_ref %1 : $AnyObject to $@opened("7CAE06CE-5F10-11E4-AF13-C82A1428F987", AnyObject) Self
+  %f2 = function_ref @use_generic_obj_guaranteed : $@convention(thin) <τ_0_0> (@guaranteed τ_0_0) -> ()
+  apply %f2<@opened("7CAE06CE-5F10-11E4-AF13-C82A1428F987", AnyObject) Self>(%2) : $@convention(thin) <τ_0_0> (@guaranteed τ_0_0) -> ()
   %3 = unchecked_ref_cast %2 : $@opened("7CAE06CE-5F10-11E4-AF13-C82A1428F987", AnyObject) Self to $Builtin.NativeObject
   return %3 : $Builtin.NativeObject
 }