Merge pull request swiftlang#39691 from meg-gupta/temprvofix

Author: swift-ci

lib/SILOptimizer/Transforms/TempRValueElimination.cpp

@@ -18,9 +18,10 @@
 
 #define DEBUG_TYPE "sil-temp-rvalue-opt"
 
-#include "swift/SIL/DebugUtils.h"
 #include "swift/SIL/BasicBlockUtils.h"
+#include "swift/SIL/DebugUtils.h"
 #include "swift/SIL/MemAccessUtils.h"
+#include "swift/SIL/OwnershipUtils.h"
 #include "swift/SIL/SILArgument.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILVisitor.h"
@@ -276,9 +277,24 @@ collectLoads(Operand *addressUse, CopyAddrInst *originalCopy,
     loadInsts.insert(user);
     return true;
   }
-  case SILInstructionKind::LoadBorrowInst:
+  case SILInstructionKind::LoadBorrowInst: {
     loadInsts.insert(user);
-    return true;
+    BorrowedValue borrow(cast<LoadBorrowInst>(user));
+    auto visitEndScope = [&](Operand *op) -> bool {
+      auto *opUser = op->getUser();
+      if (auto *endBorrow = dyn_cast<EndBorrowInst>(opUser)) {
+        if (endBorrow->getParent() != block)
+          return false;
+        loadInsts.insert(endBorrow);
+        return true;
+      }
+      // Don't look further if we see a reborrow.
+      assert(cast<BranchInst>(opUser));
+      return false;
+    };
+    auto res = borrow.visitLocalScopeEndingUses(visitEndScope);
+    return res;
+  }
   case SILInstructionKind::FixLifetimeInst:
     // If we have a fixed lifetime on our alloc_stack, we can just treat it like
     // a load and re-write it so that it is on the old memory or old src object.

test/SILOptimizer/temp_rvalue_opt_ossa.sil

@@ -21,6 +21,15 @@ struct Two {
   var b: Klass
 }
 
+struct NonTrivialStruct {
+  var val:Klass
+}
+
+public enum FakeOptional<T> {
+  case none
+  case some(T)
+}
+
 sil @unknown : $@convention(thin) () -> ()
 
 sil [ossa] @guaranteed_user : $@convention(thin) (@guaranteed Klass) -> ()
@@ -1410,3 +1419,49 @@ bb2:
   unwind
 }
 
+// This does not get optimized correctly because of the conservative treatment of load_borrow/end_borrow in MemBehavior
+// CHECK-LABEL: sil [ossa] @test_temprvoborrowboundary1 :
+// CHECK: copy_addr
+// TODO-CHECK-NOT:   alloc_stack
+// TODO-CHECK-NOT:   copy_addr
+// TODO-CHECK: [[ADDR:%.*]] = unchecked_take_enum_data_addr %0 : $*FakeOptional<NonTrivialStruct>, #FakeOptional.some!enumelt
+// TODO-CHECK: [[ELE:%.*]] = struct_element_addr [[ADDR]] : $*NonTrivialStruct, #NonTrivialStruct.val
+// TODO-CHECK: [[LD:%.*]] = load_borrow [[ELE]] : $*Klass
+// TODO-CHECK: end_borrow [[LD]] : $Klass
+// TODO-CHECK: destroy_addr [[ADDR]]
+// CHECK: } // end sil function 'test_temprvoborrowboundary1'
+sil [ossa] @test_temprvoborrowboundary1 : $@convention(thin) (@in FakeOptional<NonTrivialStruct>) -> () {
+bb0(%0 : $*FakeOptional<NonTrivialStruct>):
+  %1 = alloc_stack $NonTrivialStruct
+  %2 = unchecked_take_enum_data_addr %0 : $*FakeOptional<NonTrivialStruct>, #FakeOptional.some!enumelt
+  copy_addr [take] %2 to [initialization] %1 : $*NonTrivialStruct
+  %4 = struct_element_addr %1 : $*NonTrivialStruct, #NonTrivialStruct.val
+  %5 = load_borrow %4 : $*Klass
+  end_borrow %5 : $Klass
+  destroy_addr %1 : $*NonTrivialStruct
+  dealloc_stack %1 : $*NonTrivialStruct
+  %res = tuple ()
+  return %res : $()
+}
+
+// This does not get optimized because of the end borrow in a different block 
+// CHECK-LABEL: sil [ossa] @test_temprvoborrowboundary2 :
+// CHECK:   copy_addr
+// CHECK: } // end sil function 'test_temprvoborrowboundary2'
+sil [ossa] @test_temprvoborrowboundary2 : $@convention(thin) (@in FakeOptional<NonTrivialStruct>) -> () {
+bb0(%0 : $*FakeOptional<NonTrivialStruct>):
+  %1 = alloc_stack $NonTrivialStruct
+  %2 = unchecked_take_enum_data_addr %0 : $*FakeOptional<NonTrivialStruct>, #FakeOptional.some!enumelt
+  copy_addr [take] %2 to [initialization] %1 : $*NonTrivialStruct
+  %4 = struct_element_addr %1 : $*NonTrivialStruct, #NonTrivialStruct.val
+  %5 = load_borrow %4 : $*Klass
+  br bb1(%5 : $Klass)
+
+bb1(%6 : @guaranteed $Klass):
+  end_borrow %6 : $Klass
+  destroy_addr %1 : $*NonTrivialStruct
+  dealloc_stack %1 : $*NonTrivialStruct
+  %res = tuple ()
+  return %res : $()
+}
+