SILGen OSSA support for switch_enum/checked_cast_br, related cleanup

Use APIs for creating terminator results that handle forwarding
ownership consistently.

Add ManagedValue::forForwardedRValue(SILValue) to handle cleanups
consistently based on ownership forwarding.

Add SILGenBuilder::createForwardedTermResult(SILType type) for
creating termator results with the correct ownership and cleanups.

Add SILGenBuilder::createTermResult(SILType type, ValueOwnershipKind
ownership) that handles cleanup based on terminator result ownership.

Add SILGenBuilder::createOptionalSomeResult(SwitchEnumInst) so a lot
of code doesn't need to deal with unwrapping Optional types,
terminator results, and ownership rules.

Replace the existing "phi" APIs with a single
SILGenBuilder::createPhi(SILType, ValueOwnershipKind) that handles
cleanup based on phi ownership.

Phis and terminator results are fundamentally different and need to be handled differently everywhere. Remove the confusion where terminator results were generated with a "phi argument" API.

Author: atrick

lib/SILGen/ManagedValue.cpp

@@ -22,6 +22,25 @@
 using namespace swift;
 using namespace Lowering;
 
+ManagedValue ManagedValue::forForwardedRValue(SILGenFunction &SGF,
+                                              SILValue value) {
+  if (!value)
+    return ManagedValue();
+
+  switch (value->getOwnershipKind()) {
+  case OwnershipKind::Any:
+    llvm_unreachable("Invalid ownership for value");
+
+  case OwnershipKind::Owned:
+    return ManagedValue(value, SGF.enterDestroyCleanup(value));
+
+  case OwnershipKind::Guaranteed:
+  case OwnershipKind::None:
+  case OwnershipKind::Unowned:
+    return ManagedValue::forUnmanaged(value);
+  }
+}
+
 /// Emit a copy of this value with independent ownership.
 ManagedValue ManagedValue::copy(SILGenFunction &SGF, SILLocation loc) const {
   auto &lowering = SGF.getTypeLowering(getType());

lib/SILGen/ManagedValue.h

@@ -91,6 +91,14 @@ class ManagedValue {
     return ManagedValue(value, false, CleanupHandle::invalid());
   }
 
+  /// Create a managed value for a SILValue whose ownership is
+  /// forwarded. Creates a new cleanup for +1 values. Forwarded +0 values
+  /// require no cleanup.
+  ///
+  /// Use this for values that do not introduce a new borrow scope. This is
+  /// correct for casts and terminator results, not for phis.
+  static ManagedValue forForwardedRValue(SILGenFunction &SGF, SILValue value);
+
   /// Create a managed value for a +1 rvalue object.
   static ManagedValue forOwnedObjectRValue(SILValue value,
                                            CleanupHandle cleanup) {

lib/SILGen/ResultPlan.cpp

@@ -667,9 +667,9 @@ class ForeignAsyncInitializationPlan final : public ResultPlan {
 
       auto errorTy = SGF.getASTContext().getErrorDecl()->getDeclaredType()
         ->getCanonicalType();
-      auto errorVal
-        = SGF.B.createOwnedPhiArgument(SILType::getPrimitiveObjectType(errorTy));
-      
+      auto errorVal = SGF.B.createTermResult(
+        SILType::getPrimitiveObjectType(errorTy), OwnershipKind::Owned);
+
       SGF.emitThrow(loc, errorVal, true);
     }
 

lib/SILGen/SILGenBridging.cpp

@@ -1715,17 +1715,15 @@ void SILGenFunction::emitNativeToForeignThunk(SILDeclRef thunk) {
         {getASTContext().getOptionalSomeDecl(), hasCompletionBB},
         {getASTContext().getOptionalNoneDecl(), noCompletionBB},
       };
-      
-      B.createSwitchEnum(loc, completionBlock, nullptr, dests);
-      
+
+      auto *switchEnum =
+          B.createSwitchEnum(loc, completionBlock, nullptr, dests);
+
       B.emitBlock(noCompletionBB);
       B.createBranch(loc, doneBBOrNull);
 
       B.emitBlock(hasCompletionBB);
-      completionBlock = hasCompletionBB->createPhiArgument(
-                                  SILType::getPrimitiveObjectType(completionTy),
-                                  OwnershipKind::Guaranteed);
-      
+      completionBlock = switchEnum->createOptionalSomeResult();
     }
   };
 

lib/SILGen/SILGenBuilder.cpp

