[sil-combine] Add ownership support for switch_enum_addr transforms.

I implemented specifically the promotion from switch_enum_addr -> switch_enum
for loadable types.

NOTE: I did not add support for the inject_enum_addr based conversion to br since that
optimization deletes an edge from the CFG, potentially invalidating CFG
properties. Since OSSA needs to use DeadEndBlocks, we want to avoid changing
reachability in any way with our SILCombine transforms.

Author: gottesmm

lib/SILOptimizer/SILCombiner/SILCombinerMiscVisitors.cpp

@@ -147,20 +147,24 @@ static EnumElementDecl *getInjectEnumCaseTo(SILValue Addr) {
 }
 
 SILInstruction *SILCombiner::visitSwitchEnumAddrInst(SwitchEnumAddrInst *SEAI) {
-  if (SEAI->getFunction()->hasOwnership())
-    return nullptr;
+  SILValue Addr = SEAI->getOperand();
 
   // Convert switch_enum_addr -> br
-  // if the only thing which writes to the address is an inject_enum_addr.
-  SILValue Addr = SEAI->getOperand();
-  if (EnumElementDecl *EnumCase = getInjectEnumCaseTo(Addr)) {
-    SILBasicBlock *Dest = SEAI->getCaseDestination(EnumCase);
-    // If the only instruction which writes to Addr is an inject_enum_addr we
-    // know that there cannot be an enum payload.
-    assert(Dest->getNumArguments() == 0 &&
-           "didn't expect a payload argument");
-    Builder.createBranch(SEAI->getLoc(), Dest);
-    return eraseInstFromFunction(*SEAI);
+  //
+  // If the only thing which writes to the address is an inject_enum_addr. We
+  // only perform these optimizations when we are not in OSSA since this
+  // eliminates an edge from the CFG and we want SILCombine in OSSA to never do
+  // that, so in the future we can invalidate less.
+  if (!SEAI->getFunction()->hasOwnership()) {
+    if (EnumElementDecl *EnumCase = getInjectEnumCaseTo(Addr)) {
+      SILBasicBlock *Dest = SEAI->getCaseDestination(EnumCase);
+      // If the only instruction which writes to Addr is an inject_enum_addr we
+      // know that there cannot be an enum payload.
+      assert(Dest->getNumArguments() == 0 &&
+             "didn't expect a payload argument");
+      Builder.createBranch(SEAI->getLoc(), Dest);
+      return eraseInstFromFunction(*SEAI);
+    }
   }
 
   SILType Ty = Addr->getType();
@@ -172,15 +176,44 @@ SILInstruction *SILCombiner::visitSwitchEnumAddrInst(SwitchEnumAddrInst *SEAI) {
   //     ->
   //   %value = load %ptr
   //   switch_enum %value
-  SmallVector<std::pair<EnumElementDecl*, SILBasicBlock*>, 8> Cases;
-  for (int i = 0, e = SEAI->getNumCases(); i < e; ++i)
+  //
+  // If we are using ownership, we perform a load_borrow right before the new
+  // switch_enum and end the borrow scope right afterwards.
+  Builder.setCurrentDebugScope(SEAI->getDebugScope());
+  SmallVector<std::pair<EnumElementDecl *, SILBasicBlock *>, 8> Cases;
+  for (int i : range(SEAI->getNumCases())) {
     Cases.push_back(SEAI->getCase(i));
+  }
 
-  Builder.setCurrentDebugScope(SEAI->getDebugScope());
   SILBasicBlock *Default = SEAI->hasDefault() ? SEAI->getDefaultBB() : nullptr;
-  LoadInst *EnumVal = Builder.createLoad(SEAI->getLoc(), Addr,
-                                         LoadOwnershipQualifier::Unqualified);
-  Builder.createSwitchEnum(SEAI->getLoc(), EnumVal, Default, Cases);
+  SILValue EnumVal = Builder.emitLoadBorrowOperation(SEAI->getLoc(), Addr);
+  auto *sei = Builder.createSwitchEnum(SEAI->getLoc(), EnumVal, Default, Cases);
+
+  if (Builder.hasOwnership()) {
+    for (int i : range(sei->getNumCases())) {
+      auto c = sei->getCase(i);
+      if (c.first->hasAssociatedValues()) {
+        auto eltType = Addr->getType().getEnumElementType(
+            c.first, Builder.getModule(), Builder.getTypeExpansionContext());
+        eltType = eltType.getObjectType();
+        if (eltType.isTrivial(Builder.getFunction())) {
+          c.second->createPhiArgument(eltType, OwnershipKind::None);
+        } else {
+          c.second->createPhiArgument(eltType, OwnershipKind::Guaranteed);
+        }
+      }
+      Builder.setInsertionPoint(c.second->front().getIterator());
+      Builder.emitEndBorrowOperation(SEAI->getLoc(), EnumVal);
+    }
+
+    if (auto defaultBlock = sei->getDefaultBBOrNull()) {
+      defaultBlock.get()->createPhiArgument(EnumVal->getType(),
+                                            OwnershipKind::Guaranteed);
+      Builder.setInsertionPoint(defaultBlock.get()->front().getIterator());
+      Builder.emitEndBorrowOperation(SEAI->getLoc(), EnumVal);
+    }
+  }
+
   return eraseInstFromFunction(*SEAI);
 }
 

test/SILOptimizer/sil_combine_enums_ossa.sil

@@ -20,10 +20,10 @@ enum Numerals {
 
 sil [ossa] @external_func: $@convention(thin) () -> ()
 
-// CHECK-LABEL: eliminate_sw_enum_addr
-// XHECK-NOT: switch_enum_addr
-// XHECK: switch_enum
-// XHECK: return
+// CHECK-LABEL: sil [ossa] @eliminate_sw_enum_addr : $@convention(thin) () -> Int {
+// CHECK-NOT: switch_enum_addr
+// CHECK: switch_enum {{%.*}} :
+// CHECK: } // end sil function 'eliminate_sw_enum_addr'
 sil [ossa] @eliminate_sw_enum_addr : $@convention(thin) () -> Int {
 bb0:
   %0 = alloc_stack $Optional<SomeClass>, var, name "x" // users: %2, %4, %5, %17, %19

test/SILOptimizer/sil_combine_ossa.sil

@@ -1531,14 +1531,15 @@ bb0(%0 : @owned $E):
 }
 
 // CHECK-LABEL: sil [ossa] @unchecked_take_enum_data_addr_promotion : $@convention(thin) (@inout FakeOptional<B>) -> @owned B {
-// XHECK: bb0(%0 : $*FakeOptional<B>):
-// XHECK-NEXT: [[BORROW:%.*]] = load_borrow
-// XHECK-NEXT: switch_enum [[BORROW]]
-// XHECK: bb1([[GUARANTEED:%.*]] : @guaranteed
-// XHECK-NEXT: end_borrow [[BORROW]]
+// CHECK: bb0(%0 : $*FakeOptional<B>):
+// CHECK-NEXT: [[BORROW:%.*]] = load_borrow
+// CHECK-NEXT: switch_enum [[BORROW]]
+// CHECK: bb1([[GUARANTEED:%.*]] : @guaranteed
+// CHECK-NEXT: end_borrow [[BORROW]]
 // XHECK-NEXT: load
 // XHECK-NEXT: unchecked_enum_data
 // XHECK-NEXT: return
+// CHECK: } /
 sil [ossa] @unchecked_take_enum_data_addr_promotion : $@convention(thin) (@inout FakeOptional<B>) -> @owned B {
 bb0(%0 : $*FakeOptional<B>):
   switch_enum_addr %0 : $*FakeOptional<B>, case #FakeOptional.some!enumelt: bb1, case #FakeOptional.none!enumelt: bb2
@@ -2477,10 +2478,11 @@ struct FakeInt32 {
 }
 
 // CHECK-LABEL: sil [ossa] @loadable_switchenumaddr_promotion : $@convention(thin) (@in FakeOptional<Builtin.RawPointer>) -> () {
-// XHECK-NOT: switch_enum_addr
-// XHECK: load
-// XHECK-NEXT: switch_enum
-// XHECK-NOT: switch_enum_addr
+// CHECK-NOT: switch_enum_addr
+// CHECK: load
+// CHECK-NEXT: switch_enum
+// CHECK-NOT: switch_enum_addr
+// CHECK: } // end sil function 'loadable_switchenumaddr_promotion'
 sil [ossa] @loadable_switchenumaddr_promotion : $@convention(thin) (@in FakeOptional<Builtin.RawPointer>) -> () {
 bb0(%0 : $*FakeOptional<Builtin.RawPointer>):
   switch_enum_addr %0 : $*FakeOptional<Builtin.RawPointer>, case #FakeOptional.some!enumelt: bb1, case #FakeOptional.none!enumelt: bb2