Pass the indices for writeback-conflict diagnostics on coroutines

lib/SILGen/ArgumentSource.cpp

@@ -473,3 +473,38 @@ bool ArgumentSource::isObviouslyEqual(const ArgumentSource &other) const {
   }
   llvm_unreachable("bad kind");
 }
+
+PreparedArguments PreparedArguments::copyForDiagnostics() const {
+  if (isNull())
+    return PreparedArguments();
+
+  assert(isValid());
+  PreparedArguments result(getFormalType(), isScalar());
+  for (auto &arg : Arguments) {
+    result.Arguments.push_back(arg.copyForDiagnostics());
+  }
+  return result;
+}
+
+ArgumentSource ArgumentSource::copyForDiagnostics() const {
+  switch (StoredKind) {
+  case Kind::Invalid:
+    return ArgumentSource();
+  case Kind::LValue:
+    // We have no way to copy an l-value for diagnostics.
+    return {getKnownLValueLocation(), LValue()};
+  case Kind::RValue:
+    return {getKnownRValueLocation(), asKnownRValue().copyForDiagnostics()};
+  case Kind::Expr:
+    return asKnownExpr();
+  case Kind::Tuple: {
+    auto &tuple = Storage.get<TupleStorage>(StoredKind);
+    SmallVector<ArgumentSource, 4> copiedElements;
+    for (auto &elt : tuple.Elements) {
+      copiedElements.push_back(elt.copyForDiagnostics());
+    }
+    return {tuple.Loc, tuple.SubstType, copiedElements};
+  }
+  }
+  llvm_unreachable("bad kind");
+}

lib/SILGen/ArgumentSource.h

@@ -317,6 +317,8 @@ class ArgumentSource {
 
   bool isObviouslyEqual(const ArgumentSource &other) const;
 
+  ArgumentSource copyForDiagnostics() const;
+
   void dump() const;
   void dump(raw_ostream &os, unsigned indent = 0) const;
 
@@ -426,6 +428,8 @@ class PreparedArguments {
   PreparedArguments copy(SILGenFunction &SGF, SILLocation loc) const;
 
   bool isObviouslyEqual(const PreparedArguments &other) const;
+
+  PreparedArguments copyForDiagnostics() const;
 };
 
 } // end namespace Lowering

lib/SILGen/RValue.cpp

@@ -813,3 +813,13 @@ SILType RValue::getLoweredImplodedTupleType(SILGenFunction &SGF) const & {
     return loweredType.getAddressType();
   return loweredType.getObjectType();
 }
+
+RValue RValue::copyForDiagnostics() const {
+  assert(!isInSpecialState());
+  assert(isComplete());
+  RValue result(type);
+  for (auto value : values)
+    result.values.push_back(value);
+  result.elementsToBeAdded = 0;
+  return result;
+}

lib/SILGen/RValue.h

@@ -355,6 +355,8 @@ class RValue {
   /// Borrow all subvalues of the rvalue.
   RValue borrow(SILGenFunction &SGF, SILLocation loc) const &;
 
+  RValue copyForDiagnostics() const;
+
   static bool areObviouslySameValue(SILValue lhs, SILValue rhs);
   bool isObviouslyEqual(const RValue &rhs) const;
 

lib/SILGen/SILGenLValue.cpp

@@ -1512,14 +1512,22 @@ namespace {
 
   class EndApplyPseudoComponent : public WritebackPseudoComponent {
     CleanupHandle EndApplyHandle;
-    Optional<AccessedStorage> Access;
+    AbstractStorageDecl *Storage;
+    bool IsSuper;
+    PreparedArguments PeekedIndices;
+    Expr *IndexExprForDiagnostics;
   public:
     EndApplyPseudoComponent(const LValueTypeData &typeData,
                             CleanupHandle endApplyHandle,
-                            Optional<AccessedStorage> access)
+                            AbstractStorageDecl *storage,
+                            bool isSuper,
+                            PreparedArguments &&peekedIndices,
+                            Expr *indexExprForDiagnostics)
       : WritebackPseudoComponent(typeData),
         EndApplyHandle(endApplyHandle),
-        Access(access) {}
+        Storage(storage), IsSuper(isSuper),
+        PeekedIndices(std::move(peekedIndices)),
+        IndexExprForDiagnostics(indexExprForDiagnostics) {}
 
   private:
     void writeback(SILGenFunction &SGF, SILLocation loc,
@@ -1539,7 +1547,9 @@ namespace {
     }
 
     Optional<AccessedStorage> getAccessedStorage() const override {
-      return Access;
+      return AccessedStorage{Storage, IsSuper,
+                             PeekedIndices.isNull() ? nullptr : &PeekedIndices,
+                             IndexExprForDiagnostics};
     }
   };
 
@@ -1570,12 +1580,9 @@ namespace {
 
       ManagedValue result;
 
-      LogicalPathComponent::AccessedStorage abstractStorage = {
-        Storage, IsSuper, /*indices*/ nullptr, IndexExprForDiagnostics
-      };
-
       auto args =
         std::move(*this).prepareAccessorArgs(SGF, loc, base, Accessor);
+      auto peekedIndices = args.Indices.copyForDiagnostics();
       SmallVector<ManagedValue, 4> yields;
       auto endApplyHandle =
         SGF.emitCoroutineAccessor(
@@ -1585,7 +1592,9 @@ namespace {
       // Push a writeback that ends the access.
       std::unique_ptr<LogicalPathComponent>
         component(new EndApplyPseudoComponent(getTypeData(), endApplyHandle,
-                                              abstractStorage));
+                                              Storage, IsSuper,
+                                              std::move(peekedIndices),
+                                              IndexExprForDiagnostics));
       pushWriteback(SGF, loc, std::move(component), ManagedValue(),
                     MaterializedLValue());
 

test/SILGen/writeback_conflict_diagnostics.swift

@@ -111,3 +111,19 @@ func testMultiArrayWithoutAddressors(
   swap(&global_array_without_addressors[0][j], &array[j][i])  // ok
 }
 
+// rdar://43802132
+struct ArrayWithReadModify<T> {
+  init(value: T) { self.property = value }
+  var property: T
+  subscript(i: Int) -> T {
+    _read { yield property }
+    _modify { yield &property }
+  }
+}
+
+func testArrayWithReadModify<T>(array: ArrayWithReadModify<T>) {
+  var copy = array
+  swap(&copy[0], &copy[1])
+  swap(&copy[0], // expected-note {{concurrent writeback occurred here}}
+       &copy[0]) // expected-error {{inout writeback through subscript occurs in multiple arguments to call}}
+}