Merge pull request #30508 from eeckstein/combine-enum

SILCombine: optimize creating enums with tuple payloads.

Author: eeckstein

lib/SILOptimizer/SILCombiner/SILCombinerMiscVisitors.cpp

@@ -945,13 +945,73 @@ SILInstruction *SILCombiner::visitStrongRetainInst(StrongRetainInst *SRI) {
   return nullptr;
 }
 
+/// Create a value from stores to an address.
+///
+/// If there are only stores to \p addr, return the stored value. Also, if there
+/// are address projections, create aggregate instructions for it.
+/// If builder is null, it's just a dry-run to check if it's possible.
+static SILValue createValueFromAddr(SILValue addr, SILBuilder *builder,
+                                    SILLocation loc) {
+  SmallVector<SILValue, 4> elems;
+  enum Kind {
+    none, store, tuple
+  } kind = none;
+
+  for (Operand *use : addr->getUses()) {
+    SILInstruction *user = use->getUser();
+    if (user->isDebugInstruction())
+      continue;
+
+    auto *st = dyn_cast<StoreInst>(user);
+    if (st && kind == none && st->getDest() == addr) {
+      elems.push_back(st->getSrc());
+      kind = store;
+      // We cannot just return st->getSrc() here because we also have to check
+      // if the store destination is the only use of addr.
+      continue;
+    }
+
+    if (auto *telem = dyn_cast<TupleElementAddrInst>(user)) {
+      if (kind == none) {
+        elems.resize(addr->getType().castTo<TupleType>()->getNumElements());
+        kind = tuple;
+      }
+      if (kind == tuple) {
+        if (elems[telem->getFieldNo()])
+          return SILValue();
+        elems[telem->getFieldNo()] = createValueFromAddr(telem, builder, loc);
+        continue;
+      }
+    }
+    // TODO: handle StructElementAddrInst to create structs.
+
+    return SILValue();
+  }
+  switch (kind) {
+  case none:
+    return SILValue();
+  case store:
+    assert(elems.size() == 1);
+    return elems[0];
+  case tuple:
+    if (std::any_of(elems.begin(), elems.end(),
+                    [](SILValue v){ return !(bool)v; }))
+      return SILValue();
+    if (builder) {
+      return builder->createTuple(loc, addr->getType().getObjectType(), elems);
+    }
+    // Just return anything not null for the dry-run.
+    return elems[0];
+  }
+}
+
 /// Simplify the following two frontend patterns:
 ///
 ///   %payload_addr = init_enum_data_addr %payload_allocation
 ///   store %payload to %payload_addr
 ///   inject_enum_addr %payload_allocation, $EnumType.case
 ///
-///   inject_enum_add %nopayload_allocation, $EnumType.case
+///   inject_enum_addr %nopayload_allocation, $EnumType.case
 ///
 /// for a concrete enum type $EnumType.case to:
 ///
@@ -1144,16 +1204,6 @@ SILCombiner::visitInjectEnumAddrInst(InjectEnumAddrInst *IEAI) {
   }
   assert((EnumAddrIns == IEAI) &&
          "Found InitEnumDataAddrInst differs from IEAI");
-  // Found the DataAddrInst to this enum payload. Check if it has only use.
-  if (!hasOneNonDebugUse(DataAddrInst))
-    return nullptr;
-
-  auto *SI = dyn_cast<StoreInst>(getSingleNonDebugUser(DataAddrInst));
-  auto *AI = dyn_cast<ApplyInst>(getSingleNonDebugUser(DataAddrInst));
-  if (!SI && !AI) {
-    return nullptr;
-  }
-
   // Make sure the enum pattern instructions are the only ones which write to
   // this location
   if (!WriteSet.empty()) {
@@ -1201,18 +1251,30 @@ SILCombiner::visitInjectEnumAddrInst(InjectEnumAddrInst *IEAI) {
     }
   }
 
