[silgen] Eliminate two more cases around subscripts where we were not borrowing.

Also, the store_borrow work in the previous patch caused some additional issues
to crop up. I fixed them in this PR and added some tests in the process.

Author: gottesmm

include/swift/SIL/SILInstruction.h

@@ -4392,8 +4392,18 @@ class StoreBorrowInst
 
   ArrayRef<Operand> getAllOperands() const { return Operands.asArray(); }
   MutableArrayRef<Operand> getAllOperands() { return Operands.asArray(); }
+
+  using EndBorrowRange =
+      decltype(std::declval<ValueBase>().getUsersOfType<EndBorrowInst>());
+
+  /// Return a range over all EndBorrow instructions for this BeginBorrow.
+  EndBorrowRange getEndBorrows() const;
 };
 
+inline auto StoreBorrowInst::getEndBorrows() const -> EndBorrowRange {
+  return getUsersOfType<EndBorrowInst>();
+}
+
 /// Represents the end of a borrow scope of a value %val from a
 /// value or address %src.
 ///

lib/SIL/Utils/FieldSensitivePrunedLiveness.cpp

@@ -129,6 +129,11 @@ SubElementOffset::computeForAddress(SILValue projectionDerivedFromRoot,
       continue;
     }
 
+    if (auto *sbi = dyn_cast<StoreBorrowInst>(projectionDerivedFromRoot)) {
+      projectionDerivedFromRoot = sbi->getDest();
+      continue;
+    }
+
     if (auto *m = dyn_cast<MoveOnlyWrapperToCopyableAddrInst>(
             projectionDerivedFromRoot)) {
       projectionDerivedFromRoot = m->getOperand();

lib/SILGen/SILGenApply.cpp

@@ -3124,6 +3124,21 @@ Expr *ArgumentSource::findStorageReferenceExprForMoveOnly(
   if (kind == StorageReferenceOperationKind::Consume && !sawLoad)
     return nullptr;
 
+  // If we did not see a load or a subscript expr and our argExpr is a
+  // declref_expr, return nullptr. We have an object not something that will be
+  // in memory. This can happen with classes or with values captured by a
+  // closure.
+  //
+  // NOTE: If we see a member_ref_expr from a decl_ref_expr, we still process it
+  // since the declref_expr could be from a class.
+  if (!sawLoad && !subscriptExpr) {
+    if (auto *declRef = dyn_cast<DeclRefExpr>(argExpr)) {
+      assert(!declRef->getType()->is<LValueType>() &&
+             "Shouldn't ever have an lvalue type here!");
+      return nullptr;
+    }
+  }
+
   auto result = ::findStorageReferenceExprForBorrow(argExpr);
 
   if (!result)
@@ -3143,31 +3158,18 @@ Expr *ArgumentSource::findStorageReferenceExprForMoveOnly(
   }
 
   if (!storage)
-      return nullptr;
+    return nullptr;
   assert(type);
 
   SILType ty =
       SGF.getLoweredType(type->getWithoutSpecifierType()->getCanonicalType());
   bool isMoveOnly = ty.isPureMoveOnly();
   if (auto *pd = dyn_cast<ParamDecl>(storage)) {
-      isMoveOnly |= pd->getSpecifier() == ParamSpecifier::Borrowing;
-      isMoveOnly |= pd->getSpecifier() == ParamSpecifier::Consuming;
+    isMoveOnly |= pd->getSpecifier() == ParamSpecifier::Borrowing;
+    isMoveOnly |= pd->getSpecifier() == ParamSpecifier::Consuming;
   }
   if (!isMoveOnly)
-      return nullptr;
-
-  // It makes sense to borrow any kind of storage we refer to at this stage,
-  // but SILGenLValue does not currently handle some kinds of references well.
-  //
-  // When rejecting to do the LValue-style borrow here, it'll end up going thru
-  // the RValue-style emission, after which the extra copy will get eliminated.
-  //
-  // If we did not see a LoadExpr around the argument expression, then only
-  // do the borrow if the storage is non-local.
-  // FIXME: I don't have a principled reason for why this matters and hope that
-  // we can fix the AST we're working with.
-  if (!sawLoad && storage->getDeclContext()->isLocalContext())
-      return nullptr;
+    return nullptr;
 
   // Claim the value of this argument since we found a storage reference that
   // has a move only base.

lib/SILGen/SILGenLValue.cpp

@@ -2918,6 +2918,8 @@ class LLVM_LIBRARY_VISIBILITY SILGenBorrowedBaseVisitor
       return m->getType()->isPureMoveOnly() ||
              m->getBase()->getType()->isPureMoveOnly();
     }
+    if (auto *d = dyn_cast<DeclRefExpr>(e))
+      return e->getType()->isPureMoveOnly();
     return false;
   }
 

lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp

@@ -280,6 +280,8 @@
 using namespace swift;
 using namespace swift::siloptimizer;
 
+#pragma clang optimize off
+
 llvm::cl::opt<bool> DisableMoveOnlyAddressCheckerLifetimeExtension(
     "move-only-address-checker-disable-lifetime-extension",
     llvm::cl::init(false),
@@ -342,7 +344,7 @@ static void convertMemoryReinitToInitForm(SILInstruction *memInst,
     break;
   }
   }
-  
+
   // Insert a new debug_value instruction after the reinitialization, so that
   // the debugger knows that the variable is in a usable form again.
   insertDebugValueBefore(memInst->getNextInstruction(), debugVar,
@@ -931,6 +933,7 @@ void UseState::initializeLiveness(
   // We begin by initializing all of our init uses.
   for (auto initInstAndValue : initInsts) {
     LLVM_DEBUG(llvm::dbgs() << "Found def: " << *initInstAndValue.first);
+
     liveness.initializeDef(initInstAndValue.first, initInstAndValue.second);
   }
 
@@ -944,7 +947,7 @@ void UseState::initializeLiveness(
                              reinitInstAndValue.second);
     }
   }
-  
+
   // Then check if our markedValue is from an argument that is in,
   // in_guaranteed, inout, or inout_aliasable, consider the marked address to be
   // the initialization point.