@@ -187,23 +187,38 @@ ManagedValue SILGenBuilder::createCopyValue(SILLocation loc,
   }
 #include "swift/AST/ReferenceStorage.def"
 
-ManagedValue SILGenBuilder::createOwnedPhiArgument(SILType type) {
-  SILPhiArgument *arg =
-      getInsertionBB()->createPhiArgument(type, OwnershipKind::Owned);
-  return SGF.emitManagedRValueWithCleanup(arg);
+ManagedValue SILGenBuilder::createForwardedTermResult(SILType type) {
+  auto *succBB = getInsertionBB();
+  auto *term = cast<OwnershipForwardingTermInst>(
+      succBB->getSinglePredecessorBlock()->getTerminator());
+  auto *arg = term->createResult(succBB, type);
+  return ManagedValue::forForwardedRValue(SGF, arg);
 }
 
-ManagedValue SILGenBuilder::createGuaranteedPhiArgument(SILType type) {
-  SILPhiArgument *arg =
-      getInsertionBB()->createPhiArgument(type, OwnershipKind::Guaranteed);
-  return SGF.emitManagedBorrowedArgumentWithCleanup(arg);
+ManagedValue SILGenBuilder::createTermResult(SILType type,
+                                             ValueOwnershipKind ownership) {
+  // Despite the name, 'arg' is a terminator result, not a phi.
+  auto *arg = getInsertionBB()->createPhiArgument(type, ownership);
+  return ManagedValue::forForwardedRValue(SGF, arg);
 }
 
-ManagedValue
-SILGenBuilder::createGuaranteedTransformingTerminatorArgument(SILType type) {
-  SILPhiArgument *arg =
-      getInsertionBB()->createPhiArgument(type, OwnershipKind::Guaranteed);
-  return ManagedValue::forUnmanaged(arg);
+ManagedValue SILGenBuilder::createPhi(SILType type,
+                                      ValueOwnershipKind ownership) {
+  SILPhiArgument *arg = getInsertionBB()->createPhiArgument(type, ownership);
+  switch (ownership) {
+  case OwnershipKind::Any:
+    llvm_unreachable("Invalid ownership for value");
+
+  case OwnershipKind::Owned:
+    return SGF.emitManagedRValueWithCleanup(arg);
+
+  case OwnershipKind::Guaranteed:
+    return SGF.emitManagedBorrowedArgumentWithCleanup(arg);
+
+  case OwnershipKind::None:
+  case OwnershipKind::Unowned:
+    return ManagedValue::forUnmanaged(arg);
+  }
 }
 
 ManagedValue SILGenBuilder::createAllocRef(

lib/SILGen/SILGenBuilder.h

@@ -125,17 +125,34 @@ class SILGenBuilder : public SILBuilder {
                                              ManagedValue originalValue);
 #include "swift/AST/ReferenceStorage.def"
 
-  ManagedValue createOwnedPhiArgument(SILType type);
-  ManagedValue createGuaranteedPhiArgument(SILType type);
+  /// Create the block argument for an OwnershipForwardingTermInst result
+  /// (checked_cast_br or switch_enum). Allows creating terminator
+  /// results separately from creating the terminator. Alternatively, the result
+  /// can be directly created as follows:
+  ///
+  ///   ManagedValue::forForwardedRValue(term->createResult(succBB, resultTy))
+  ///
+  /// For a switch_enum with a default payload, use:
+  ///
+  ///   ManagedValue::forForwardedRValue(switchEnum->createDefaultResult())
+  ///
+  ManagedValue createForwardedTermResult(SILType type);
 
-  /// For arguments from terminators that are "transforming terminators". These
-  /// types of guaranteed arguments are validated as part of the operand of the
-  /// transforming terminator since transforming terminators are guaranteed to
-  /// be the only predecessor of our parent block.
+  /// Create a terminator result with specified ownership.
   ///
-  /// NOTE: Two examples of transforming terminators are switch_enum,
-  /// checked_cast_br.
-  ManagedValue createGuaranteedTransformingTerminatorArgument(SILType type);
+  /// Typically for an apply's return or error value, which should use
+  /// OwnershipKind::Owned. For OwnershipForwardingTermInst (checked_cast_br or
+  /// switch_enum), use createForwardedTermResult instead.
+  ManagedValue createTermResult(SILType type, ValueOwnershipKind ownership);
+
+  /// Create the block argument for an Optional switch_enum.
+  ManagedValue createOptionalSomeResult(SwitchEnumInst *switchEnum) {
+    return ManagedValue::forForwardedRValue(
+        SGF, switchEnum->createOptionalSomeResult());
+  }
+
+  // Create the block argument for a phi.
+  ManagedValue createPhi(SILType type, ValueOwnershipKind ownership);
 
   using SILBuilder::createMarkUninitialized;
   ManagedValue createMarkUninitialized(ValueDecl *decl, ManagedValue operand,

lib/SILGen/SILGenBuiltin.cpp

@@ -1533,8 +1533,8 @@ static ManagedValue emitBuiltinWithUnsafeContinuation(
 
     auto errorTy = SGF.getASTContext().getErrorDecl()->getDeclaredType()
       ->getCanonicalType();
-    auto errorVal
-      = SGF.B.createOwnedPhiArgument(SILType::getPrimitiveObjectType(errorTy));
+    auto errorVal = SGF.B.createTermResult(
+        SILType::getPrimitiveObjectType(errorTy), OwnershipKind::Owned);
 
     SGF.emitThrow(loc, errorVal, true);
   }

lib/SILGen/SILGenConvert.cpp

@@ -198,7 +198,9 @@ SILGenFunction::emitPreconditionOptionalHasValue(SILLocation loc,
 
   // If we have an object, make sure the object is at +1. All switch_enum of
   // objects is done at +1.
-  if (optional.getType().isAddress()) {
+  bool isAddress = optional.getType().isAddress();
+  SwitchEnumInst *switchEnum = nullptr;
+  if (isAddress) {
     // We forward in the creation routine for
     // unchecked_take_enum_data_addr. switch_enum_addr is a +0 operation.
     B.createSwitchEnumAddr(loc, optional.getValue(),
@@ -208,9 +210,9 @@ SILGenFunction::emitPreconditionOptionalHasValue(SILLocation loc,
     optional = optional.ensurePlusOne(*this, loc);
     hadCleanup = true;
     hadLValue = false;
-    B.createSwitchEnum(loc, optional.forward(*this),
-                       /*defaultDest*/ nullptr,
-                       {{someDecl, contBB}, {noneDecl, failBB}});
+    switchEnum = B.createSwitchEnum(loc, optional.forward(*this),
+                                    /*defaultDest*/ nullptr,
+                                    {{someDecl, contBB}, {noneDecl, failBB}});
   }
   B.emitBlock(failBB);
 
@@ -241,13 +243,12 @@ SILGenFunction::emitPreconditionOptionalHasValue(SILLocation loc,
   B.emitBlock(contBB);
 
   ManagedValue result;
-  SILType payloadType = optional.getType().getOptionalObjectType();
-
-  if (payloadType.isObject()) {
-    result = B.createOwnedPhiArgument(payloadType);
-  } else {
+  if (isAddress) {
+    SILType payloadType = optional.getType().getOptionalObjectType();
     result =
         B.createUncheckedTakeEnumDataAddr(loc, optional, someDecl, payloadType);
+  } else {
+    result = B.createOptionalSomeResult(switchEnum);
   }
 
   if (hadCleanup) {
@@ -415,15 +416,14 @@ SILGenFunction::emitOptionalToOptional(SILLocation loc,
   // If the result is address-only, we need to return something in memory,
   // otherwise the result is the BBArgument in the merge point.
   // TODO: use the SGFContext passed in.
-  ManagedValue finalResult;
-  if (resultTL.isAddressOnly() && silConv.useLoweredAddresses()) {
-    finalResult = emitManagedBufferWithCleanup(
+  ManagedValue resultAddress;
+  bool addressOnly = resultTL.isAddressOnly() && silConv.useLoweredAddresses();
+  if (addressOnly) {
+    resultAddress = emitManagedBufferWithCleanup(
         emitTemporaryAllocation(loc, resultTy), resultTL);
-  } else {
-    SILGenSavedInsertionPoint IP(*this, contBB);
-    finalResult = B.createOwnedPhiArgument(resultTL.getLoweredType());
   }
 
+  ValueOwnershipKind resultOwnership = OwnershipKind::Any;
   SEBuilder.addOptionalSomeCase(
       isPresentBB, contBB, [&](ManagedValue input, SwitchCaseFullExpr &&scope) {
         // If we have an address only type, we want to match the old behavior of
@@ -439,36 +439,42 @@ SILGenFunction::emitOptionalToOptional(SILLocation loc,
 
         ManagedValue result = transformValue(*this, loc, input, noOptResultTy,
                                              SGFContext());
-
-        if (!(resultTL.isAddressOnly() && silConv.useLoweredAddresses())) {
+        resultOwnership = result.getValue()->getOwnershipKind();
+        if (!addressOnly) {
           SILValue some = B.createOptionalSome(loc, result).forward(*this);
           return scope.exitAndBranch(loc, some);
         }
 
         RValue R(*this, loc, noOptResultTy.getASTType(), result);
         ArgumentSource resultValueRV(loc, std::move(R));
         emitInjectOptionalValueInto(loc, std::move(resultValueRV),
-                                    finalResult.getValue(), resultTL);
+                                    resultAddress.getValue(), resultTL);
         return scope.exitAndBranch(loc);
       });
 
   SEBuilder.addOptionalNoneCase(
       isNotPresentBB, contBB,
       [&](ManagedValue input, SwitchCaseFullExpr &&scope) {
-        if (!(resultTL.isAddressOnly() && silConv.useLoweredAddresses())) {
+        if (!addressOnly) {
           SILValue none =
               B.createManagedOptionalNone(loc, resultTy).forward(*this);
           return scope.exitAndBranch(loc, none);
         }
 
-        emitInjectOptionalNothingInto(loc, finalResult.getValue(), resultTL);
+        emitInjectOptionalNothingInto(loc, resultAddress.getValue(), resultTL);
         return scope.exitAndBranch(loc);
       });
 
   std::move(SEBuilder).emit();
 
   B.emitBlock(contBB);
-  return finalResult;
+  if (addressOnly)
+    return resultAddress;
+
+  // This phi's ownership is derived from the transformed value's
+  // ownership, not the input ownership. Transformation can convert a value with
+  // no ownership to an owned value.
+  return B.createPhi(resultTL.getLoweredType(), resultOwnership);
 }
 
 SILGenFunction::OpaqueValueRAII::~OpaqueValueRAII() {
@@ -712,8 +718,9 @@ ManagedValue SILGenFunction::emitExistentialErasure(
         { ctx.getOptionalSomeDecl(), isPresentBB },
         { ctx.getOptionalNoneDecl(), isNotPresentBB }
       };
-      B.createSwitchEnum(loc, potentialNSError.forward(*this),
-                         /*default*/ nullptr, cases);
+      auto *switchEnum =
+          B.createSwitchEnum(loc, potentialNSError.forward(*this),
+                             /*default*/ nullptr, cases);
 
       // If we did get an NSError, emit the existential erasure from that
       // NSError.
@@ -726,9 +733,7 @@ ManagedValue SILGenFunction::emitExistentialErasure(
 
         // Receive the error value.  It's typed as an 'AnyObject' for
         // layering reasons, so perform an unchecked cast down to NSError.
-        SILType anyObjectTy =
-            potentialNSError.getType().getOptionalObjectType();
-        ManagedValue nsError = B.createOwnedPhiArgument(anyObjectTy);
+        auto nsError = B.createOptionalSomeResult(switchEnum);
         nsError = B.createUncheckedRefCast(loc, nsError, 
                                            getLoweredType(nsErrorType));
 

lib/SILGen/SILGenDistributed.cpp

@@ -622,19 +622,18 @@ void SILGenFunction::emitDistributedActorFactory(FuncDecl *fd) {
     auto resolve =
         switchBB->createPhiArgument(optionalReturnTy, OwnershipKind::Owned);
 
-    B.createSwitchEnum(
+    auto *switchEnum = B.createSwitchEnum(
         loc, resolve, nullptr,
         {{C.getOptionalSomeDecl(), resolvedBB},
          {std::make_pair(C.getOptionalNoneDecl(), makeProxyBB)}});
+    switchEnum->createOptionalSomeResult();
   }
 
   // ==== Case 'some') return the resolved instance
   {
     B.emitBlock(resolvedBB);
 
-    auto local = resolvedBB->createPhiArgument(returnTy, OwnershipKind::Owned);
-
-    B.createBranch(loc, returnBB, {local});
+    B.createBranch(loc, returnBB, {resolvedBB->getArgument(0)});
   }
 
   // ==== Case 'none') Create the remote instance
@@ -1016,4 +1015,4 @@ void SILGenFunction::emitDistributedThunk(SILDeclRef thunk) {
     Cleanups.emitCleanupsForReturn(CleanupLocation(loc), IsForUnwind);
     B.createThrow(loc, error);
   }
-}
+}