SILCombine: make the alloc_stack-enum optimization a bit more tolerant.

When eliminating an enum from an alloc_stack, accept "dead" inject_enum_addr instructions, which inject different enum cases.

Author: eeckstein

lib/SILOptimizer/SILCombiner/SILCombinerMiscVisitors.cpp

@@ -453,6 +453,10 @@ bool SILCombiner::optimizeStackAllocatedEnum(AllocStackInst *AS) {
     return false;
 
   EnumElementDecl *element = nullptr;
+  unsigned numInits =0;
+  unsigned numTakes = 0;
+  SILBasicBlock *initBlock = nullptr;
+  SILBasicBlock *takeBlock = nullptr;
   SILType payloadType;
 
   // First step: check if the stack location is only used to hold one specific
@@ -463,6 +467,8 @@ bool SILCombiner::optimizeStackAllocatedEnum(AllocStackInst *AS) {
       case SILInstructionKind::DebugValueAddrInst:
       case SILInstructionKind::DestroyAddrInst:
       case SILInstructionKind::DeallocStackInst:
+      case SILInstructionKind::InjectEnumAddrInst:
+        // We'll check init_enum_addr below.
         break;
       case SILInstructionKind::InitEnumDataAddrInst: {
         auto *ieda = cast<InitEnumDataAddrInst>(user);
@@ -472,20 +478,17 @@ bool SILCombiner::optimizeStackAllocatedEnum(AllocStackInst *AS) {
         element = el;
         assert(!payloadType || payloadType == ieda->getType());
         payloadType = ieda->getType();
-        break;
-      }
-      case SILInstructionKind::InjectEnumAddrInst: {
-        auto *el = cast<InjectEnumAddrInst>(user)->getElement();
-        if (element && el != element)
-          return false;
-        element = el;
+        numInits++;
+        initBlock = user->getParent();
         break;
       }
       case SILInstructionKind::UncheckedTakeEnumDataAddrInst: {
         auto *el = cast<UncheckedTakeEnumDataAddrInst>(user)->getElement();
         if (element && el != element)
           return false;
         element = el;
+        numTakes++;
+        takeBlock = user->getParent();
         break;
       }
       default:
@@ -495,6 +498,24 @@ bool SILCombiner::optimizeStackAllocatedEnum(AllocStackInst *AS) {
   if (!element || !payloadType)
     return false;
 
+  // If the enum has a single init-take pair in a single block, we know that
+  // the enum cannot contain any valid payload outside that init-take pair.
+  //
+  // This also means that we can ignore any inject_enum_addr of another enum
+  // case, because this can only inject a case without a payload.
+  bool singleInitTakePair =
+    (numInits == 1 && numTakes == 1 && initBlock == takeBlock);
+  if (!singleInitTakePair) {
+    // No single init-take pair: We cannot ignore inject_enum_addrs with a
+    // mismatching case.
+    for (auto *use : AS->getUses()) {
+      if (auto *inject = dyn_cast<InjectEnumAddrInst>(use->getUser())) {
+        if (inject->getElement() != element)
+          return false;
+      }
+    }
+  }
+
   // Second step: replace the enum alloc_stack with a payload alloc_stack.
   auto *newAlloc = Builder.createAllocStack(
       AS->getLoc(), payloadType, AS->getVarInfo(), AS->hasDynamicLifetime());
@@ -508,6 +529,22 @@ bool SILCombiner::optimizeStackAllocatedEnum(AllocStackInst *AS) {
         eraseInstFromFunction(*user);
         break;
       case SILInstructionKind::DestroyAddrInst:
+        if (singleInitTakePair) {
+          // It's not possible that the enum has a payload at the destroy_addr,
+          // because it must have already been taken by the take of the
+          // single init-take pair.
+          // We _have_ to remove the destroy_addr, because we also remove all
+          // inject_enum_addrs which might inject a payload-less case before
+          // the destroy_addr.
+          eraseInstFromFunction(*user);
+        } else {
+          // The enum payload can still be valid at the destroy_addr, so we have
+          // to keep the destroy_addr. Just replace the enum with the payload
+          // (and because it's not a singleInitTakePair, we can be sure that the
+          // enum cannot have any other case than the payload case).
+          use->set(newAlloc);
+        }
+        break;
       case SILInstructionKind::DeallocStackInst:
         use->set(newAlloc);
         break;

test/SILOptimizer/cast_optimizer_conditional_conformance.sil

@@ -98,8 +98,6 @@ bb6:
 // CHECK: [[IEDA:%.*]] = init_enum_data_addr [[OPT]] : $*Optional<HasFoo>, #Optional.some!enumelt
 // CHECK: checked_cast_addr_br take_always S1<Int8> in [[VAL]] : $*S1<Int8> to HasFoo in [[IEDA]] : $*HasFoo, bb1, bb2
 // bbs...
-// CHECK: switch_enum_addr [[OPT]] : $*Optional<HasFoo>, case #Optional.some!enumelt: bb4, case #Optional.none!enumelt: bb5
-// bbs...
 // CHECK: [[UNWRAP:%.*]] = unchecked_take_enum_data_addr [[OPT]] : $*Optional<HasFoo>, #Optional.some!enumelt
 // CHECK: copy_addr [take] [[UNWRAP]] to [initialization] [[EXIS]] : $*HasFoo
 // CHECK: [[OPEN:%.*]] = open_existential_addr immutable_access [[EXIS]] : $*HasFoo to $*@opened("4E16CBC0-FD9F-11E8-A311-D0817AD9F6DD") HasFoo

test/SILOptimizer/sil_combine_enums.sil

@@ -504,6 +504,7 @@ enum MP {
 }
 
 sil @take_s : $@convention(thin) (@in S) -> ()
+sil @take_t : $@convention(thin) (@in T) -> ()
 sil @use_mp : $@convention(thin) (@in_guaranteed MP) -> ()
 
 // CHECK-LABEL: sil @expand_alloc_stack_of_enum1
@@ -568,6 +569,37 @@ bb3:
   return %11 : $()
 }
 
+// CHECK-LABEL: sil @expand_alloc_stack_of_enum_multiple_cases
+// CHECK:        [[A:%[0-9]+]] = alloc_stack $T
+// CHECK:      bb1:
+// CHECK-NEXT:   store %0 to [[A]]
+// CHECK:        apply {{%[0-9]+}}([[A]])
+// CHECK:      bb2:
+// CHECK-NEXT:   br bb3
+// CHECK:      bb3:
+// CHECK: } // end sil function 'expand_alloc_stack_of_enum_multiple_cases'
+sil @expand_alloc_stack_of_enum_multiple_cases : $@convention(method) (T) -> () {
+bb0(%0 : $T):
+  %1 = alloc_stack $X
+  cond_br undef, bb1, bb2
+bb1:
+  %2 = init_enum_data_addr %1 : $*X, #X.some!enumelt
+  store %0 to %2 : $*T
+  inject_enum_addr %1 : $*X, #X.some!enumelt
+  %7 = unchecked_take_enum_data_addr %1 : $*X, #X.some!enumelt
+  %8 = function_ref @take_t : $@convention(thin) (@in T) -> ()
+  %9 = apply %8(%7) : $@convention(thin) (@in T) -> ()
+  br bb3
+bb2:
+  inject_enum_addr %1 : $*X, #X.none!enumelt
+  destroy_addr %1 : $*X
+  br bb3
+bb3:
+  dealloc_stack %1 : $*X
+  %11 = tuple ()
+  return %11 : $()
+}
+
 // CHECK-LABEL: sil @dont_expand_alloc_stack_of_enum_multiple_cases
 // CHECK:   alloc_stack $MP
 // CHECK: } // end sil function 'dont_expand_alloc_stack_of_enum_multiple_cases'
@@ -592,6 +624,28 @@ bb3:
   return %11 : $()
 }
 
+// CHECK-LABEL: sil @dont_expand_alloc_stack_of_enum_multiple_cases2
+// CHECK:   alloc_stack $X
+// CHECK: } // end sil function 'dont_expand_alloc_stack_of_enum_multiple_cases2'
+sil @dont_expand_alloc_stack_of_enum_multiple_cases2 : $@convention(method) (T) -> () {
+bb0(%0 : $T):
+  %1 = alloc_stack $X
+  cond_br undef, bb1, bb2
+bb1:
+  %2 = init_enum_data_addr %1 : $*X, #X.some!enumelt
+  store %0 to %2 : $*T
+  inject_enum_addr %1 : $*X, #X.some!enumelt
+  br bb3
+bb2:
+  inject_enum_addr %1 : $*X, #X.none!enumelt
+  br bb3
+bb3:
+  destroy_addr %1 : $*X
+  dealloc_stack %1 : $*X
+  %11 = tuple ()
+  return %11 : $()
+}
+
 // CHECK-LABEL: sil @dont_expand_alloc_stack_of_enum_unknown_use
 // CHECK:   alloc_stack $MP
 // CHECK: } // end sil function 'dont_expand_alloc_stack_of_enum_unknown_use'