TempRValueOpt: fix the handling of load [take]

load [take] was not considered as a use and it was not detected if it's in a different basic block.
This fixes a miscompile in case there is a modification of the copy-source before a load [take].

rdar://problem/69757314

Author: eeckstein

lib/SILOptimizer/Transforms/TempRValueElimination.cpp

@@ -241,7 +241,8 @@ collectLoads(Operand *addressUse, CopyAddrInst *originalCopy,
   case SILInstructionKind::UncheckedAddrCastInst:
     return collectLoadsFromProjection(cast<SingleValueInstruction>(user),
                                       originalCopy, loadInsts);
-  case SILInstructionKind::LoadInst:
+
+  case SILInstructionKind::LoadInst: {
     // Loads are the end of the data flow chain. The users of the load can't
     // access the temporary storage.
     //
@@ -251,14 +252,17 @@ collectLoads(Operand *addressUse, CopyAddrInst *originalCopy,
     // function. So if we have such a load [take], we /must/ have a
     // reinitialization or an alloc_stack that does not fit the pattern we are
     // expecting from SILGen. Be conservative and return false.
-    if (auto *li = dyn_cast<LoadInst>(user)) {
-      if (li->getOwnershipQualifier() == LoadOwnershipQualifier::Take) {
-        return false;
-      }
+    auto *li = cast<LoadInst>(user);
+    if (li->getOwnershipQualifier() == LoadOwnershipQualifier::Take &&
+        // Only accept load [take] if it takes the whole temporary object.
+        // load [take] from a projection would destroy only a part of the
+        // temporary and we don't handle this.
+        address != originalCopy->getDest()) {
+      return false;
     }
     loadInsts.insert(user);
     return true;
-
+  }
   case SILInstructionKind::LoadBorrowInst:
     loadInsts.insert(user);
     return true;
@@ -432,15 +436,6 @@ bool TempRValueOptPass::tryOptimizeCopyIntoTemp(CopyAddrInst *copyInst) {
     if (isa<DestroyAddrInst>(user) || isa<DeallocStackInst>(user))
       continue;
 
-    // Same for load [take] on the top level temp object. SILGen always takes
-    // whole values from temporaries. If we have load [take] on projections from
-    // our base, we fail since those would be re-initializations.
-    if (auto *li = dyn_cast<LoadInst>(user)) {
-      if (li->getOwnershipQualifier() == LoadOwnershipQualifier::Take) {
-        continue;
-      }
-    }
-
     if (!collectLoads(useOper, copyInst, loadInsts))
       return false;
   }

test/SILOptimizer/temp_rvalue_opt_ossa.sil

@@ -128,6 +128,22 @@ bb0(%0 : $*B, %1 : $*GS<B>):
   return %9999 : $()
 }
 
+// CHECK-LABEL: sil [ossa] @store_before_load_take
+// CHECK:   [[STK:%[0-9]+]] = alloc_stack
+// CHECK:   copy_addr [take] %0 to [initialization] [[STK]]
+// CHECK:   store
+// CHECK:   load [take] [[STK]]
+// CHECK:   return
+// CHECK: } // end sil function 'store_before_load_take'
+sil [ossa] @store_before_load_take : $@convention(thin) (@inout Builtin.NativeObject, @owned Builtin.NativeObject) -> @owned Builtin.NativeObject {
+bb0(%0 : $*Builtin.NativeObject, %1 : @owned $Builtin.NativeObject):
+  %stk = alloc_stack $Builtin.NativeObject
+  copy_addr [take] %0 to [initialization] %stk : $*Builtin.NativeObject
+  store %1 to [init] %0 : $*Builtin.NativeObject
+  %obj = load [take] %stk : $*Builtin.NativeObject
+  dealloc_stack %stk : $*Builtin.NativeObject
+  return %obj : $Builtin.NativeObject
+}
 // CHECK-LABEL: sil [ossa] @load_in_wrong_block
 // CHECK:      bb0(%0 : $*GS<B>):
 // CHECK-NEXT:   alloc_stack
@@ -693,38 +709,28 @@ bb0(%0 : @owned $GS<Builtin.NativeObject>):
   return %obj : $GS<Builtin.NativeObject>
 }
 
-// CHECK-LABEL: sil [ossa] @hoist_load_copy_to_src_copy_addr_site_two_takes : $@convention(thin) (@in GS<Builtin.NativeObject>) -> @owned GS<Builtin.NativeObject> {
-// CHECK: bb0([[ARG:%.*]] :
-// CHECK:   [[COPY_1:%.*]] = load [copy] [[ARG]]
-// CHECK:   [[COPY_2:%.*]] = load [copy] [[ARG]]
-// CHECK:   destroy_addr [[ARG]]
-// CHECK:   cond_br undef, bb1, bb2
-//
+// CHECK-LABEL: sil [ossa] @dont_optimize_with_load_in_different_block
+// CHECK:   [[STK:%[0-9]+]] = alloc_stack
+// CHECK:   copy_addr %0 to [initialization] [[STK]]
 // CHECK: bb1:
-// CHECK:   destroy_value [[COPY_1]]
-// CHECK:   br bb3([[COPY_2]] :
-//
+// CHECK:   load [take] [[STK]]
 // CHECK: bb2:
-// CHECK:   destroy_value [[COPY_2]]
-// CHECK:   br bb3([[COPY_1]] :
-//
-// CHECK: bb3([[RESULT:%.*]] : @owned
-// CHECK:   apply {{%.*}}([[RESULT]])
-// CHECK:   return [[RESULT]]
-// CHECK: } // end sil function 'hoist_load_copy_to_src_copy_addr_site_two_takes'
-sil [ossa] @hoist_load_copy_to_src_copy_addr_site_two_takes : $@convention(thin) (@in GS<Builtin.NativeObject>) -> @owned GS<Builtin.NativeObject> {
-bb0(%0 : $*GS<Builtin.NativeObject>):
+// CHECK:   copy_addr %1 to %0
+// CHECK:   load [take] [[STK]]
+// CHECK: } // end sil function 'dont_optimize_with_load_in_different_block'
+sil [ossa] @dont_optimize_with_load_in_different_block : $@convention(thin) (@inout GS<Builtin.NativeObject>, @in_guaranteed GS<Builtin.NativeObject>) -> @owned GS<Builtin.NativeObject> {
+bb0(%0 : $*GS<Builtin.NativeObject>, %1 : $*GS<Builtin.NativeObject>):
   %f = function_ref @use_gsbase_builtinnativeobject : $@convention(thin) (@guaranteed GS<Builtin.NativeObject>) -> ()
   %stk = alloc_stack $GS<Builtin.NativeObject>
   copy_addr %0 to [initialization] %stk : $*GS<Builtin.NativeObject>
-  destroy_addr %0 : $*GS<Builtin.NativeObject>
   cond_br undef, bb1, bb2
 
 bb1:
   %obj = load [take] %stk : $*GS<Builtin.NativeObject>
   br bb3(%obj : $GS<Builtin.NativeObject>)
 
 bb2:
+  copy_addr %1 to %0 : $*GS<Builtin.NativeObject>
   %obj2 = load [take] %stk : $*GS<Builtin.NativeObject>
   br bb3(%obj2 : $GS<Builtin.NativeObject>)