[sil-cloner] Strip out the argument of default blocks of switch_enum/checked_cast_br when cloning from ossa into non-ossa code.

Author: gottesmm

include/swift/SIL/SILCloner.h

@@ -715,12 +715,43 @@ void SILCloner<ImplClass>::visitBlocksDepthFirst(SILBasicBlock *startBB) {
 template<typename ImplClass>
 void
 SILCloner<ImplClass>::doFixUp(SILFunction *F) {
+  // If our source function is in ossa form, but the function into which we are
+  // cloning is not in ossa, after we clone, eliminate default arguments.
+  if (!getBuilder().hasOwnership() && F->hasOwnership()) {
+    for (auto &Block : getBuilder().getFunction()) {
+      auto *Term = Block.getTerminator();
+      if (auto *CCBI = dyn_cast<CheckedCastBranchInst>(Term)) {
+        // Check if we have a default argument.
+        auto *FailureBlock = CCBI->getFailureBB();
+        assert(FailureBlock->getNumArguments() <= 1 &&
+               "We should either have no args or a single default arg");
+        if (0 == FailureBlock->getNumArguments())
+          continue;
+        FailureBlock->getArgument(0)->replaceAllUsesWith(CCBI->getOperand());
+        FailureBlock->eraseArgument(0);
+        continue;
+      }
+
+      if (auto *SEI = dyn_cast<SwitchEnumInst>(Term)) {
+        if (auto DefaultBlock = SEI->getDefaultBBOrNull()) {
+          assert(DefaultBlock.get()->getNumArguments() <= 1 &&
+                 "We should either have no args or a single default arg");
+          if (0 == DefaultBlock.get()->getNumArguments())
+            continue;
+          DefaultBlock.get()->getArgument(0)->replaceAllUsesWith(
+              SEI->getOperand());
+          DefaultBlock.get()->eraseArgument(0);
+          continue;
+        }
+      }
+    }
+  }
+
   // Remove any code after unreachable instructions.
 
-  // NOTE: It is unfortunate that it essentially duplicates
-  // the code from sil-combine, but doing so allows for
-  // avoiding any cross-layer invocations between SIL and
-  // SILOptimizer layers.
+  // NOTE: It is unfortunate that it essentially duplicates the code from
+  // sil-combine, but doing so allows for avoiding any cross-layer invocations
+  // between SIL and SILOptimizer layers.
 
   for (auto *BB : BlocksWithUnreachables) {
     for (auto &I : *BB) {

include/swift/SIL/SILInstruction.h

@@ -7207,10 +7207,14 @@ class SwitchEnumInstBase : public TermInst {
   bool hasDefault() const {
     return SILInstruction::Bits.SwitchEnumInstBase.HasDefault;
   }
+
   SILBasicBlock *getDefaultBB() const {
     assert(hasDefault() && "doesn't have a default");
     return getSuccessorBuf()[getNumCases()];
   }
+
+  NullablePtr<SILBasicBlock> getDefaultBBOrNull() const;
+
   ProfileCounter getDefaultCount() const {
     assert(hasDefault() && "doesn't have a default");
     return getSuccessorBuf()[getNumCases()].getCount();

lib/SIL/SILInstructions.cpp

@@ -1566,6 +1566,12 @@ SwitchEnumInstBase::getUniqueCaseForDestination(SILBasicBlock *BB) {
   return D;
 }
 
+NullablePtr<SILBasicBlock> SwitchEnumInstBase::getDefaultBBOrNull() const {
+  if (!hasDefault())
+    return nullptr;
+  return getDefaultBB();
+}
+
 SwitchEnumInst *SwitchEnumInst::create(
     SILDebugLocation Loc, SILValue Operand, SILBasicBlock *DefaultBB,
     ArrayRef<std::pair<EnumElementDecl *, SILBasicBlock *>> CaseBBs,

test/SILOptimizer/mandatory_inlining_ossa_to_non_ossa.sil

@@ -14,6 +14,12 @@ sil @in_guaranteed_user : $@convention(thin) (@in_guaranteed Builtin.NativeObjec
 sil @guaranteed_user : $@convention(thin) (@guaranteed Builtin.NativeObject) -> ()
 sil @owned_user : $@convention(thin) (@owned Builtin.NativeObject) -> ()
 
+class Klass {}
+enum FakeOptional<T> {
+case some(T)
+case none
+}
+
 ///////////
 // Tests //
 ///////////
@@ -231,3 +237,59 @@ bb0(%0 : $Builtin.NativeObject):
   return %9999 : $()
 }
 
+sil [transparent] [ossa] @term_ossa_callee : $@convention(thin) (@owned Builtin.NativeObject, @owned FakeOptional<Klass>) -> () {
+bb0(%0 : @owned $Builtin.NativeObject, %1 : @owned $FakeOptional<Klass>):
+  checked_cast_br %0 : $Builtin.NativeObject to $Klass, bb1, bb2
+
+bb1(%1a : @owned $Klass):
+  destroy_value %1a : $Klass
+  br bb3
+
+bb2(%2 : @owned $Builtin.NativeObject):
+  destroy_value %2 : $Builtin.NativeObject
+  br bb3
+
+bb3:
+  switch_enum %1 : $FakeOptional<Klass>, case #FakeOptional.some: bb4, default bb5
+
+bb4(%3 : @owned $Klass):
+  destroy_value %3 : $Klass
+  br bb6
+
+bb5(%4 : @owned $FakeOptional<Klass>):
+  destroy_value %4 : $FakeOptional<Klass>
+  br bb6
+
+bb6:
+  %9999 = tuple()
+  return %9999 : $()
+}
+
+// CHECK-LABEL: sil @term_non_ossa_caller : $@convention(thin) (@owned Builtin.NativeObject, @owned FakeOptional<Klass>) -> () {
+// CHECK: bb0([[ARG0:%.*]] : $Builtin.NativeObject, [[ARG1:%.*]] : $FakeOptional<Klass>):
+// CHECK:   checked_cast_br [[ARG0]] : $Builtin.NativeObject to $Klass, [[YES:bb[0-9]+]], [[NO:bb[0-9]+]]
+//
+// CHECK: [[YES]]([[SUCC:%.*]] :
+// CHECK:   release_value [[SUCC]]
+// CHECK:   br [[CONT:bb[0-9]+]]
+//
+// CHECK: [[NO]]:
+// CHECK:   release_value [[ARG0]]
+// CHECK:   br [[CONT]]
+//
+// CHECK: [[CONT]]:
+// CHECK:   switch_enum [[ARG1]] : $FakeOptional<Klass>, case #FakeOptional.some!enumelt.1: [[SOME:bb[0-9]+]], default [[DEFAULT:bb[0-9]+]]
+//
+// CHECK: [[SOME]]([[PAYLOAD:%.*]] :
+// CHECK:   release_value [[PAYLOAD]]
+//
+// CHECK: [[DEFAULT]]
+// CHECK:   release_value [[ARG1]]
+// CHECK: } // end sil function 'term_non_ossa_caller'
+sil @term_non_ossa_caller : $@convention(thin) (@owned Builtin.NativeObject, @owned FakeOptional<Klass>) -> () {
+bb0(%0 : $Builtin.NativeObject, %1 : $FakeOptional<Klass>):
+  %2 = function_ref @term_ossa_callee : $@convention(thin) (@owned Builtin.NativeObject, @owned FakeOptional<Klass>) -> ()
+  apply %2(%0, %1) : $@convention(thin) (@owned Builtin.NativeObject, @owned FakeOptional<Klass>) -> ()
+  %9999 = tuple()
+  return %9999 : $()
+}