Mem2Reg: Fix lifetime completion for enum case values.

include/swift/SIL/OSSACompleteLifetime.h

@@ -62,15 +62,22 @@ class OSSACompleteLifetime {
   /// TODO: Remove this option.
   bool ForceLivenessVerification;
 
+  /// If true, lifetimes are completed with `end_lifetime` instead of `destroy_value`.
+  /// This can be used e.g. for enum values of non-trivial enum types, which are
+  /// known to be a trivial case.
+  bool nonDestroyingEnd;
+
 public:
   OSSACompleteLifetime(
       SILFunction *function, const DominanceInfo *domInfo,
       DeadEndBlocks &deadEndBlocks,
       HandleTrivialVariable_t handleTrivialVariable = IgnoreTrivialVariable,
-      bool forceLivenessVerification = false)
+      bool forceLivenessVerification = false,
+      bool nonDestroyingEnd = false)
       : domInfo(domInfo), deadEndBlocks(deadEndBlocks),
         completedValues(function), handleTrivialVariable(handleTrivialVariable),
-        ForceLivenessVerification(forceLivenessVerification) {}
+        ForceLivenessVerification(forceLivenessVerification),
+        nonDestroyingEnd(nonDestroyingEnd) {}
 
   /// The kind of boundary at which to complete the lifetime.
   ///

lib/SIL/Utils/OSSACompleteLifetime.cpp

