[sil-cse] Types of basic block arguments should be updated during the CSE of open_existential_ref (#3559)

Opened archetypes may be used by types of basic block arguments. Therefore, when performing the CSE of open_existential_ref instructions, one should also update the type of any basic block argument that may have used the archetype opened by the open_existential_ref being replaced.

rdar://problem/27386065

Author: swiftix

lib/SILOptimizer/Transforms/CSE.cpp

@@ -591,6 +591,42 @@ namespace {
   };
 }
 
+/// Update SIL basic block's arguments types which refer to opened
+/// archetypes. Replace such types by performing type substitutions
+/// according to the provided type substitution map.
+static void updateBasicBlockArgTypes(SILBasicBlock *BB,
+                                     TypeSubstitutionMap &TypeSubstMap) {
+  // Check types of all BB arguments.
+  for (auto &Arg : BB->getBBArgs()) {
+    if (!Arg->getType().getSwiftRValueType()->hasOpenedExistential())
+      continue;
+    // Type of this BB argument uses an opened existential.
+    // Try to apply substitutions to it and if it produces a different type,
+    // use this type as new type of the BB argument.
+    auto NewArgType = Arg->getType().subst(
+        BB->getModule(), BB->getModule().getSwiftModule(), TypeSubstMap);
+    if (NewArgType == Arg->getType())
+      continue;
+    // Replace the type of this BB argument. The type of a BBArg
+    // can only be changed using replaceBBArg, if the BBArg has no uses.
+    // So, make it look as if it has no uses.
+
+    // First collect all uses, before changing the type.
+    SmallVector<Operand *, 4> OriginalArgUses;
+    for (auto *ArgUse : Arg->getUses()) {
+      OriginalArgUses.push_back(ArgUse);
+    }
+    // Then replace all uses by an undef.
+    Arg->replaceAllUsesWith(SILUndef::get(Arg->getType(), BB->getModule()));
+    // Replace the type of the BB argument.
+    BB->replaceBBArg(Arg->getIndex(), NewArgType, Arg->getDecl());
+    // Restore all uses to refer to the BB argument with updated type.
+    for (auto ArgUse : OriginalArgUses) {
+      ArgUse->set(Arg);
+    }
+  }
+}
+
 /// Handle CSE of open_existential_ref instructions.
 /// Returns true if uses of open_existential_ref can
 /// be replaced by a dominating instruction.
@@ -601,6 +637,10 @@ bool CSE::processOpenExistentialRef(SILInstruction *Inst, ValueBase *V,
                                     SILBasicBlock::iterator &I) {
   assert(isa<OpenExistentialRefInst>(Inst));
   llvm::SmallSetVector<SILInstruction *, 16> Candidates;
+  auto OldOpenedArchetype = getOpenedArchetypeOf(Inst);
+  auto NewOpenedArchetype = getOpenedArchetypeOf(dyn_cast<SILInstruction>(V));
+  TypeSubstitutionMap TypeSubstMap;
+  TypeSubstMap[OldOpenedArchetype.getPointer()] = NewOpenedArchetype;
   // Collect all candidates that may contain opened archetypes
   // that need to be replaced.
   for (auto Use : Inst->getUses()) {
@@ -614,10 +654,21 @@ bool CSE::processOpenExistentialRef(SILInstruction *Inst, ValueBase *V,
       }
       Candidates.insert(User);
     }
+    if (!isa<TermInst>(User))
+      continue;
+    // The current use of the opened archetype is a terminator instruction.
+    // Check if any of the successor BBs uses this opened archetype in the
+    // types of its basic block arguments. If this is the case, replace
+    // those uses by the new opened archetype.
+    auto Successors = User->getParent()->getSuccessorBlocks();
+    for (auto Successor : Successors) {
+      if (Successor->bbarg_empty())
+        continue;
+      // If a BB has any arguments, update their types if necessary.
+      updateBasicBlockArgTypes(Successor, TypeSubstMap);
+    }
   }
   // Now process candidates.
