Merge pull request #66094 from gottesmm/pr-9cc91f1b6b3ff60686adb825c4a444221c3d798a

[move-only] Teach SILGenApply how to emit subscripts with borrowed base values.

Author: gottesmm

lib/SILGen/ArgumentSource.h

@@ -280,7 +280,7 @@ class ArgumentSource {
 
   ArgumentSource copyForDiagnostics() const;
 
-  void dump() const;
+  LLVM_DUMP_METHOD void dump() const;
   void dump(raw_ostream &os, unsigned indent = 0) const;
 
 private:

lib/SILGen/SILGenApply.cpp

@@ -3029,7 +3029,8 @@ static StorageRefResult findStorageReferenceExprForBorrow(Expr *e) {
       return result.withTransitiveRoot(te);
 
   } else if (auto ioe = dyn_cast<InOutExpr>(e)) {
-    return ioe;
+    if (auto result = findStorageReferenceExprForBorrow(ioe->getSubExpr()))
+      return result.withTransitiveRoot(ioe);
   }
 
   return StorageRefResult();
@@ -3049,6 +3050,24 @@ Expr *ArgumentSource::findStorageReferenceExprForMoveOnly(
     sawLoad = true;
   }
 
+  // If we have a subscript, strip it off and make sure that our base is
+  // something that we can process. If we do and we succeed below, we return the
+  // subscript instead.
+  SubscriptExpr *subscriptExpr = nullptr;
+  if ((subscriptExpr = dyn_cast<SubscriptExpr>(argExpr))) {
+    auto *decl = cast<SubscriptDecl>(subscriptExpr->getDecl().getDecl());
+    if (decl->getReadImpl() != ReadImplKind::Read) {
+      subscriptExpr = nullptr;
+    } else {
+      argExpr = subscriptExpr->getBase();
+    }
+
+    // If there's a load on the base of the subscript expr, look past it.
+    if (auto *li = dyn_cast<LoadExpr>(argExpr)) {
+      argExpr = li->getSubExpr();
+    }
+  }
+
   // If we're consuming instead, then the load _must_ have been there.
   if (kind == StorageReferenceOperationKind::Consume && !sawLoad)
     return nullptr;
@@ -3097,6 +3116,12 @@ Expr *ArgumentSource::findStorageReferenceExprForMoveOnly(
   // 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.
+  if (subscriptExpr)
+    return subscriptExpr;
+
   return result.getTransitiveRoot();
 }
 

lib/SILGen/SILGenDecl.cpp

@@ -1964,13 +1964,19 @@ InitializationPtr SILGenFunction::emitLocalVariableWithCleanup(
 std::unique_ptr<TemporaryInitialization>
 SILGenFunction::emitTemporary(SILLocation loc, const TypeLowering &tempTL) {
   SILValue addr = emitTemporaryAllocation(loc, tempTL.getLoweredType());
+  if (addr->getType().isMoveOnly())
+    addr = B.createMarkMustCheckInst(
+        loc, addr, MarkMustCheckInst::CheckKind::ConsumableAndAssignable);
   return useBufferAsTemporary(addr, tempTL);
 }
 
 std::unique_ptr<TemporaryInitialization>
 SILGenFunction::emitFormalAccessTemporary(SILLocation loc,
                                           const TypeLowering &tempTL) {
   SILValue addr = emitTemporaryAllocation(loc, tempTL.getLoweredType());
+  if (addr->getType().isMoveOnly())
+    addr = B.createMarkMustCheckInst(
+        loc, addr, MarkMustCheckInst::CheckKind::ConsumableAndAssignable);
   CleanupHandle cleanup =
       enterDormantFormalAccessTemporaryCleanup(addr, loc, tempTL);
   return std::unique_ptr<TemporaryInitialization>(

lib/SILGen/SILGenLValue.cpp

@@ -366,6 +366,7 @@ LogicalPathComponent::projectForRead(SILGenFunction &SGF, SILLocation loc,
   }
 
   TemporaryInitializationPtr tempInit;
+  ManagedValue result;
   RValue rvalue;
 
   // If the RValue type has an openedExistential, then the RValue must be
@@ -377,9 +378,11 @@ LogicalPathComponent::projectForRead(SILGenFunction &SGF, SILLocation loc,
 
     // Create a temporary, whose type may depend on the 'get'.
     tempInit = SGF.emitFormalAccessTemporary(loc, RValueTL);
+    result = tempInit->getManagedAddress();
   } else {
     // Create a temporary for a static (non-dependent) RValue type.
     tempInit = SGF.emitFormalAccessTemporary(loc, RValueTL);
+    result = tempInit->getManagedAddress();
 
     // Emit a 'get' directly into the temporary.
     rvalue = std::move(*this).get(SGF, loc, base, SGFContext(tempInit.get()));
@@ -390,7 +393,8 @@ LogicalPathComponent::projectForRead(SILGenFunction &SGF, SILLocation loc,
   if (!rvalue.isInContext())
     std::move(rvalue).forwardInto(SGF, loc, tempInit.get());
 
-  return tempInit->getManagedAddress();
+  assert(result);
+  return result;
 }
 
 ManagedValue LogicalPathComponent::project(SILGenFunction &SGF,

lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp

@@ -2523,29 +2523,27 @@ void MoveOnlyAddressCheckerPImpl::rewriteUses(
   // scope, we would have emitted the scope expansion error during diagnostics.
   for (auto pair : addressUseState.borrows) {
     if (auto *li = dyn_cast<LoadInst>(pair.first)) {
-      if (li->getOwnershipQualifier() == LoadOwnershipQualifier::Copy) {
-        // If we had a load [copy], borrow then we know that all of its destroys
-        // must have been destroy_value. So we can just gather up those
-        // destroy_value and use then to create a new load_borrow scope.
-        SILBuilderWithScope builder(li);
-        auto *lbi = builder.createLoadBorrow(li->getLoc(), li->getOperand());
-        // We use this auxillary list to avoid iterator invalidation of
-        // li->getConsumingUse();
-        StackList<DestroyValueInst *> toDelete(lbi->getFunction());
-        for (auto *consumeUse : li->getConsumingUses()) {
-          auto *dvi = cast<DestroyValueInst>(consumeUse->getUser());
-          SILBuilderWithScope destroyBuilder(dvi);
-          destroyBuilder.createEndBorrow(dvi->getLoc(), lbi);
-          toDelete.push_back(dvi);
-          changed = true;
-        }
-        while (!toDelete.empty())
-          toDelete.pop_back_val()->eraseFromParent();
-
-        li->replaceAllUsesWith(lbi);
-        li->eraseFromParent();
-        continue;
+      // If we had a load -> load_borrow then we know that all of its destroys
+      // must have been destroy_value. So we can just gather up those
+      // destroy_value and use then to create a new load_borrow scope.
+      SILBuilderWithScope builder(li);
+      auto *lbi = builder.createLoadBorrow(li->getLoc(), li->getOperand());
+      // We use this auxillary list to avoid iterator invalidation of
+      // li->getConsumingUse();
+      StackList<DestroyValueInst *> toDelete(lbi->getFunction());
+      for (auto *consumeUse : li->getConsumingUses()) {
+        auto *dvi = cast<DestroyValueInst>(consumeUse->getUser());
+        SILBuilderWithScope destroyBuilder(dvi);
+        destroyBuilder.createEndBorrow(dvi->getLoc(), lbi);
+        toDelete.push_back(dvi);
+        changed = true;
       }
+      while (!toDelete.empty())
+        toDelete.pop_back_val()->eraseFromParent();
+
+      li->replaceAllUsesWith(lbi);
+      li->eraseFromParent();
+      continue;
     }
 
     llvm::dbgs() << "Borrow: " << *pair.first;