-  if (SI) {
-    assert((SI->getDest() == DataAddrInst) &&
-           "Can't find StoreInst with DataAddrInst as its destination");
+  // Check if we can replace all stores to the enum data with an enum of the
+  // stored value. We can also handle tuples as payloads, e.g.
+  //
+  //   %payload_addr = init_enum_data_addr %enum_addr
+  //   %elem0_addr = tuple_element_addr %payload_addr, 0
+  //   %elem1_addr = tuple_element_addr %payload_addr, 1
+  //   store %payload0 to %elem0_addr
+  //   store %payload1 to %elem1_addr
+  //   inject_enum_addr %enum_addr, $EnumType.case
+  //
+  if (createValueFromAddr(DataAddrInst, nullptr, DataAddrInst->getLoc())) {
+    SILValue en =
+      createValueFromAddr(DataAddrInst, &Builder, DataAddrInst->getLoc());
+    assert(en);
+
     // In that case, create the payload enum/store.
     EnumInst *E = Builder.createEnum(
-        DataAddrInst->getLoc(), SI->getSrc(), DataAddrInst->getElement(),
+        DataAddrInst->getLoc(), en, DataAddrInst->getElement(),
         DataAddrInst->getOperand()->getType().getObjectType());
     Builder.createStore(DataAddrInst->getLoc(), E, DataAddrInst->getOperand(),
                         StoreOwnershipQualifier::Unqualified);
     // Cleanup.
-    eraseInstFromFunction(*SI);
-    eraseInstFromFunction(*DataAddrInst);
+    eraseUsesOfInstruction(DataAddrInst);
+    recursivelyDeleteTriviallyDeadInstructions(DataAddrInst, true);
     return eraseInstFromFunction(*IEAI);
   }
 
@@ -1230,7 +1292,9 @@ SILCombiner::visitInjectEnumAddrInst(InjectEnumAddrInst *IEAI) {
   //  %1 = enum $EnumType, $EnumType.case, %load
   //  store %1 to %nopayload_addr
   //
-  assert(AI && "Must have an apply");
+  auto *AI = dyn_cast_or_null<ApplyInst>(getSingleNonDebugUser(DataAddrInst));
+  if (!AI)
+    return nullptr;
   unsigned ArgIdx = 0;
   Operand *EnumInitOperand = nullptr;
   for (auto &Opd : AI->getArgumentOperands()) {

test/SILOptimizer/opt_enumerate.swift

@@ -0,0 +1,26 @@
+// RUN: %target-swift-frontend %s -O -module-name=test -emit-sil | %FileCheck %s
+// REQUIRES: swift_stdlib_no_asserts,optimized_stdlib
+
+var gg = 0
+
+@inline(never)
+func take(_ x: Int, _ y: Int) {
+  gg = x + y
+}
+
+// CHECK-LABEL: sil @$s4test23check_cond_fail_messageySiSaySiGF
+// CHECK:  cond_fail {{.*}} "Index out of range"
+// CHECK: // end sil function '$s4test23check_cond_fail_messageySiSaySiGF'
+public func check_cond_fail_message(_ array: [Int]) -> Int {
+  return array[2]
+}
+
+// CHECK-LABEL: sil @$s4test22eliminate_bounds_checkyySaySiGF
+// CHECK-NOT:  cond_fail {{.*}} "Index out of range"
+// CHECK: // end sil function '$s4test22eliminate_bounds_checkyySaySiGF'
+public func eliminate_bounds_check(_ array: [Int]) {
+  for (index, x) in array.enumerated() {
+      take(x, index)
+  }
+}
+

test/SILOptimizer/sil_combine_enums.sil

@@ -473,3 +473,26 @@ bb3(%14 : $C):
 // CHECK: return
   return %16 : $T
 }
+
+// CHECK-LABEL: sil @test_inject_tuple
+// CHECK:   [[A:%[0-9]+]] = alloc_stack $Optional<(Int, Int)>
+// CHECK:   [[T:%[0-9]+]] = tuple (%0 : $Int, %1 : $Int)
+// CHECK:   [[E:%[0-9]+]] = enum $Optional<(Int, Int)>, #Optional.some!enumelt.1, [[T]] : $(Int, Int)
+// CHECK:   store [[E]] to [[A]]
+// CHECK:   [[L:%[0-9]+]] = load [[A]]
+// CHECK:   return [[L]]
+// CHECK: } // end sil function 'test_inject_tuple'
+sil @test_inject_tuple : $@convention(thin) (Int, Int) -> Optional<(Int, Int)> {
+bb0(%0 : $Int, %1 : $Int):
+  %17 = alloc_stack $Optional<(Int, Int)>
+  %45 = init_enum_data_addr %17 : $*Optional<(Int, Int)>, #Optional.some!enumelt.1
+  %46 = tuple_element_addr %45 : $*(Int, Int), 0
+  %47 = tuple_element_addr %45 : $*(Int, Int), 1
+  store %0 to %46 : $*Int
+  store %1 to %47 : $*Int
+  inject_enum_addr %17 : $*Optional<(Int, Int)>, #Optional.some!enumelt.1
+  %r = load %17 : $*Optional<(Int, Int)>
+  dealloc_stack %17 : $*Optional<(Int, Int)>
+  return %r : $Optional<(Int, Int)>
+}
+