Merge pull request #71726 from jckarter/borrowing-switch-9

MoveOnlyAddressChecker: Relax more checks to accommodate borrowing switch codegen.

Author: jckarter

lib/SIL/Utils/FieldSensitivePrunedLiveness.cpp

@@ -114,7 +114,13 @@ SubElementOffset::computeForAddress(SILValue projectionDerivedFromRoot,
   unsigned finalSubElementOffset = 0;
   SILModule &mod = *rootAddress->getModule();
 
+  LLVM_DEBUG(llvm::dbgs() << "computing element offset for root:\n";
+             rootAddress->print(llvm::dbgs()));
+
   while (1) {
+    LLVM_DEBUG(llvm::dbgs() << "projection: ";
+               projectionDerivedFromRoot->print(llvm::dbgs()));
+
     // If we got to the root, we're done.
     if (rootAddress == projectionDerivedFromRoot)
       return {SubElementOffset(finalSubElementOffset)};
@@ -195,6 +201,16 @@ SubElementOffset::computeForAddress(SILValue projectionDerivedFromRoot,
       projectionDerivedFromRoot = initData->getOperand();
       continue;
     }
+    
+    // Look through wrappers.
+    if (auto c2m = dyn_cast<CopyableToMoveOnlyWrapperAddrInst>(projectionDerivedFromRoot)) {
+      projectionDerivedFromRoot = c2m->getOperand();
+      continue;
+    }
+    if (auto m2c = dyn_cast<MoveOnlyWrapperToCopyableValueInst>(projectionDerivedFromRoot)) {
+      projectionDerivedFromRoot = m2c->getOperand();
+      continue;
+    }
 
     // If we do not know how to handle this case, just return None.
     //

lib/SILGen/SILGenPattern.cpp

@@ -1398,6 +1398,28 @@ void PatternMatchEmission::bindVariable(Pattern *pattern, VarDecl *var,
   }
 }
 
+namespace {
+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());
+    }
+  }
+};
+}
+
 /// Bind a borrow binding into the current scope.
 void PatternMatchEmission::bindBorrow(Pattern *pattern, VarDecl *var,
                                       ConsumableManagedValue value) {
@@ -1419,6 +1441,14 @@ void PatternMatchEmission::bindBorrow(Pattern *pattern, VarDecl *var,
   if (bindValue.getType().isObject()) {
     // Create a notional copy for the borrow checker to use.
     bindValue = bindValue.copy(SGF, pattern);
+  } else {
+    // Treat use of an address value in-place as a separate nested read access.
+    auto access = SGF.B.createBeginAccess(pattern, bindValue.getValue(),
+                                          SILAccessKind::Read,
+                                          SILAccessEnforcement::Static,
+                                          false, false);
+    SGF.Cleanups.pushCleanup<EndAccessCleanup>(access);
+    bindValue = ManagedValue::forBorrowedAddressRValue(access);
   }
   // 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
@@ -1789,16 +1819,13 @@ emitTupleDispatch(ArrayRef<RowToSpecialize> rows, ConsumableManagedValue src,
         return getManagedSubobject(SGF, member, fieldTL,
                                    src.getFinalConsumption());
       }
-      case CastConsumptionKind::CopyOnSuccess: {
+      case CastConsumptionKind::CopyOnSuccess:
+      case CastConsumptionKind::BorrowAlways: {
         // We translate copy_on_success => borrow_always.
         auto memberMV = ManagedValue::forBorrowedAddressRValue(member);
         return {SGF.B.createLoadBorrow(loc, memberMV),
                 CastConsumptionKind::BorrowAlways};
       }
-      case CastConsumptionKind::BorrowAlways: {
-        llvm_unreachable(
-            "Borrow always can only occur along object only code paths");
-      }
       }
       llvm_unreachable("covered switch");
     }());
@@ -3202,26 +3229,6 @@ 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());

lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp

@@ -1128,14 +1128,21 @@ void UseState::initializeLiveness(
     liveness.initializeDef(SILValue(address), liveness.getTopLevelSpan());
   }
 
