Merge pull request #71489 from jckarter/borrowing-switch-8

MoveOnlyAddressUtils: Fixes for borrowing switch over address-only types.

Author: jckarter

include/swift/SIL/SILBuilder.h

@@ -1444,9 +1444,11 @@ class SILBuilder {
 
   MarkUnresolvedNonCopyableValueInst *createMarkUnresolvedNonCopyableValueInst(
       SILLocation loc, SILValue src,
-      MarkUnresolvedNonCopyableValueInst::CheckKind kind) {
+      MarkUnresolvedNonCopyableValueInst::CheckKind kind,
+      MarkUnresolvedNonCopyableValueInst::IsStrict_t strict
+        = MarkUnresolvedNonCopyableValueInst::IsNotStrict) {
     return insert(new (getModule()) MarkUnresolvedNonCopyableValueInst(
-        getSILDebugLocation(loc), src, kind));
+        getSILDebugLocation(loc), src, kind, strict));
   }
 
   MarkUnresolvedReferenceBindingInst *createMarkUnresolvedReferenceBindingInst(

include/swift/SIL/SILInstruction.h

@@ -8716,15 +8716,31 @@ class MarkUnresolvedNonCopyableValueInst
     /// like class initializers.
     InitableButNotConsumable,
   };
+  
+  /// During SILGen, we have not yet done escape analysis on local variables,
+  /// so we conservatively emit them as boxed and let the AllocBoxToStack
+  /// pass promote unescaped local variables. As part of this promotion,
+  /// non-strict `NoConsumeOrAssign` accesses can be promoted to
+  /// `ConsumableAndAssignable` since the variable is locally owned
+  /// if it doesn't escape. "Strict" accesses on the other hand preserve
+  /// their stricter access constraints. This is useful for representing things
+  /// like `borrow` bindings.
+  enum IsStrict_t : bool {
+    IsNotStrict = false,
+    IsStrict = true,
+  };
 
 private:
   CheckKind kind;
+  IsStrict_t strict;
 
   MarkUnresolvedNonCopyableValueInst(SILDebugLocation DebugLoc,
-                                     SILValue operand, CheckKind checkKind)
+                                     SILValue operand, CheckKind checkKind,
+                                     IsStrict_t strict = IsNotStrict)
       : UnaryInstructionBase(DebugLoc, operand, operand->getType(),
                              operand->getOwnershipKind()),
-        kind(checkKind) {
+        kind(checkKind),
+        strict(strict) {
     assert(operand->getType().isMoveOnly() &&
            "mark_unresolved_non_copyable_value can only take a move only typed "
            "value");
@@ -8746,6 +8762,10 @@ class MarkUnresolvedNonCopyableValueInst
       return true;
     }
   }
+  
+  IsStrict_t isStrict() const {
+    return strict;
+  }
 };
 
 /// A marker instruction that states a given alloc_box or alloc_stack is a

lib/SIL/IR/SILPrinter.cpp

