[capture-promotion] When checking if a (struct_element_addr (project_box box)) is written to, check that all of the operands are loads, instead of returning early when we find one.

lib/SILOptimizer/IPO/CapturePromotion.cpp

@@ -694,6 +694,13 @@ static SILArgument *getBoxFromIndex(SILFunction *F, unsigned Index) {
   return Entry.getArgument(Index);
 }
 
+static bool isNonMutatingLoad(SILInstruction *I) {
+  auto *LI = dyn_cast<LoadInst>(I);
+  if (!LI)
+    return false;
+  return LI->getOwnershipQualifier() != LoadOwnershipQualifier::Take;
+}
+
 /// \brief Given a partial_apply instruction and the argument index into its
 /// callee's argument list of a box argument (which is followed by an argument
 /// for the address of the box's contents), return true if the closure is known
@@ -726,16 +733,15 @@ isNonMutatingCapture(SILArgument *BoxArg) {
   // function that mirrors isNonEscapingUse.
   auto checkAddrUse = [](SILInstruction *AddrInst) {
     if (auto *SEAI = dyn_cast<StructElementAddrInst>(AddrInst)) {
-      for (auto *UseOper : SEAI->getUses()) {
-        if (isa<LoadInst>(UseOper->getUser()))
-          return true;
-      }
-    } else if (isa<LoadInst>(AddrInst) || isa<DebugValueAddrInst>(AddrInst)
-               || isa<MarkFunctionEscapeInst>(AddrInst)
-               || isa<EndAccessInst>(AddrInst)) {
-      return true;
+      return all_of(SEAI->getUses(),
+                    [](Operand *Op) -> bool {
+                      return isNonMutatingLoad(Op->getUser());
+                    });
     }
-    return false;
+
+    return isNonMutatingLoad(AddrInst) || isa<DebugValueAddrInst>(AddrInst)
+           || isa<MarkFunctionEscapeInst>(AddrInst)
+           || isa<EndAccessInst>(AddrInst);
   };
   for (auto *Projection : Projections) {
     for (auto *UseOper : Projection->getUses()) {
@@ -744,7 +750,10 @@ isNonMutatingCapture(SILArgument *BoxArg) {
           if (!checkAddrUse(AccessUseOper->getUser()))
             return false;
         }
-      } else if (!checkAddrUse(UseOper->getUser()))
+        continue;
+      }
+
+      if (!checkAddrUse(UseOper->getUser()))
         return false;
     }
   }

test/SILOptimizer/capture_promotion.swift

@@ -1,4 +1,4 @@
-// RUN: %target-swift-frontend %s -emit-sil -o - -verify | %FileCheck %s
+// RUN: %target-swift-frontend %s -emit-sil -o - | %FileCheck %s
 
 class Foo {
   func foo() -> Int {

test/SILOptimizer/capture_promotion_ownership.sil

@@ -332,3 +332,53 @@ bb0(%0: @trivial $*Int, %1 : @owned $<τ_0_0> { var τ_0_0 } <Foo>, %2 : @owned
   %16 = tuple()
   return %16 : $()
 }
+
+// This test makes sure that we properly handle non load uses of
+// struct_element_addr that always have load users with the load occuring before
+// and after the element in the use list.
+
+sil @mutate_int : $@convention(thin) (@inout Int) -> ()
+
+// CHECK-LABEL: sil @test_closure_multiple_uses_of_struct_element_addr : $@convention(thin) () -> @owned @callee_owned () -> () {
+// CHECK: [[FUNC:%.*]] = function_ref @closure_multiple_uses_of_struct_element_addr : $@convention(thin)
+// CHECK: partial_apply [[FUNC]](
+// CHECK: } // end sil function 'test_closure_multiple_uses_of_struct_element_addr'
+sil @test_closure_multiple_uses_of_struct_element_addr : $@convention(thin) () -> @owned @callee_owned () -> () {
+bb0:
+  %0 = function_ref @baz_init : $@convention(thin) (@thin Baz.Type) -> @owned Baz
+  %1 = metatype $@thin Baz.Type
+
+  %2 = alloc_box $<τ_0_0> { var τ_0_0 } <Baz>
+  %3 = project_box %2 : $<τ_0_0> { var τ_0_0 } <Baz>, 0
+  %4 = apply %0(%1) : $@convention(thin) (@thin Baz.Type) -> @owned Baz
+  store %4 to [init] %3 : $*Baz
+
+  %5 = alloc_box $<τ_0_0> { var τ_0_0 } <Baz>
+  %6 = project_box %5 : $<τ_0_0> { var τ_0_0 } <Baz>, 0
+  %7 = apply %0(%1) : $@convention(thin) (@thin Baz.Type) -> @owned Baz
+  store %7 to [init] %6 : $*Baz
+
+  %8 = function_ref @closure_multiple_uses_of_struct_element_addr : $@convention(thin) (@owned <τ_0_0> { var τ_0_0 } <Baz>, @owned <τ_0_0> { var τ_0_0 } <Baz>) -> ()
+  %9 = partial_apply %8(%2, %5) : $@convention(thin) (@owned <τ_0_0> { var τ_0_0 } <Baz>, @owned <τ_0_0> { var τ_0_0 } <Baz>) -> ()
+
+  return %9 : $@callee_owned () -> ()
+}
+
+sil @closure_multiple_uses_of_struct_element_addr : $@convention(thin) (@owned <τ_0_0> { var τ_0_0 } <Baz>, @owned <τ_0_0> { var τ_0_0 } <Baz>) -> () {
+bb0(%0 : @owned $<τ_0_0> { var τ_0_0 } <Baz>, %1 : @owned $<τ_0_0> { var τ_0_0 } <Baz>):
+  %2 = project_box %0 : $<τ_0_0> { var τ_0_0 } <Baz>, 0
+  %3 = struct_element_addr %2 : $*Baz, #Baz.x
+  %4 = load [trivial] %3 : $*Int
+  %5 = function_ref @mutate_int : $@convention(thin) (@inout Int) -> ()
+  apply %5(%3) : $@convention(thin) (@inout Int) -> ()
+
+  %6 = project_box %1 : $<τ_0_0> { var τ_0_0 } <Baz>, 0
+  %7 = struct_element_addr %6 : $*Baz, #Baz.x
+  apply %5(%7) : $@convention(thin) (@inout Int) -> ()
+  %8 = load [trivial] %7 : $*Int
+
+  destroy_value %1 : $<τ_0_0> { var τ_0_0 } <Baz>
+  destroy_value %0 : $<τ_0_0> { var τ_0_0 } <Baz>
+  %9999 = tuple()
+  return %9999 : $()
+}

test/SILOptimizer/capture_promotion_ownership.swift

@@ -1,4 +1,4 @@
-// RUN: %target-swift-frontend %s -enable-sil-ownership -disable-sil-linking -emit-sil -o - -verify | %FileCheck %s
+// RUN: %target-swift-frontend %s -enable-sil-ownership -disable-sil-linking -emit-sil -o - | %FileCheck %s
 
 // NOTE: We add -disable-sil-linking to the compile line to ensure that we have
 // access to declarations for standard library types, but not definitions. This