[sil-combine] Leave the new load at the old load site when promoting unchecked_take_enum_data_addr + load -> load + unchecked_enum_data.

lib/SILOptimizer/SILCombiner/SILCombinerMiscVisitors.cpp

@@ -1751,7 +1751,9 @@ SILInstruction *SILCombiner::visitUncheckedTakeEnumDataAddrInst(
   SILType payloadType = tedai->getType().getObjectType();
 
   // Go back through a second time now that we know all of our users are
-  // loads. Perform the transformation on each load.
+  // loads. Perform the transformation on each load at the load's use site. The
+  // reason that we have to do this is that otherwise we would be hoisting the
+  // loads causing us to need to consider additional ARC issues.
   while (!tedai->use_empty()) {
     auto *use = *tedai->use_begin();
     auto *user = use->getUser();
@@ -1763,40 +1765,52 @@ SILInstruction *SILCombiner::visitUncheckedTakeEnumDataAddrInst(
     }
 
     // Insert a new Load of the enum and extract the data from that.
+    //
+    // NOTE: This is potentially hoisting the load, so we need to insert
+    // compensating destroys.
     auto *svi = cast<SingleValueInstruction>(user);
     SILValue newValue;
     if (auto *oldLoad = dyn_cast<LoadInst>(svi)) {
+      SILBuilderWithScope localBuilder(oldLoad, Builder);
       // If the old load is trivial and our enum addr is non-trivial, we need to
       // use a load_borrow here. We know that the unchecked_enum_data will
       // produce a trivial value meaning that we can just do a
       // load_borrow/immediately end the lifetime here.
       if (oldLoad->getOwnershipQualifier() == LoadOwnershipQualifier::Trivial &&
           !enumAddr->getType().isTrivial(Builder.getFunction())) {
-        Builder.emitScopedBorrowOperation(loc, enumAddr, [&](SILValue newLoad) {
-          newValue = Builder.createUncheckedEnumData(loc, newLoad, enumElt,
-                                                     payloadType);
-        });
+        localBuilder.emitScopedBorrowOperation(
+            loc, enumAddr, [&](SILValue newLoad) {
+              newValue = localBuilder.createUncheckedEnumData(
+                  loc, newLoad, enumElt, payloadType);
+            });
       } else {
-        auto newLoad = Builder.emitLoadValueOperation(
+        auto newLoad = localBuilder.emitLoadValueOperation(
             loc, enumAddr, oldLoad->getOwnershipQualifier());
-        newValue =
-            Builder.createUncheckedEnumData(loc, newLoad, enumElt, payloadType);
+        newValue = localBuilder.createUncheckedEnumData(loc, newLoad, enumElt,
+                                                        payloadType);
       }
     } else if (auto *lbi = cast<LoadBorrowInst>(svi)) {
-      auto newLoad = Builder.emitLoadBorrowOperation(loc, enumAddr);
+      SILBuilderWithScope localBuilder(lbi, Builder);
+      auto newLoad = localBuilder.emitLoadBorrowOperation(loc, enumAddr);
       for (auto ui = lbi->consuming_use_begin(), ue = lbi->consuming_use_end();
            ui != ue; ui = lbi->consuming_use_begin()) {
         // We already checked that all of our uses here are end_borrow above.
         assert(isa<EndBorrowInst>(ui->getUser()) &&
                "Expected only end_borrow consuming uses");
         ui->set(newLoad);
       }
-      newValue =
-          Builder.createUncheckedEnumData(loc, newLoad, enumElt, payloadType);
+      // Any lifetime ending uses of our original load_borrow have been
+      // rewritten by the previous loop to be on the new load_borrow. The reason
+      // that we must do this is end_borrows only are placed on borrow
+      // introducing guaranteed values and our unchecked_enum_data (unlike the
+      // old load_borrow of the same type) is not one.
+      newValue = localBuilder.createUncheckedEnumData(loc, newLoad, enumElt,
+                                                      payloadType);
     }
     assert(newValue);
 
-    // Replace all uses of the old load with the data and erase the old load.
+    // Replace all uses of the old load with the newValue and erase the old
+    // load.
     replaceInstUsesWith(*svi, newValue);
     eraseInstFromFunction(*svi);
   }

test/SILOptimizer/sil_combine_ossa.sil

@@ -68,6 +68,11 @@ sil @use_anyobject_guaranteed : $@convention(thin) (@guaranteed AnyObject) -> ()
 sil @use_generic_obj_guaranteed : $@convention(thin) <τ_0_0> (@guaranteed τ_0_0) -> ()
 sil [ossa] @unknown : $@convention(thin) () -> ()
 
+enum KlassNativeObjEither {
+case lhs(Klass)
+case rhs(Builtin.NativeObject)
+}
+
 //////////////////////
 // Simple DCE Tests //
 //////////////////////
@@ -4824,3 +4829,50 @@ bb1:
 bb2:
   unreachable
 }
+
+// Make sure that when we promote the unchecked_take_enum_data_uses in bb1, bb2,
+// we keep the loads in bb1, bb2 instead of hoisting them to the
+// unchecked_take_enum_data_addr that we are optimizing.
+//
+// CHECK-LABEL: sil [ossa] @unchecked_take_enum_data_addr_promotion_host_load : $@convention(thin) (@inout FakeOptional<KlassNativeObjEither>) -> @owned KlassNativeObjEither {
+// CHECK: bb0([[ARG:%.*]] : $*
+// CHECK:   [[LOADED_ARG:%.*]] = load_borrow [[ARG]]
+// CHECK:   [[FOR_SWITCH:%.*]] = unchecked_enum_data [[LOADED_ARG]]
+// CHECK:   switch_enum [[FOR_SWITCH]] :
+//
+// CHECK: bb1(
+// CHECK-NEXT: end_borrow [[LOADED_ARG]]
+// CHECK-NEXT: [[LOADED_ARG_2:%.*]] = load [copy] [[ARG]]
+// CHECK-NEXT: [[LOADED_ARG_2_EXT:%.*]] = unchecked_enum_data [[LOADED_ARG_2]]
+// CHECK-NEXT: br bb3([[LOADED_ARG_2_EXT]] :
+//
+// CHECK: bb2(
+// CHECK-NEXT: end_borrow [[LOADED_ARG]]
+// CHECK-NEXT: [[LOADED_ARG_3:%.*]] = load [take] [[ARG]]
+// CHECK-NEXT: [[LOADED_ARG_3_EXT:%.*]] = unchecked_enum_data [[LOADED_ARG_3]]
+// CHECK-NEXT: [[NONE:%.*]] = enum $FakeOptional<
+// CHECK-NEXT: store [[NONE]] to [init] [[ARG]]
+// CHECK-NEXT: br bb3([[LOADED_ARG_3_EXT]] :
+//
+// CHECK: bb3([[RESULT:%.*]] : @owned
+// CHECK-NEXT: return [[RESULT]]
+// CHECK: } // end sil function 'unchecked_take_enum_data_addr_promotion_host_load'
+sil [ossa] @unchecked_take_enum_data_addr_promotion_host_load : $@convention(thin) (@inout FakeOptional<KlassNativeObjEither>) -> @owned KlassNativeObjEither {
+bb0(%0 : $*FakeOptional<KlassNativeObjEither>):
+  %1 = unchecked_take_enum_data_addr %0 : $*FakeOptional<KlassNativeObjEither>, #FakeOptional.some!enumelt
+  switch_enum_addr %1 : $*KlassNativeObjEither, case #KlassNativeObjEither.lhs!enumelt: bb1, case #KlassNativeObjEither.rhs!enumelt: bb2
+
+bb1:
+  %2 = load [copy] %1 : $*KlassNativeObjEither
+  br bb3(%2 : $KlassNativeObjEither)
+
+bb2:
+  %3 = load [take] %1 : $*KlassNativeObjEither
+  %nil = enum $FakeOptional<KlassNativeObjEither>, #FakeOptional.none!enumelt
+  store %nil to [init] %0 : $*FakeOptional<KlassNativeObjEither>
+  br bb3(%3 : $KlassNativeObjEither)
+
+bb3(%4 : @owned $KlassNativeObjEither):
+  return %4 : $KlassNativeObjEither
+}
+