Fix MoveOnlyWrappedTypeEliminator handling of store_borrow.

Defer visiting an instruction until its operands have been visited. Otherwise,
this pass will crash during ownership verification with invalid operand
ownership.

Fixes rdar://152879038 ([moveonly] MoveOnlyWrappedTypeEliminator ownership
verifier crashes on @_addressableSelf)

Author: atrick

lib/SILOptimizer/Mandatory/MoveOnlyWrappedTypeEliminator.cpp

@@ -60,6 +60,28 @@ namespace {
 
 struct SILMoveOnlyWrappedTypeEliminatorVisitor
     : SILInstructionVisitor<SILMoveOnlyWrappedTypeEliminatorVisitor, bool> {
+
+  // Instructions waiting to be visited.
+  llvm::SmallSetVector<SILInstruction *, 8> &pendingInsts;
+  llvm::SmallVector<SILInstruction *, 8> deferredInsts;
+
+  SILMoveOnlyWrappedTypeEliminatorVisitor(llvm::SmallSetVector<SILInstruction *, 8> &pendingInsts):
+    pendingInsts(pendingInsts)
+  {}
+
+  // If 'user's operand is not yet visited, push it onto the end of
+  // 'pendingInsts' and pretend we didn't see it yet. This is only relevant for
+  // non-address operands in which ownership should be consistent.
+  bool waitFor(SILInstruction *user, SILValue operand) {
+    if (auto *def = operand.getDefiningInstruction()) {
+      if (pendingInsts.contains(def)) {
+        deferredInsts.push_back(user);
+        return false;
+      }
+    }
+    return true;
+  }
+
   bool visitSILInstruction(SILInstruction *inst) {
     llvm::errs() << "Unhandled SIL Instruction: " << *inst;
     llvm_unreachable("error");
@@ -79,13 +101,19 @@ struct SILMoveOnlyWrappedTypeEliminatorVisitor
   }
 
   bool visitStoreInst(StoreInst *si) {
+    if (!waitFor(si, si->getSrc())) {
+      return false;
+    }
     if (!si->getSrc()->getType().isTrivial(*si->getFunction()))
       return false;
     si->setOwnershipQualifier(StoreOwnershipQualifier::Trivial);
     return true;
   }
 
   bool visitStoreBorrowInst(StoreBorrowInst *si) {
+    if (!waitFor(si, si->getSrc())) {
+      return false;
+    }
     if (!si->getSrc()->getType().isTrivial(*si->getFunction()))
       return false;
     SILBuilderWithScope b(si);
@@ -332,6 +360,9 @@ bool SILMoveOnlyWrappedTypeEliminator::process() {
     return true;
   };
 
+  // (1) Check each value's type for the MoveOnly wrapper, (2) strip the wrapper
+  // type, and (3) add all uses to 'touchedInsts' in forward order for
+  // efficiency.
   for (auto &bb : *fn) {
     for (auto *arg : bb.getArguments()) {
       bool relevant = visitValue(arg);
@@ -377,9 +408,18 @@ bool SILMoveOnlyWrappedTypeEliminator::process() {
     madeChange = true;
   }
 
-  SILMoveOnlyWrappedTypeEliminatorVisitor visitor;
-  while (!touchedInsts.empty()) {
-    visitor.visit(touchedInsts.pop_back_val());
+  SILMoveOnlyWrappedTypeEliminatorVisitor visitor{touchedInsts};
+  while(true) {
+    while (!touchedInsts.empty()) {
+      visitor.visit(touchedInsts.pop_back_val());
+    }
+    if (visitor.deferredInsts.empty())
+      break;
+
+    for (auto *inst : visitor.deferredInsts) {
+      touchedInsts.insert(inst);
+    }
+    visitor.deferredInsts.clear();
   }
 
   return madeChange;

test/SILOptimizer/moveonly_addressors.swift

@@ -1,9 +1,10 @@
-// RUN: %target-swift-frontend  -enable-experimental-feature BuiltinModule -parse-stdlib -module-name Swift -DADDRESS_ONLY -emit-sil -verify %s
-// RUN: %target-swift-frontend  -enable-experimental-feature BuiltinModule -parse-stdlib -module-name Swift -DLOADABLE -emit-sil -verify %s
-// RUN: %target-swift-frontend  -enable-experimental-feature BuiltinModule -parse-stdlib -module-name Swift -DTRIVIAL -emit-sil -verify %s
-// RUN: %target-swift-frontend  -enable-experimental-feature BuiltinModule -parse-stdlib -module-name Swift -DEMPTY -emit-sil -verify %s
+// RUN: %target-swift-frontend  -enable-experimental-feature BuiltinModule -enable-experimental-feature AddressableParameters -parse-stdlib -module-name Swift -DADDRESS_ONLY -emit-sil -verify %s
+// RUN: %target-swift-frontend  -enable-experimental-feature BuiltinModule -enable-experimental-feature AddressableParameters -parse-stdlib -module-name Swift -DLOADABLE -emit-sil -verify %s
+// RUN: %target-swift-frontend  -enable-experimental-feature BuiltinModule -enable-experimental-feature AddressableParameters -parse-stdlib -module-name Swift -DTRIVIAL -emit-sil -verify %s
+// RUN: %target-swift-frontend  -enable-experimental-feature BuiltinModule -enable-experimental-feature AddressableParameters -parse-stdlib -module-name Swift -DEMPTY -emit-sil -verify %s
 
 // REQUIRES: swift_feature_BuiltinModule
+// REQUIRES: swift_feature_AddressableParameters
 
 // TODO: Use the real stdlib types once `UnsafePointer` supports noncopyable
 // types.
@@ -54,6 +55,13 @@ struct S {
     unsafeAddress { return makeUpAPointer() }
     unsafeMutableAddress { return makeUpAPointer() }
   }
+
+  var pointer: Builtin.RawPointer {
+    @_addressableSelf
+    get {
+      Builtin.addressOfBorrow(self)
+    }
+  }
 }
 
 struct SNC: ~Copyable {
@@ -122,3 +130,10 @@ func test(mut_snc snc: inout SNC) {
   mod(&snc.mutableData)
   take(snc.mutableData) // expected-error{{missing reinitialization of inout parameter 'snc.mutableData' after consume}} expected-note{{consumed here}}
 }
+
+// Test calling passing a no-implicit-copy value as an indirect funtion argument. This requires the
+// MoveOnlyWrappedTypeEliminator to handle a store_borrow with multiple move-only operands without crashing in the
+// verifier with invalid operand ownership.
+func testAddressableNoImplicitCopy(s: consuming S) -> Builtin.RawPointer {
+  s.pointer
+}

test/SILOptimizer/moveonly_type_eliminator.sil

@@ -26,6 +26,8 @@ struct TrivialPair {
 
 sil @use_trivial : $@convention(thin) (Trivial) -> ()
 
+sil @readIndirectTrivial : $@convention(method) (@in_guaranteed Trivial) -> ()
+
 class Klass {
   func foo()
 }
@@ -674,3 +676,38 @@ bb0(%0 : $@thin Trivial.Type):
   %13 = tuple ()
   return %13 : $()
 }
+
+// Test deferred visitation of a store_borrow until both operands have been visited.
+//
+// CHECK-LABEL: sil hidden [ossa] @testLoadStoreBorrow : $@convention(thin) (Trivial) -> () {
+// CHECK: bb0(%0 : @noImplicitCopy @_eagerMove $Trivial):
+// CHECK: [[VAR:%[0-9]+]] = alloc_stack [var_decl] [moveable_value_debuginfo] $Trivial, var, name "c"
+// CHECK: store %0 to [trivial] [[VAR]] : $*Trivial
+// CHECK: [[ACCESS:%[0-9]+]] = begin_access [read] [static] [[VAR]] : $*Trivial
+// CHECK: [[LD:%[0-9]+]] = load [trivial] [[ACCESS]] : $*Trivial
+// CHECK: [[ARG:%[0-9]+]] = alloc_stack [moveable_value_debuginfo] $Trivial
+// CHECK: store [[LD]] to [trivial] [[ARG]] : $*Trivial
+// CHECK-NOT: borrow
+// CHECK-LABEL: } // end sil function 'testLoadStoreBorrow'
+sil hidden [ossa] @testLoadStoreBorrow : $@convention(thin) (Trivial) -> () {
+bb0(%0 : @noImplicitCopy @_eagerMove $Trivial):
+  %1 = alloc_stack [var_decl] $@moveOnly Trivial, var, name "c", type $@moveOnly Trivial
+  %2 = moveonlywrapper_to_copyable_addr %1
+  store %0 to [trivial] %2
+  %4 = begin_access [read] [static] %1
+  %5 = load_borrow %4
+  %6 = alloc_stack $@moveOnly Trivial
+  // This store_borrow has two move-only operands.
+  %7 = store_borrow %5 to %6
+  %8 = function_ref @readIndirectTrivial : $@convention(method) (@in_guaranteed Trivial) -> ()
+  %9 = moveonlywrapper_to_copyable_addr %7
+  %10 = apply %8(%9) : $@convention(method) (@in_guaranteed Trivial) -> ()
+  end_borrow %7
+  end_borrow %5
+  end_access %4
+  dealloc_stack %6
+  destroy_addr %1
+  dealloc_stack %1
+  %99 = tuple ()
+  return %99 : $()
+}