@@ -1050,6 +1053,26 @@ void UseState::initializeLiveness(
   LLVM_DEBUG(llvm::dbgs() << "Liveness with just inits:\n";
              liveness.print(llvm::dbgs()));
 
+  for (auto initInstAndValue : initInsts) {
+    // If our init inst is a store_borrow, treat the end_borrow as liveness
+    // uses.
+    //
+    // NOTE: We do not need to check for access scopes here since store_borrow
+    // can only apply to alloc_stack today.
+    if (auto *sbi = dyn_cast<StoreBorrowInst>(initInstAndValue.first)) {
+      // We can only store_borrow if our mark_must_check is a
+      // no_consume_or_assign.
+      assert(address->getCheckKind() ==
+                 MarkMustCheckInst::CheckKind::NoConsumeOrAssign &&
+             "store_borrow implies no_consume_or_assign since we cannot "
+             "consume a borrowed inited value");
+      for (auto *ebi : sbi->getEndBorrows()) {
+        liveness.updateForUse(ebi, initInstAndValue.second,
+                              false /*lifetime ending*/);
+      }
+    }
+  }
+
   // Now at this point, we have defined all of our defs so we can start adding
   // uses to the liveness.
   for (auto reinitInstAndValue : reinitInsts) {
@@ -1979,6 +2002,22 @@ bool GatherUsesVisitor::visitUse(Operand *op) {
   if (isa<EndAccessInst>(user))
     return true;
 
+  // This visitor looks through store_borrow instructions but does visit the
+  // end_borrow of the store_borrow. If we see such an end_borrow, register the
+  // store_borrow instead. Since we use sets, if we visit multiple end_borrows,
+  // we will only record the store_borrow once.
+  if (auto *ebi = dyn_cast<EndBorrowInst>(user)) {
+    if (auto *sbi = dyn_cast<StoreBorrowInst>(ebi->getOperand())) {
+      LLVM_DEBUG(llvm::dbgs() << "Found store_borrow: " << *sbi);
+      auto leafRange = TypeTreeLeafTypeRange::get(op->get(), getRootAddress());
+      if (!leafRange)
+        return false;
+
+      useState.recordInitUse(user, op->get(), *leafRange);
+      return true;
+    }
+  }
+
   if (auto *di = dyn_cast<DebugValueInst>(user)) {
     // Save the debug_value if it is attached directly to this mark_must_check.
     // If the underlying storage we're checking is immutable, then the access

lib/SILOptimizer/Mandatory/MoveOnlyUtils.cpp

@@ -220,6 +220,9 @@ bool noncopyable::memInstMustInitialize(Operand *memOper) {
   case SILInstructionKind::Store##Name##Inst:                                  \
     return cast<Store##Name##Inst>(memInst)->isInitializationOfDest();
 #include "swift/AST/ReferenceStorage.def"
+
+  case SILInstructionKind::StoreBorrowInst:
+    return true;
   }
 }
 

test/SILGen/moveonly.swift

@@ -818,13 +818,10 @@ func enumSwitchTest1(_ e: borrowing EnumSwitchTests.E) {
 //
 // CHECK: [[GLOBAL:%.*]] = global_addr @$s8moveonly9letGlobalAA16NonTrivialStructVvp :
 // CHECK: [[MARKED_GLOBAL:%.*]] = mark_must_check [no_consume_or_assign] [[GLOBAL]]
-// FIXME: this copy probably shouldn't be here when accessing through the letGlobal, but maybe it's cleaned up?
-// CHECK: [[LOADED_VAL:%.*]] = load [copy] [[MARKED_GLOBAL]] : $*NonTrivialStruct
-// CHECK: [[LOADED_BORROWED_VAL:%.*]] = begin_borrow [[LOADED_VAL]]
-// CHECK: [[LOADED_GEP:%.*]] = struct_extract [[LOADED_BORROWED_VAL]] : $NonTrivialStruct, #NonTrivialStruct.nonTrivialStruct2
+// CHECK: [[LOADED_VAL:%.*]] = load_borrow [[MARKED_GLOBAL]] : $*NonTrivialStruct
+// CHECK: [[LOADED_GEP:%.*]] = struct_extract [[LOADED_VAL]] : $NonTrivialStruct, #NonTrivialStruct.nonTrivialStruct2
 // CHECK: apply {{%.*}}([[LOADED_GEP]])
-// CHECK: end_borrow [[LOADED_BORROWED_VAL]]
-// CHECK: destroy_value [[LOADED_VAL]]
+// CHECK: end_borrow [[LOADED_VAL]]
 // CHECK: } // end sil function '$s8moveonly16testGlobalBorrowyyF'
 func testGlobalBorrow() {
     borrowVal(varGlobal)
@@ -856,13 +853,11 @@ func testGlobalBorrow() {
 //
 // CHECK: [[GLOBAL:%.*]] = global_addr @$s8moveonly9letGlobalAA16NonTrivialStructVvp :
 // CHECK: [[MARKED_GLOBAL:%.*]] = mark_must_check [no_consume_or_assign] [[GLOBAL]]
-// CHECK: [[LOADED_VAL:%.*]] = load [copy] [[MARKED_GLOBAL]]
-// CHECK: [[LOADED_BORROWED_VAL:%.*]] = begin_borrow [[LOADED_VAL]]
-// CHECK: [[LOADED_GEP:%.*]] = struct_extract [[LOADED_BORROWED_VAL]]
+// CHECK: [[LOADED_VAL:%.*]] = load_borrow [[MARKED_GLOBAL]]
+// CHECK: [[LOADED_GEP:%.*]] = struct_extract [[LOADED_VAL]]
 // CHECK: [[LOADED_GEP_COPY:%.*]] = copy_value [[LOADED_GEP]]
-// CHECK: end_borrow [[LOADED_BORROWED_VAL]]
-// CHECK: destroy_value [[LOADED_VAL]]
 // CHECK: apply {{%.*}}([[LOADED_GEP_COPY]])
+// CHECK: end_borrow [[LOADED_VAL]]
 //
 // CHECK: } // end sil function '$s8moveonly17testGlobalConsumeyyF'
 func testGlobalConsume() {
@@ -1833,11 +1828,8 @@ public func testSubscriptReadModify_BaseLoadable_ResultAddressOnly_Var() {
 // CHECK: [[MARK:%.*]] = mark_must_check [no_consume_or_assign] [[PROJECT]]
 // CHECK: [[LOAD_BORROW:%.*]] = load_borrow [[MARK]]
 // CHECK: ([[CORO_RESULT:%.*]], [[CORO_TOKEN:%.*]]) = begin_apply {{%.*}}({{%.*}}, [[LOAD_BORROW]])
-// CHECK: [[TEMP:%.*]] = alloc_stack $AddressOnlyProtocol
-// CHECK: copy_addr [[CORO_RESULT]] to [init] [[TEMP]]
+// CHECK: apply {{%.*}}([[CORO_RESULT]])
 // CHECK: end_apply [[CORO_TOKEN]]
-// CHECK: apply {{%.*}}([[TEMP]])
-// CHECK: destroy_addr [[TEMP]]
 // CHECK: end_borrow [[LOAD_BORROW]]
 // CHECK: } // end sil function '$s8moveonly58testSubscriptReadModify_BaseLoadable_ResultAddressOnly_LetyyF'
 public func testSubscriptReadModify_BaseLoadable_ResultAddressOnly_Let() {

test/SILOptimizer/moveonly_addresschecker.sil

@@ -844,3 +844,18 @@ bb0(%0 : $Int, %1a : $*NonCopyableNativeObjectPair):
   %16 = tuple ()
   return %16 : $()
 }
+
+sil [ossa] @testSupportStoreBorrow : $@convention(thin) (@guaranteed NonTrivialStruct) -> () {
+bb0(%0 : @guaranteed $NonTrivialStruct):
+  %1 = alloc_stack $NonTrivialStruct
+  %1a = mark_must_check [no_consume_or_assign] %1 : $*NonTrivialStruct
+  %borrow = store_borrow %0 to %1a : $*NonTrivialStruct
+  %f = function_ref @useNonTrivialStruct : $@convention(thin) (@guaranteed NonTrivialStruct) -> ()
+  %l = load_borrow %borrow : $*NonTrivialStruct
+  apply %f(%l) : $@convention(thin) (@guaranteed NonTrivialStruct) -> ()
+  end_borrow %l : $NonTrivialStruct
+  end_borrow %borrow : $*NonTrivialStruct
+  dealloc_stack %1 : $*NonTrivialStruct
+  %9999 = tuple ()
+  return %9999 : $()
+}

test/SILOptimizer/moveonly_addresschecker_diagnostics.sil

@@ -1,8 +1,6 @@
 // RUN: %target-sil-opt -sil-move-only-address-checker -enable-experimental-feature MoveOnlyPartialConsumption -enable-experimental-feature MoveOnlyClasses -enable-sil-verify-all %s -verify
 // RUN: %target-sil-opt -sil-move-only-address-checker -enable-experimental-feature MoveOnlyPartialConsumption -enable-experimental-feature MoveOnlyClasses -enable-sil-verify-all -move-only-diagnostics-silently-emit-diagnostics %s | %FileCheck %s
 
-// TODO: Add FileCheck
-
 // This file contains specific SIL test cases that we expect to emit
 // diagnostics. These are cases where we want to make it easy to validate
 // independent of potential changes in the frontend's emission that this
@@ -77,6 +75,7 @@ sil @get_aggstruct : $@convention(thin) () -> @owned AggStruct
 sil @nonConsumingUseKlass : $@convention(thin) (@guaranteed Klass) -> ()
 sil @nonConsumingUseNonTrivialStruct : $@convention(thin) (@guaranteed NonTrivialStruct) -> ()
 sil @consumingUseNonTrivialStruct : $@convention(thin) (@owned NonTrivialStruct) -> ()
+sil @inUseNonTrivialStruct : $@convention(thin) (@in NonTrivialStruct) -> ()
 sil @classConsume : $@convention(thin) (@owned Klass) -> () // user: %18
 sil @copyableClassConsume : $@convention(thin) (@owned CopyableKlass) -> () // user: %24
 sil @copyableClassUseMoveOnlyWithoutEscaping : $@convention(thin) (@guaranteed CopyableKlass) -> () // user: %16
@@ -529,3 +528,115 @@ bb0(%0 : @owned $NonTrivialStruct):
   %9999 = tuple()
   return %9999 : $()
 }
+
+// CHECK-LABEL: sil [ossa] @testSupportStoreBorrow1 : $@convention(thin) (@guaranteed NonTrivialStruct) -> () {
+// CHECK: bb0([[ARG:%.*]] : @guaranteed
+// CHECK:   [[STACK:%.*]] = alloc_stack
+// CHECK:   [[BORROW:%.*]] = store_borrow [[ARG]] to [[STACK]]
+// CHECK:   [[LOAD:%.*]] = load_borrow [[BORROW]]
+// CHECK:   apply {{%.*}}([[LOAD]])
+// CHECK:   end_borrow [[LOAD]]
+// CHECK:   end_borrow [[BORROW]]
+// CHECK:   dealloc_stack [[STACK]]
+// CHECK: } // end sil function 'testSupportStoreBorrow1'
+sil [ossa] @testSupportStoreBorrow1 : $@convention(thin) (@guaranteed NonTrivialStruct) -> () {
+bb0(%0 : @guaranteed $NonTrivialStruct):
+  %1 = alloc_stack $NonTrivialStruct
+  %1a = mark_must_check [no_consume_or_assign] %1 : $*NonTrivialStruct
+  %borrow = store_borrow %0 to %1a : $*NonTrivialStruct
+  %f = function_ref @nonConsumingUseNonTrivialStruct : $@convention(thin) (@guaranteed NonTrivialStruct) -> ()
+  %l = load_borrow %borrow : $*NonTrivialStruct
+  apply %f(%l) : $@convention(thin) (@guaranteed NonTrivialStruct) -> ()
+  end_borrow %l : $NonTrivialStruct
+  end_borrow %borrow : $*NonTrivialStruct
+  dealloc_stack %1 : $*NonTrivialStruct
+  %9999 = tuple ()
+  return %9999 : $()
+}
+
+// CHECK-LABEL: sil [ossa] @testSupportStoreBorrow2 : $@convention(thin) (@guaranteed NonTrivialStruct) -> () {
+// CHECK: bb0([[ARG:%.*]] : @guaranteed
+// CHECK:   [[STACK:%.*]] = alloc_stack
+// CHECK:   [[BORROW:%.*]] = store_borrow [[ARG]] to [[STACK]]
+// CHECK:   [[LOAD:%.*]] = load_borrow [[BORROW]]
+// CHECK:   apply {{%.*}}([[LOAD]])
+// CHECK:   end_borrow [[LOAD]]
+// CHECK:   end_borrow [[BORROW]]
+// CHECK:   dealloc_stack [[STACK]]
+// CHECK: } // end sil function 'testSupportStoreBorrow2'
+sil [ossa] @testSupportStoreBorrow2 : $@convention(thin) (@guaranteed NonTrivialStruct) -> () {
+bb0(%0 : @guaranteed $NonTrivialStruct):
+  %1 = alloc_stack $NonTrivialStruct
+  %1a = mark_must_check [no_consume_or_assign] %1 : $*NonTrivialStruct
+  %borrow = store_borrow %0 to %1a : $*NonTrivialStruct
+  %f = function_ref @nonConsumingUseNonTrivialStruct : $@convention(thin) (@guaranteed NonTrivialStruct) -> ()
+  %l = load [copy] %borrow : $*NonTrivialStruct
+  apply %f(%l) : $@convention(thin) (@guaranteed NonTrivialStruct) -> ()
+  destroy_value %l : $NonTrivialStruct
+  end_borrow %borrow : $*NonTrivialStruct
+  dealloc_stack %1 : $*NonTrivialStruct
+  %9999 = tuple ()
+  return %9999 : $()
+}
+
+// CHECK-LABEL: sil [ossa] @testSupportStoreBorrow3 : $@convention(thin) (@guaranteed NonTrivialStruct) -> () {
+// CHECK: bb0([[ARG:%.*]] : @guaranteed
+// CHECK:   [[STACK:%.*]] = alloc_stack
+// CHECK:   [[BORROW:%.*]] = store_borrow [[ARG]] to [[STACK]]
+// CHECK:   [[LOAD:%.*]] = load_borrow [[BORROW]]
+// CHECK:   apply {{%.*}}([[LOAD]])
+// CHECK:   end_borrow [[LOAD]]
+// CHECK:   end_borrow [[BORROW]]
+// CHECK:   dealloc_stack [[STACK]]
+// CHECK: } // end sil function 'testSupportStoreBorrow3'
+sil [ossa] @testSupportStoreBorrow3 : $@convention(thin) (@guaranteed NonTrivialStruct) -> () {
+bb0(%0 : @guaranteed $NonTrivialStruct):
+  %1 = alloc_stack $NonTrivialStruct
+  %1a = mark_must_check [no_consume_or_assign] %1 : $*NonTrivialStruct
+  %borrow = store_borrow %0 to %1a : $*NonTrivialStruct
+  %f = function_ref @nonConsumingUseNonTrivialStruct : $@convention(thin) (@guaranteed NonTrivialStruct) -> ()
+  %l = load_borrow %borrow : $*NonTrivialStruct
+  %l2 = copy_value %l : $NonTrivialStruct
+  apply %f(%l2) : $@convention(thin) (@guaranteed NonTrivialStruct) -> ()
+  destroy_value %l2 : $NonTrivialStruct
+  end_borrow %l : $NonTrivialStruct
+  end_borrow %borrow : $*NonTrivialStruct
+  dealloc_stack %1 : $*NonTrivialStruct
+  %9999 = tuple ()
+  return %9999 : $()
+}
+
+sil [ossa] @testSupportStoreBorrow4 : $@convention(thin) (@guaranteed NonTrivialStruct) -> () {
+bb0(%0 : @guaranteed $NonTrivialStruct):
+  %1 = alloc_stack $NonTrivialStruct
+  %1a = mark_must_check [no_consume_or_assign] %1 : $*NonTrivialStruct
+  // expected-error @-1 {{noncopyable 'unknown' cannot be consumed when captured by an escaping closure}}
+  %borrow = store_borrow %0 to %1a : $*NonTrivialStruct
+  %f = function_ref @consumingUseNonTrivialStruct : $@convention(thin) (@owned NonTrivialStruct) -> ()
+  %l = load_borrow %borrow : $*NonTrivialStruct
+  %l2 = copy_value %l : $NonTrivialStruct
+  apply %f(%l2) : $@convention(thin) (@owned NonTrivialStruct) -> ()
+  end_borrow %l : $NonTrivialStruct
+  end_borrow %borrow : $*NonTrivialStruct
+  dealloc_stack %1 : $*NonTrivialStruct
+  %9999 = tuple ()
+  return %9999 : $()
+}
+
+sil [ossa] @testSupportStoreBorrow5 : $@convention(thin) (@guaranteed NonTrivialStruct) -> () {
+bb0(%0 : @guaranteed $NonTrivialStruct):
+  %1 = alloc_stack $NonTrivialStruct
+  %1a = mark_must_check [no_consume_or_assign] %1 : $*NonTrivialStruct
+  // expected-error @-1 {{'unknown' is borrowed and cannot be consumed}}
+  %borrow = store_borrow %0 to %1a : $*NonTrivialStruct
+  %f = function_ref @consumingUseNonTrivialStruct : $@convention(thin) (@owned NonTrivialStruct) -> ()
+  %a = alloc_stack $NonTrivialStruct
+  copy_addr %borrow to [init] %a : $*NonTrivialStruct
+  // expected-note @-1 {{consumed here}}
+  destroy_addr %a : $*NonTrivialStruct
+  dealloc_stack %a : $*NonTrivialStruct
+  end_borrow %borrow : $*NonTrivialStruct
+  dealloc_stack %1 : $*NonTrivialStruct
+  %9999 = tuple ()
+  return %9999 : $()
+}

test/SILOptimizer/moveonly_addressonly_subscript_diagnostics.swift

@@ -203,7 +203,6 @@ public func testSubscriptReadModify_BaseLoadable_ResultAddressOnly_Var() {
 public func testSubscriptReadModify_BaseLoadable_ResultAddressOnly_Let() {
     let m = LoadableSubscriptReadModifyTester()
     m[0].nonMutatingFunc()
-    // expected-error @-1 {{copy of noncopyable typed value}}
 }
 
 public func testSubscriptReadModify_BaseLoadable_ResultAddressOnly_InOut(m: inout LoadableSubscriptReadModifyTester) {