Merge pull request #71406 from jckarter/borrowing-switch-5

SILGen: Emit borrowing switch subjects under a formal access.

Author: jckarter

lib/SILGen/ArgumentSource.h

@@ -222,12 +222,6 @@ class ArgumentSource {
     return Storage.get<LValueStorage>(StoredKind).Loc;
   }
 
-  /// The kind of operation under which we are querying a storage reference.
-  enum class StorageReferenceOperationKind {
-    Borrow,
-    Consume
-  };
-
   Expr *findStorageReferenceExprForBorrow() &&;
   Expr *findStorageReferenceExprForMoveOnly(SILGenFunction &SGF,
                                       StorageReferenceOperationKind refKind) &&;

lib/SILGen/SILGenApply.cpp

@@ -3196,13 +3196,8 @@ static StorageRefResult findStorageReferenceExprForBorrow(Expr *e) {
   return StorageRefResult();
 }
 
-Expr *ArgumentSource::findStorageReferenceExprForMoveOnly(
-    SILGenFunction &SGF, StorageReferenceOperationKind kind) && {
-  if (!isExpr())
-    return nullptr;
-
-  auto argExpr = asKnownExpr();
-
+Expr *SILGenFunction::findStorageReferenceExprForMoveOnly(Expr *argExpr,
+                                           StorageReferenceOperationKind kind) {
   // If there's a load around the outer part of this arg expr, look past it.
   bool sawLoad = false;
   if (auto *li = dyn_cast<LoadExpr>(argExpr)) {
@@ -3270,7 +3265,7 @@ Expr *ArgumentSource::findStorageReferenceExprForMoveOnly(
   assert(type);
 
   SILType ty =
-      SGF.getLoweredType(type->getWithoutSpecifierType()->getCanonicalType());
+      getLoweredType(type->getWithoutSpecifierType()->getCanonicalType());
   bool isMoveOnly = ty.isMoveOnly(/*orWrapped=*/false);
   if (auto *pd = dyn_cast<ParamDecl>(storage)) {
     isMoveOnly |= pd->getSpecifier() == ParamSpecifier::Borrowing;
@@ -3279,10 +3274,6 @@ Expr *ArgumentSource::findStorageReferenceExprForMoveOnly(
   if (!isMoveOnly)
     return nullptr;
 
-  // Claim the value of this argument since we found a storage reference that
-  // has a move only base.
-  (void)std::move(*this).asKnownExpr();
-
   // If we saw a subscript expr and the base of the subscript expr passed our
   // tests above, we can emit the call to the subscript directly as a borrowed
   // lvalue. Return the subscript expr here so that we emit it appropriately.
@@ -3292,11 +3283,22 @@ Expr *ArgumentSource::findStorageReferenceExprForMoveOnly(
   return result.getTransitiveRoot();
 }
 
-Expr *
-ArgumentSource::findStorageReferenceExprForBorrowExpr(SILGenFunction &SGF) && {
+Expr *ArgumentSource::findStorageReferenceExprForMoveOnly(
+    SILGenFunction &SGF, StorageReferenceOperationKind kind) && {
   if (!isExpr())
     return nullptr;
 
+  auto lvExpr = SGF.findStorageReferenceExprForMoveOnly(asKnownExpr(), kind);
+  if (lvExpr) {
+    // Claim the value of this argument since we found a storage reference that
+    // has a move only base.
+    (void)std::move(*this).asKnownExpr();
+  }
+  return lvExpr;
+}
+
+Expr *
+SILGenFunction::findStorageReferenceExprForBorrowExpr(Expr *argExpr) {
   // We support two patterns:
   //
   // (load_expr (borrow_expr))
@@ -3305,8 +3307,6 @@ ArgumentSource::findStorageReferenceExprForBorrowExpr(SILGenFunction &SGF) && {
   //
   // The first happens if a borrow is used on a non-self argument. The second
   // happens if we pass self as a borrow.
-  auto *argExpr = asKnownExpr();
-
   if (auto *parenExpr = dyn_cast<ParenExpr>(argExpr))
     argExpr = parenExpr->getSubExpr();
 
@@ -3317,14 +3317,21 @@ ArgumentSource::findStorageReferenceExprForBorrowExpr(SILGenFunction &SGF) && {
   if (!borrowExpr)
     return nullptr;
 
-  Expr *lvExpr = ::findStorageReferenceExprForBorrow(borrowExpr->getSubExpr())
+  return ::findStorageReferenceExprForBorrow(borrowExpr->getSubExpr())
                      .getTransitiveRoot();
+}
 
+Expr *
+ArgumentSource::findStorageReferenceExprForBorrowExpr(SILGenFunction &SGF) && {
+  if (!isExpr())
+    return nullptr;
+
+  auto lvExpr = SGF.findStorageReferenceExprForBorrowExpr(asKnownExpr());
   // Claim the value of this argument.
   if (lvExpr) {
     (void)std::move(*this).asKnownExpr();
   }
-
+  
   return lvExpr;
 }
 
@@ -3686,7 +3693,7 @@ class ArgEmitter {
 
     // Try to find an expression we can emit as a borrowed l-value.
     auto lvExpr = std::move(arg).findStorageReferenceExprForMoveOnly(
-        SGF, ArgumentSource::StorageReferenceOperationKind::Borrow);
+        SGF, StorageReferenceOperationKind::Borrow);
     if (!lvExpr)
       return false;
 
@@ -3760,7 +3767,7 @@ class ArgEmitter {
 
     // Try to find an expression we can emit as a consumed l-value.
     auto lvExpr = std::move(arg).findStorageReferenceExprForMoveOnly(
-        SGF, ArgumentSource::StorageReferenceOperationKind::Consume);
+                                  SGF, StorageReferenceOperationKind::Consume);
     if (!lvExpr)
       return false;
 
@@ -3791,7 +3798,7 @@ class ArgEmitter {
   void
   emitDirect(ArgumentSource &&arg, SILType loweredSubstArgType,
              AbstractionPattern origParamType, SILParameterInfo param,
-             llvm::Optional<AnyFunctionType::Param> origParam = llvm::None) {
+             llvm::Optional<AnyFunctionType::Param> origParam = llvm::None) {             
     ManagedValue value;
     auto loc = arg.getLocation();
 

lib/SILGen/SILGenFunction.h

@@ -272,6 +272,12 @@ struct MaterializedLValue {
       callbackStorage(callbackStorage) {}
 };
 
+/// The kind of operation under which we are querying a storage reference.
+enum class StorageReferenceOperationKind {
+  Borrow,
+  Consume
+};
+
 /// SILGenFunction - an ASTVisitor for producing SIL from function bodies.
 class LLVM_LIBRARY_VISIBILITY SILGenFunction
   : public ASTVisitor<SILGenFunction>
@@ -1518,6 +1524,12 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   /// Emit the given expression as an r-value.
   RValue emitRValue(Expr *E, SGFContext C = SGFContext());
 
+  /// Given an expression, find the subexpression that can be emitted as a borrow formal access, if
+  /// any.
+  Expr *findStorageReferenceExprForMoveOnly(Expr *argExpr,
+                                            StorageReferenceOperationKind kind);
+  Expr *findStorageReferenceExprForBorrowExpr(Expr *argExpr);
+
   /// Emit the given expression as a +1 r-value.
   ///
   /// *NOTE* This creates the +1 r-value and then pushes that +1 r-value through

lib/SILGen/SILGenPattern.cpp

@@ -3178,6 +3178,9 @@ void SILGenFunction::emitSwitchStmt(SwitchStmt *S) {
              llvm::dbgs() << '\n');
 
   auto subjectTy = S->getSubjectExpr()->getType();
+  auto subjectLoweredTy = getLoweredType(subjectTy);
+  auto subjectLoweredAddress =
+    silConv.useLoweredAddresses() && subjectLoweredTy.isAddressOnly(F);
 
   // If the subject expression is uninhabited, we're already dead.
   // Emit an unreachable in place of the switch statement.
@@ -3254,27 +3257,88 @@ void SILGenFunction::emitSwitchStmt(SwitchStmt *S) {
   emission.emitAddressOnlyAllocations();
 
   SILBasicBlock *contBB = createBasicBlock();
-  JumpDest contDest(contBB, Cleanups.getCleanupsDepth(), CleanupLocation(S));
-
+  // Depending on the switch ownership behavior, we want to either:
+  //
+  // - allow the subject to (potentially or explicitly) be forwarded into the
+  //   case blocks. In this case we want to include the cleanups for the subject
+  //   in the case blocks so that the subject and its components can be
+  //   consumed by each case block.
+  //
+  // or:
+  //
+  // - evaluate the subject under a formal access scope (a borrow or inout).
+  //   In this case the lifetime of the access should cover all the way to the
+  //   exits out of the switch.
+  //
+  // When we break out of a case block, we take the subject's remnants with us
+  // in the former case, but not the latter.q
+  CleanupsDepth subjectDepth = Cleanups.getCleanupsDepth();
   LexicalScope switchScope(*this, CleanupLocation(S));
+  llvm::Optional<FormalEvaluationScope> switchFormalAccess;
+  
+  bool subjectUndergoesFormalAccess;
+  switch (ownership) {
+  // For normal copyable subjects, we allow the value to be forwarded into
+  // the cases, since we can copy it as needed to evaluate the pattern match.
+  case ValueOwnership::Default:
+  // Similarly, if the subject is an explicitly consumed noncopyable value,
+  // we can forward ownership of the subject's parts into matching case blocks.
+  case ValueOwnership::Owned:
+    subjectUndergoesFormalAccess = false;
+    break;
 
-  // Enter a break/continue scope.  If we wanted a continue
-  // destination, it would probably be out here.
-  BreakContinueDestStack.push_back({S, contDest, JumpDest(S)});
+  // Borrowed and inout pattern matches both undergo a formal access.
+  case ValueOwnership::Shared:
+  case ValueOwnership::InOut:
+    subjectUndergoesFormalAccess = true;
+    break;
+  }
 
   PatternMatchContext switchContext = { emission };
   SwitchStack.push_back(&switchContext);
 
-  // Emit the subject value. If at +1, dispatching will consume it. If it is at
-  // +0, we just forward down borrows.
-  //
-  // TODO: For an inout switch, we'd start a formal access to an lvalue here.
-  ManagedValue subjectMV = emitRValueAsSingleValue(
-      S->getSubjectExpr(),
-      ownership <= ValueOwnership::Shared
-        ? SGFContext::AllowGuaranteedPlusZero
-        : SGFContext());
-
+  // Emit the subject value.
+  ManagedValue subjectMV;
+  switch (ownership) {
+  case ValueOwnership::Default: {
+    // A regular copyable pattern match. Emit as a regular rvalue.
+    // If at +1, dispatching will consume it. If it is at +0, we just forward
+    // down borrows.
+    subjectMV = emitRValueAsSingleValue(
+        S->getSubjectExpr(),
+        SGFContext::AllowGuaranteedPlusZero);
+    break;
+  }
+  case ValueOwnership::Shared: {
+    // A borrowing pattern match. See if we can emit the subject under a read
+    // formal access.
+    switchFormalAccess.emplace(*this);
+    auto subjectExpr = S->getSubjectExpr();
+    if (auto subjectLVExpr = findStorageReferenceExprForMoveOnly(subjectExpr,
+                                       StorageReferenceOperationKind::Borrow)) {
+      LValue sharedLV = emitLValue(subjectLVExpr,
+        subjectLoweredAddress ? SGFAccessKind::BorrowedAddressRead
+                              : SGFAccessKind::BorrowedObjectRead);
+      subjectMV = emitBorrowedLValue(S->getSubjectExpr(), std::move(sharedLV));
+    } else {
+      // Emit the value as an allowed-+0 rvalue if it doesn't have special
+      // lvalue treatment.
+      subjectMV = emitRValueAsSingleValue(S->getSubjectExpr(),
+                                          SGFContext::AllowGuaranteedPlusZero);
+    }
+    break;
+  }
+  case ValueOwnership::InOut: {
+    // A mutating pattern match. Emit the subject under a modify access.
+    llvm_unreachable("not yet implemented");
+  }
+  case ValueOwnership::Owned: {
+    // A consuming pattern match. Emit as a +1 rvalue.
+    subjectMV = emitRValueAsSingleValue(S->getSubjectExpr());
+    break;
+  }
+  }
+  
   // Inline constructor for subject.
   auto subject = ([&]() -> ConsumableManagedValue {
     if (subjectMV.getType().isMoveOnly()) {
@@ -3292,12 +3356,15 @@ void SILGenFunction::emitSwitchStmt(SwitchStmt *S) {
 
         case ValueOwnership::Shared:
           emission.setNoncopyableBorrowingOwnership();
+          if (!subjectMV.isPlusZero()) {
+            subjectMV = subjectMV.borrow(*this, S);
+          }
           if (subjectMV.getType().isAddress() &&
               subjectMV.getType().isLoadable(F)) {
             // Load a borrow if the type is loadable.
-            subjectMV = B.createLoadBorrow(S, subjectMV);
-          } else if (!subjectMV.isPlusZero()) {
-            subjectMV = subjectMV.borrow(*this, S);
+            subjectMV = subjectUndergoesFormalAccess
+              ? B.createFormalAccessLoadBorrow(S, subjectMV)
+              : B.createLoadBorrow(S, subjectMV);
           }
           return {subjectMV, CastConsumptionKind::BorrowAlways};
 
@@ -3385,6 +3452,20 @@ void SILGenFunction::emitSwitchStmt(SwitchStmt *S) {
     return {subjectMV.copy(*this, S), CastConsumptionKind::TakeAlways};
   }());
 
+  CleanupsDepth caseBodyDepth = Cleanups.getCleanupsDepth();
+  llvm::Optional<LexicalScope> caseBodyScope;
+  if (subjectUndergoesFormalAccess) {
+    caseBodyScope.emplace(*this, CleanupLocation(S));
+  }
+
+  // Enter a break/continue scope. As discussed above, the depth we jump to
+  // depends on whether the subject is under a formal access.
+  JumpDest contDest(contBB,
+                    subjectUndergoesFormalAccess ? caseBodyDepth
+                                                 : subjectDepth,
+                    CleanupLocation(S));
+  BreakContinueDestStack.push_back({S, contDest, JumpDest(S)});
+
   // If we need to diagnose an unexpected enum case or unexpected enum case
   // value, we need access to a value metatype for the subject. Emit this state
   // now before we emit the actual switch to ensure that the subject has not
@@ -3446,7 +3527,14 @@ void SILGenFunction::emitSwitchStmt(SwitchStmt *S) {
   }
 
   assert(!B.hasValidInsertionPoint());
-  switchScope.pop();
+  // Disable the cleanups for values that should be consumed by the case
+  // bodies.
+  caseBodyScope.reset();
+  // If the subject isn't under a formal access, that includes the subject
+  // itself.
+  if (!subjectUndergoesFormalAccess) {
+    switchScope.pop();
+  }
 
   // Then emit the case blocks shared by multiple pattern cases.
   emission.emitSharedCaseBlocks(
@@ -3463,6 +3551,12 @@ void SILGenFunction::emitSwitchStmt(SwitchStmt *S) {
   } else {
     B.emitBlock(contBB);
   }
+  
+  // End the formal access to the subject now (if there was one).
+  if (subjectUndergoesFormalAccess) {
+    switchFormalAccess.reset();
+    switchScope.pop();
+  }
 
   // Now that we have emitted everything, see if our unexpected enum case info
   // metatypes were actually used. If not, delete them.