@@ -2114,6 +2114,9 @@ class SILPrinter : public SILInstructionVisitor<SILPrinter> {
   void visitMarkUnresolvedNonCopyableValueInst(
       MarkUnresolvedNonCopyableValueInst *I) {
     using CheckKind = MarkUnresolvedNonCopyableValueInst::CheckKind;
+    if (I->isStrict()) {
+      *this << "[strict] ";
+    }
     switch (I->getCheckKind()) {
     case CheckKind::Invalid:
       llvm_unreachable("Invalid?!");

lib/SIL/Parser/ParseSIL.cpp

@@ -3511,6 +3511,16 @@ bool SILParser::parseSpecificSILInstruction(SILBuilder &B,
       P.diagnose(InstLoc.getSourceLoc(), diag);
       return true;
     }
+    
+    auto Strict = MarkUnresolvedNonCopyableValueInst::IsNotStrict;
+    if (AttrName.equals("strict")) {
+      Strict = MarkUnresolvedNonCopyableValueInst::IsStrict;
+      if (!parseSILOptional(AttrName, *this)) {
+        auto diag = diag::sil_markmustcheck_requires_attribute;
+        P.diagnose(InstLoc.getSourceLoc(), diag);
+        return true;
+      }
+    }
 
     using CheckKind = MarkUnresolvedNonCopyableValueInst::CheckKind;
     CheckKind CKind =
@@ -3535,7 +3545,8 @@ bool SILParser::parseSpecificSILInstruction(SILBuilder &B,
     if (parseSILDebugLocation(InstLoc, B))
       return true;
 
-    auto *MVI = B.createMarkUnresolvedNonCopyableValueInst(InstLoc, Val, CKind);
+    auto *MVI = B.createMarkUnresolvedNonCopyableValueInst(InstLoc, Val, CKind,
+                                                           Strict);
     ResultVal = MVI;
     break;
   }

lib/SIL/Utils/FieldSensitivePrunedLiveness.cpp

@@ -202,6 +202,9 @@ SubElementOffset::computeForAddress(SILValue projectionDerivedFromRoot,
     // really do not want to abort. Instead, our caller can choose to abort if
     // they get back a None. This ensures that we do not abort in cases where we
     // just want to emit to the user a "I do not understand" error.
+    LLVM_DEBUG(llvm::dbgs() << "unhandled projection derived from root:\n";
+               projectionDerivedFromRoot->print(llvm::dbgs()));
+
     return llvm::None;
   }
 }

lib/SILGen/SILGenBuilder.cpp

@@ -1090,13 +1090,14 @@ ManagedValue SILGenBuilder::createGuaranteedCopyableToMoveOnlyWrapperValue(
 
 ManagedValue SILGenBuilder::createMarkUnresolvedNonCopyableValueInst(
     SILLocation loc, ManagedValue value,
-    MarkUnresolvedNonCopyableValueInst::CheckKind kind) {
+    MarkUnresolvedNonCopyableValueInst::CheckKind kind,
+    MarkUnresolvedNonCopyableValueInst::IsStrict_t strict) {
   assert((value.isPlusOne(SGF) || value.isLValue() ||
           value.getType().isAddress()) &&
          "Argument must be at +1 or be an address!");
   CleanupCloner cloner(*this, value);
   auto *mdi = SILBuilder::createMarkUnresolvedNonCopyableValueInst(
-      loc, value.forward(getSILGenFunction()), kind);
+      loc, value.forward(getSILGenFunction()), kind, strict);
   return cloner.clone(mdi);
 }
 

lib/SILGen/SILGenBuilder.h

@@ -493,7 +493,9 @@ class SILGenBuilder : public SILBuilder {
   using SILBuilder::createMarkUnresolvedNonCopyableValueInst;
   ManagedValue createMarkUnresolvedNonCopyableValueInst(
       SILLocation loc, ManagedValue value,
-      MarkUnresolvedNonCopyableValueInst::CheckKind kind);
+      MarkUnresolvedNonCopyableValueInst::CheckKind kind,
+      MarkUnresolvedNonCopyableValueInst::IsStrict_t strict
+        = MarkUnresolvedNonCopyableValueInst::IsNotStrict);
 
   using SILBuilder::emitCopyValueOperation;
   ManagedValue emitCopyValueOperation(SILLocation Loc, ManagedValue v);

lib/SILGen/SILGenPattern.cpp

@@ -1420,8 +1420,12 @@ void PatternMatchEmission::bindBorrow(Pattern *pattern, VarDecl *var,
     // Create a notional copy for the borrow checker to use.
     bindValue = bindValue.copy(SGF, pattern);
   }
+  // We mark the borrow check as "strict" because we don't want to allow
+  // consumes through the binding, even if the original value manages to be
+  // stack promoted during AllocBoxToStack or anything like that.
   bindValue = SGF.B.createMarkUnresolvedNonCopyableValueInst(pattern, bindValue,
-              MarkUnresolvedNonCopyableValueInst::CheckKind::NoConsumeOrAssign);
+              MarkUnresolvedNonCopyableValueInst::CheckKind::NoConsumeOrAssign,
+              MarkUnresolvedNonCopyableValueInst::IsStrict);
 
   SGF.VarLocs[var] = SILGenFunction::VarLoc::get(bindValue.getValue());
 }
@@ -3198,6 +3202,26 @@ static void switchCaseStmtSuccessCallback(SILGenFunction &SGF,
   SGF.Cleanups.emitBranchAndCleanups(sharedDest, caseBlock, args);
 }
 
+class EndAccessCleanup final : public Cleanup {
+  SILValue beginAccess;
+public:
+  EndAccessCleanup(SILValue beginAccess)
+    : beginAccess(beginAccess)
+  {}
+  
+  void emit(SILGenFunction &SGF, CleanupLocation loc, ForUnwind_t forUnwind)
+  override {
+    SGF.B.createEndAccess(loc, beginAccess, /*aborted*/ false);
+  }
+  
+  void dump(SILGenFunction &SGF) const override {
+    llvm::errs() << "EndAccessCleanup\n";
+    if (beginAccess) {
+      beginAccess->print(llvm::errs());
+    }
+  }
+};
+
 void SILGenFunction::emitSwitchStmt(SwitchStmt *S) {
   LLVM_DEBUG(llvm::dbgs() << "emitting switch stmt\n";
              S->dump(llvm::dbgs());
@@ -3385,12 +3409,22 @@ void SILGenFunction::emitSwitchStmt(SwitchStmt *S) {
           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 = subjectUndergoesFormalAccess
-              ? B.createFormalAccessLoadBorrow(S, subjectMV)
-              : B.createLoadBorrow(S, subjectMV);
+          if (subjectMV.getType().isAddress()) {
+            if (subjectMV.getType().isLoadable(F)) {
+              // Load a borrow if the type is loadable.
+              subjectMV = subjectUndergoesFormalAccess
+                ? B.createFormalAccessLoadBorrow(S, subjectMV)
+                : B.createLoadBorrow(S, subjectMV);
+            } else {
+              // Initiate a read access on the memory, to ensure that even
+              // if the underlying memory is mutable or consumable, the pattern
+              // match is not allowed to modify it.
+              auto access = B.createBeginAccess(S, subjectMV.getValue(),
+                SILAccessKind::Read,
+                SILAccessEnforcement::Static, false, false);
+              Cleanups.pushCleanup<EndAccessCleanup>(access);
+              subjectMV = ManagedValue::forBorrowedAddressRValue(access);
+            }
           }
           return {subjectMV, CastConsumptionKind::BorrowAlways};