@@ -70,6 +70,7 @@ llvm::cl::opt<bool> VerifyLifetimeCompletion("verify-lifetime-completion",
 
 static SILInstruction *endOSSALifetime(SILValue value,
                                        OSSACompleteLifetime::LifetimeEnd end,
+                                       bool nonDestroyingEnd,
                                        SILBuilder &builder,
                                        DeadEndBlocks &deb) {
   auto loc =
@@ -82,7 +83,11 @@ static SILInstruction *endOSSALifetime(SILValue value,
     if (value->getType().is<SILBoxType>()) {
       return builder.createDeallocBox(loc, value, isDeadEnd);
     }
-    return builder.createDestroyValue(loc, value, DontPoisonRefs, isDeadEnd);
+    if (nonDestroyingEnd && !isDeadEnd) {
+      return builder.createEndLifetime(loc, value);
+    } else {
+      return builder.createDestroyValue(loc, value, DontPoisonRefs, isDeadEnd);
+    }
   }
   if (auto scopedAddress = ScopedAddressValue(value)) {
     return scopedAddress.createScopeEnd(builder.getInsertionPoint(), loc);
@@ -95,6 +100,7 @@ static SILInstruction *endOSSALifetime(SILValue value,
 
 static bool endLifetimeAtLivenessBoundary(SILValue value,
                                           const SSAPrunedLiveness &liveness,
+                                          bool nonDestroyingEnd,
                                           DeadEndBlocks &deb) {
   PrunedLivenessBoundary boundary;
   liveness.computeBoundary(boundary);
@@ -105,17 +111,17 @@ static bool endLifetimeAtLivenessBoundary(SILValue value,
         PrunedLiveness::LifetimeEndingUse) {
       changed = true;
       SILBuilderWithScope::insertAfter(
-          lastUser, [value, &deb](SILBuilder &builder) {
+          lastUser, [value, &deb, nonDestroyingEnd](SILBuilder &builder) {
             endOSSALifetime(value, OSSACompleteLifetime::LifetimeEnd::Boundary,
-                            builder, deb);
+                            nonDestroyingEnd, builder, deb);
           });
     }
   }
   for (SILBasicBlock *edge : boundary.boundaryEdges) {
     changed = true;
     SILBuilderWithScope builder(edge->begin());
-    endOSSALifetime(value, OSSACompleteLifetime::LifetimeEnd::Boundary, builder,
-                    deb);
+    endOSSALifetime(value, OSSACompleteLifetime::LifetimeEnd::Boundary,\
+                    nonDestroyingEnd, builder, deb);
   }
   for (SILNode *deadDef : boundary.deadDefs) {
     SILInstruction *next = nullptr;
@@ -126,8 +132,8 @@ static bool endLifetimeAtLivenessBoundary(SILValue value,
     }
     changed = true;
     SILBuilderWithScope builder(next);
-    endOSSALifetime(value, OSSACompleteLifetime::LifetimeEnd::Boundary, builder,
-                    deb);
+    endOSSALifetime(value, OSSACompleteLifetime::LifetimeEnd::Boundary,
+                    nonDestroyingEnd, builder, deb);
   }
   return changed;
 }
@@ -465,12 +471,13 @@ void OSSACompleteLifetime::visitAvailabilityBoundary(
 
 static bool endLifetimeAtAvailabilityBoundary(SILValue value,
                                               const SSAPrunedLiveness &liveness,
+                                              bool nonDestroyingEnd,
                                               DeadEndBlocks &deb) {
   bool changed = false;
   OSSACompleteLifetime::visitAvailabilityBoundary(
       value, liveness, [&](auto *unreachable, auto end) {
         SILBuilderWithScope builder(unreachable);
-        endOSSALifetime(value, end, builder, deb);
+        endOSSALifetime(value, end, nonDestroyingEnd, builder, deb);
         changed = true;
       });
   return changed;
@@ -479,15 +486,16 @@ static bool endLifetimeAtAvailabilityBoundary(SILValue value,
 static bool endLifetimeAtBoundary(SILValue value,
                                   SSAPrunedLiveness const &liveness,
                                   OSSACompleteLifetime::Boundary boundary,
+                                  bool nonDestroyingEnd,
                                   DeadEndBlocks &deadEndBlocks) {
   bool changed = false;
   switch (boundary) {
   case OSSACompleteLifetime::Boundary::Liveness:
-    changed |= endLifetimeAtLivenessBoundary(value, liveness, deadEndBlocks);
+    changed |= endLifetimeAtLivenessBoundary(value, liveness, nonDestroyingEnd, deadEndBlocks);
     break;
   case OSSACompleteLifetime::Boundary::Availability:
     changed |=
-        endLifetimeAtAvailabilityBoundary(value, liveness, deadEndBlocks);
+        endLifetimeAtAvailabilityBoundary(value, liveness, nonDestroyingEnd, deadEndBlocks);
     break;
   }
   return changed;
@@ -548,7 +556,7 @@ bool OSSACompleteLifetime::analyzeAndUpdateLifetime(
     ASSERT(false && "caller must check for pointer escapes");
   }
   return endLifetimeAtBoundary(scopedAddress.value, liveness, boundary,
-                               deadEndBlocks);
+                               nonDestroyingEnd, deadEndBlocks);
 }
 
 /// End the lifetime of \p value at unreachable instructions.
@@ -572,7 +580,7 @@ bool OSSACompleteLifetime::analyzeAndUpdateLifetime(SILValue value,
     LinearLiveness liveness(value);
     liveness.compute();
     return endLifetimeAtBoundary(value, liveness.getLiveness(), boundary,
-                                 deadEndBlocks);
+                                 nonDestroyingEnd, deadEndBlocks);
   }
   InteriorLiveness liveness(value);
   liveness.compute(domInfo, handleInnerScope);
@@ -586,7 +594,7 @@ bool OSSACompleteLifetime::analyzeAndUpdateLifetime(SILValue value,
     ASSERT(false && "caller must check for pointer escapes");
   }
   return endLifetimeAtBoundary(value, liveness.getLiveness(), boundary,
-                               deadEndBlocks);
+                               nonDestroyingEnd, deadEndBlocks);
 }
 
 namespace swift::test {
@@ -598,16 +606,23 @@ namespace swift::test {
 static FunctionTest OSSACompleteLifetimeTest(
     "ossa_complete_lifetime", [](auto &function, auto &arguments, auto &test) {
       SILValue value = arguments.takeValue();
-      OSSACompleteLifetime::Boundary kind =
-          llvm::StringSwitch<OSSACompleteLifetime::Boundary>(
-              arguments.takeString())
-              .Case("liveness", OSSACompleteLifetime::Boundary::Liveness)
-              .Case("availability",
-                    OSSACompleteLifetime::Boundary::Availability);
+      OSSACompleteLifetime::Boundary kind = OSSACompleteLifetime::Boundary::Liveness;
+      bool nonDestroyingEnd = false;
+      StringRef kindArg = arguments.takeString();
+      if (kindArg == "liveness_no_destroy") {
+        nonDestroyingEnd = true;
+      } else if (kindArg == "availability") {
+        kind = OSSACompleteLifetime::Boundary::Availability;
+      } else {
+        assert(kindArg == "liveness");
+      }
       auto *deb = test.getDeadEndBlocks();
       llvm::outs() << "OSSA lifetime completion on " << kind
                    << " boundary: " << value;
-      OSSACompleteLifetime completion(&function, /*domInfo*/ nullptr, *deb);
+      OSSACompleteLifetime completion(&function, /*domInfo*/ nullptr, *deb,
+                                      OSSACompleteLifetime::IgnoreTrivialVariable,
+                                      /*forceLivenessVerification=*/false,
+                                      nonDestroyingEnd);
       completion.completeOSSALifetime(value, kind);
       function.print(llvm::outs());
     });

lib/SILOptimizer/Transforms/SILMem2Reg.cpp

@@ -1823,8 +1823,15 @@ void StackAllocationPromoter::run(BasicBlockSetVector &livePhiBlocks) {
   deleter.forceDeleteWithUsers(asi);
 
   // Now, complete lifetimes!
+  // We know that the incomplete values are trivial enum cases. Therefore we
+  // complete their lifetimes with `end_lifetime` instead of `destroy_value`.
+  // This is especially important for embedded swift where we are not allowed
+  // to insert destroys which were not there before.
   OSSACompleteLifetime completion(function, domInfo,
-                                  *deadEndBlocksAnalysis->get(function));
+                                  *deadEndBlocksAnalysis->get(function),
+                                  OSSACompleteLifetime::IgnoreTrivialVariable,
+                                  /*forceLivenessVerification=*/false,
+                                  /*nonDestroyingEnd=*/true);
 
   // We may have incomplete lifetimes for enum locations on trivial paths.
   // After promoting them, complete lifetime here.

test/SILOptimizer/mem2reg_ossa_nontrivial.sil

@@ -58,6 +58,11 @@ enum KlassOptional {
   case none
 }
 
+enum Result<S, F> {
+  case success(S)
+  case failure(F)
+}
+
 sil [ossa] @get_nontrivialstruct : $@convention(thin) () -> @owned NonTrivialStruct
 sil [ossa] @get_nontrivialenum : $@convention(thin) () -> @owned NonTrivialEnum
 sil [ossa] @get_optionalnontrivialstruct : $@convention(thin) () -> @owned FakeOptional<NonTrivialStruct>
@@ -1549,3 +1554,40 @@ bb9:
   dealloc_stack %5
   br bb3
 }
+
+// CHECK-LABEL: sil [ossa] @test_end_lifetime_of_trivial_enum_case :
+// CHECK-NOT:     alloc_stack
+// CHECK:       bb2(%{{[0-9]+}} : $Int): 
+// CHECK:         fix_lifetime
+// CHECK:         end_borrow
+// CHECK:         end_lifetime %1
+// CHECK:       bb3:
+// CHECK-LABEL: } // end sil function 'test_end_lifetime_of_trivial_enum_case'
+sil [ossa] @test_end_lifetime_of_trivial_enum_case : $@convention(thin) (@inout X, @owned Result<X, Int>) -> () {
+bb0(%0 : $*X, %1 : @owned $Result<X, Int>):
+  %2 = alloc_stack $Result<X, Int>
+  store %1 to [init] %2
+  %3 = load_borrow %2
+  switch_enum %3, case #Result.success!enumelt: bb1, case #Result.failure!enumelt: bb2
+
+bb1(%10 : @guaranteed $X):
+  end_borrow %3
+  %12 = load [take] %2
+  %13 = unchecked_enum_data %12 : $Result<X, Int>, #Result.success!enumelt
+  store %13 to [assign] %0
+  br bb3
+
+bb2(%20 : $Int):
+  end_borrow %3
+  %22 = load_borrow %2
+  %23 = unchecked_enum_data %22, #Result.failure!enumelt
+  fix_lifetime %23
+  end_borrow %22
+  br bb3
+
+bb3:
+  dealloc_stack %2
+  %r = tuple ()
+  return %r
+}
+

test/SILOptimizer/ossa_lifetime_completion.sil

@@ -17,6 +17,12 @@ public enum FakeOptional<T> {
   case some(T)
 }
 
+public enum E<T> {
+  case a(T)
+  case b(T)
+  case c
+}
+
 sil @swift_asyncLet_finish : $@convention(thin) @async (Builtin.RawPointer, Builtin.RawPointer) -> ()
 sil @swift_asyncLet_get : $@convention(thin) @async (Builtin.RawPointer, Builtin.RawPointer) -> ()
 sil @getC : $@convention(thin) () -> (@owned C)
@@ -123,6 +129,38 @@ bb3:
   return %r : $()
 }
 
+// CHECK-LABEL: sil [ossa] @enumTest_no_destroy : $@convention(method) (@guaranteed E<C>) -> () {
+// CHECK:       bb2(%{{[0-9]+}} : @guaranteed $C):
+// CHECK:         destroy_value [dead_end]
+// CHECK-NEXT:    unreachable
+// CHECK:       bb3:
+// CHECK:         end_lifetime
+// CHECK-NEXT:    br bb4
+// CHECK-LABEL: } // end sil function 'enumTest_no_destroy'
+sil [ossa] @enumTest_no_destroy : $@convention(method) (@guaranteed E<C>) -> () {
+bb0(%0 : @guaranteed $E<C>):
+  specify_test "ossa_complete_lifetime @instruction[0] liveness_no_destroy"
+  %copy = copy_value %0 : $E<C>
+  %borrow = begin_borrow %copy : $E<C>
+  switch_enum %borrow : $E<C>, case #E.a!enumelt: bb1, case #E.b!enumelt: bb2, case #E.c!enumelt: bb3
+
+bb1(%10 : @guaranteed $C):
+  end_borrow %borrow : $E<C>
+  destroy_value %copy : $E<C>
+  br bb4
+
+bb2(%15 : @guaranteed $C):
+  unreachable
+
+bb3:
+  end_borrow %borrow : $E<C>
+  br bb4
+
+bb4:
+  %r = tuple ()
+  return %r : $()
+}
+
 sil @use_guaranteed : $@convention(thin) (@guaranteed C) -> ()
 
 sil [ossa] @argTest : $@convention(method) (@owned C) -> () {