-  auto OldOpenedArchetype = getOpenedArchetypeOf(Inst);
-  auto NewOpenedArchetype = getOpenedArchetypeOf(dyn_cast<SILInstruction>(V));
   // TODO: Move it to CSE instance to avoid recreating it every time?
   SILOpenedArchetypesTracker OpenedArchetypesTracker(*Inst->getFunction());
   // Register the new archetype to be used.

test/SILOptimizer/cse_objc.sil

@@ -119,6 +119,41 @@ bb0(%0 : $Walkable):
   return %9 : $()                                 // id: %10
 }
 
+// Check that the open_existential_ref is CSEed not only in instructions, but also in the types
+// of BB arguments. The type of BB2 argument should use the same opened existential as the
+// one opened in the entry block.
+
+// CHECK-LABEL: sil @cse_open_existential_in_bbarg : $@convention(thin) (@owned AnyObject) -> ()
+// CHECK: open_existential_ref %0 : $AnyObject to $[[OPENED:@opened\(.*\) AnyObject]]
+// CHECK: bb1(%{{[0-9]*}} : $@convention(objc_method) ([[OPENED]]) -> @autoreleased Optional<NSString>):
+// CHECK-NOT: open_existential_ref
+// CHECK: bb2(%{{[0-9]*}} : $@convention(objc_method) ([[OPENED]]) -> @autoreleased Optional<NSString>):
+// CHECK-NOT: open_existential_ref
+// CHECK: return
+sil @cse_open_existential_in_bbarg : $@convention(thin) (@owned AnyObject) -> () {
+bb0(%0 : $AnyObject):
+  %1 = open_existential_ref %0 : $AnyObject to $@opened("CCEAC0E2-4BB2-11E6-86F8-B8E856428C60") AnyObject
+  dynamic_method_br %1 : $@opened("CCEAC0E2-4BB2-11E6-86F8-B8E856428C60") AnyObject, #NSProxy.description!getter.1.foreign, bb1, bb3
+
+bb1(%3 : $@convention(objc_method) (@opened("CCEAC0E2-4BB2-11E6-86F8-B8E856428C60") AnyObject) -> @autoreleased Optional<NSString>):
+  strong_retain %1 : $@opened("CCEAC0E2-4BB2-11E6-86F8-B8E856428C60") AnyObject
+  %5 = partial_apply %3(%1) : $@convention(objc_method) (@opened("CCEAC0E2-4BB2-11E6-86F8-B8E856428C60") AnyObject) -> @autoreleased Optional<NSString>
+  %6 = apply %5() : $@callee_owned () -> @owned Optional<NSString>
+  %7 = open_existential_ref %0 : $AnyObject to $@opened("CCEDAF1E-4BB2-11E6-86F8-B8E856428C60") AnyObject
+  dynamic_method_br %7 : $@opened("CCEDAF1E-4BB2-11E6-86F8-B8E856428C60") AnyObject, #NSProxy.description!getter.1.foreign, bb2, bb3
+
+bb2(%9 : $@convention(objc_method) (@opened("CCEDAF1E-4BB2-11E6-86F8-B8E856428C60") AnyObject) -> @autoreleased Optional<NSString>):
+  strong_retain %7 : $@opened("CCEDAF1E-4BB2-11E6-86F8-B8E856428C60") AnyObject
+  %10 = partial_apply %9(%7) : $@convention(objc_method) (@opened("CCEDAF1E-4BB2-11E6-86F8-B8E856428C60") AnyObject) -> @autoreleased Optional<NSString>
+  %11 = apply %10() : $@callee_owned () -> @owned Optional<NSString>
+  br bb3
+
+bb3:
+  strong_release %0 : $AnyObject
+  %13 = tuple ()
+  return %13 : $()
+}
+
 sil_vtable Bar {
   #Bar.init!initializer.1: _TFC4test3BarcfMS0_FT_S0_	// test.Bar.init (test.Bar.Type)() -> test.Bar
   #Bar.walk!1: _TFC4test3Bar4walkfS0_FT_T_	// test.Bar.walk (test.Bar)() -> ()