-  // Assume a strict check of a temporary is initialized before the check.
-  if (auto *asi = dyn_cast<BeginAccessInst>(address->getOperand());
-      asi && address->isStrict()) {
+  // Assume a strict-checked value initialized before the check.
+  if (address->isStrict()) {
     LLVM_DEBUG(llvm::dbgs()
-               << "Adding strict-marked begin_access as init!\n");
+               << "Adding strict marker as init!\n");
     recordInitUse(address, address, liveness.getTopLevelSpan());
     liveness.initializeDef(SILValue(address), liveness.getTopLevelSpan());
   }
+  
+  // Assume a value wrapped in a MoveOnlyWrapper is initialized.
+  if (auto *m2c = dyn_cast<CopyableToMoveOnlyWrapperAddrInst>(address->getOperand())) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "Adding copyable_to_move_only_wrapper as init!\n");
+    recordInitUse(address, address, liveness.getTopLevelSpan());
+    liveness.initializeDef(SILValue(address), liveness.getTopLevelSpan());    
+  }
 
   // Now that we have finished initialization of defs, change our multi-maps
   // from their array form to their map form.
@@ -2570,9 +2577,14 @@ bool GatherUsesVisitor::visitUse(Operand *op) {
 
 #ifndef NDEBUG
   if (user->mayWriteToMemory()) {
-    llvm::errs() << "Found a write classified as a liveness use?!\n";
-    llvm::errs() << "Use: " << *user;
-    llvm_unreachable("standard failure");
+    // TODO: `unchecked_take_enum_addr` should inherently be understood as
+    // non-side-effecting when it's nondestructive.
+    auto ue = dyn_cast<UncheckedTakeEnumDataAddrInst>(user);
+    if (!ue || ue->isDestructive()) {
+      llvm::errs() << "Found a write classified as a liveness use?!\n";
+      llvm::errs() << "Use: " << *user;
+      llvm_unreachable("standard failure");
+    }
   }
 #endif
   useState.recordLivenessUse(user, *leafRange);

test/SILOptimizer/moveonly_borrowing_switch.swift

@@ -236,9 +236,119 @@ func testOuterAO(consuming bas: consuming AOBas) { // expected-error{{'bas' used
         break
     }
 
-    switch bas { // expected-note{{used here}}
-    case _borrowing x: // expected-warning{{}}
+    switch bas {
+    case _borrowing x: // expected-warning{{}} expected-note{{used here}}
         break
     }
 }
 
+enum E<T>: ~Copyable {
+    case a(T)
+}
+
+extension E {
+    func f() {
+        switch self {
+        case .a:
+            print("a")
+        }
+    }
+
+    func g() {
+        switch self {
+        case .a(_borrowing t): // expected-warning{{}}
+            print("a")
+        }
+    }
+}
+
+struct Box: ~Copyable {
+    let ptr: UnsafeMutablePointer<Int>
+}
+
+struct ChonkyBox: ~Copyable {
+    let container: Any
+}
+
+enum List<Element>: ~Copyable {
+    case end
+    case more(Element, Box)
+}
+
+enum ChonkyList<Element>: ~Copyable {
+    case end
+    case more(Element, ChonkyBox)
+}
+
+extension List {
+    var isEmpty: Bool {
+        switch self {
+        case .end: true
+        case .more: false
+        }
+    }
+
+    var peek: Box {
+        _read {
+            switch self {
+            case .end:
+                fatalError()
+            case .more(_, _borrowing box):
+                yield box
+            }
+        }
+    }
+
+/*
+    TODO: type mismatch because of `@moveOnly` wrapper. yield needs to peel it
+    off
+
+    var head: Element {
+        _read {
+            switch self {
+            case .end:
+                fatalError()
+            case .more(_borrowing head, _):
+                yield head
+            }
+        }
+    }
+*/
+}
+
+extension ChonkyList {
+    var isEmpty: Bool {
+        switch self {
+        case .end: true
+        case .more: false
+        }
+    }
+
+    var peek: ChonkyBox {
+        _read {
+            switch self {
+            case .end:
+                fatalError()
+            case .more(_, _borrowing box):
+                yield box
+            }
+        }
+    }
+
+/*
+    TODO: type mismatch because of `@moveOnly` wrapper. yield needs to peel it
+    off
+
+    var head: Element {
+        _read {
+            switch self {
+            case .end:
+                fatalError()
+            case .more(_borrowing head, _):
+                yield head
+            }
+        }
+    }
+*/
+}
+