[TempRValueOpts] Opt some lexical alloc_stacks.

Previously, whenever an alloc_stack [lexical] was seen, optimization
bailed conservatively.

In the fullness of time, we should optimize an alloc_stack [lexical] so
long as:
(1) the root of the source of the store/copy_addr is lexical
(2) the range in which the alloc_stack [lexical] contains the value
    store'd/copy_addr'd into it (i.e. the range ending at
    destroy_addr/etc) is contained within the live range of that lexical
    root

Here, an incremental step in that direction is taken: if the base of the
accessed storage of the source of a copy_addr is a guaranteed function
argument, then we know
(1) the base is lexical--guaranteed function arguments are always
    lexical
(2) the range in which the alloc_stack contains the value copy_addr'd in
    is trivially contained within the range of that lexical root--the
    live range of a guaranteed function argument is the whole function

Added TODOs for the full optimization.

Author: nate-chandler

lib/SILOptimizer/Transforms/TempRValueElimination.cpp

@@ -501,11 +501,29 @@ void TempRValueOptPass::tryOptimizeCopyIntoTemp(CopyAddrInst *copyInst) {
   if (!tempObj)
     return;
 
-  // If the storage corresponds to a source-level var, do not optimize here.
-  // Mem2Reg will transform the lexical alloc_stack into a lexical begin_borrow
-  // which will ensure that the value's lifetime isn't observably shortened.
-  if (tempObj->isLexical())
-    return;
+  // If the temporary storage is lexical, it either came from a source-level var
+  // or was marked lexical because it was passed to a function that has been
+  // inlined.
+  // TODO: [begin_borrow_addr] Once we can mark addresses as being borrowed, we
+  //       won't need to mark alloc_stacks lexical during inlining.  At that
+  //       point, the above comment should change, but the implementation
+  //       remains the same.
+  //
+  // In either case, we can eliminate the temporary if the source of the copy is
+  // lexical and it is live for longer than the temporary.
+  if (tempObj->isLexical()) {
+    // TODO: Determine whether the base of the copy_addr's source is lexical and
+    //       its live range contains the range in which the alloc_stack
+    //       contains the value copied into it via the copy_addr.
+    //
+    // For now, only look for guaranteed arguments.
+    auto storage = AccessStorageWithBase::compute(copyInst->getSrc());
+    if (!storage.base)
+      return;
+    if (auto *arg = dyn_cast<SILFunctionArgument>(storage.base))
+      if (arg->getOwnershipKind() != OwnershipKind::Guaranteed)
+        return;
+  }
 
   bool isOSSA = copyInst->getFunction()->hasOwnership();
 
@@ -645,10 +663,20 @@ TempRValueOptPass::tryOptimizeStoreIntoTemp(StoreInst *si) {
     return std::next(si->getIterator());
   }
 
-  // If the storage corresponds to a source-level var, do not optimize here.
-  // Mem2Reg will transform the lexical alloc_stack into a lexical begin_borrow
-  // which will ensure that the value's lifetime isn't observably shortened.
+  // If the temporary storage is lexical, it either came from a source-level var
+  // or was marked lexical because it was passed to a function that has been
+  // inlined.
+  // TODO: [begin_borrow_addr] Once we can mark addresses as being borrowed, we
+  //       won't need to mark alloc_stacks lexical during inlining.  At that
+  //       point, the above comment should change, but the implementation
+  //       remains the same.
+  //
+  // In either case, we can eliminate the temporary if the source of the store
+  // is lexical and it is live for longer than the temporary.
   if (tempObj->isLexical()) {
+    // TODO: Find the lexical root of the source, if any, and allow optimization
+    //       if its live range contains the range in which the alloc_stack
+    //       contains the value stored into it.
     return std::next(si->getIterator());
   }
 

test/SILOptimizer/temp_rvalue_opt_ossa.sil

@@ -19,6 +19,10 @@ class Klass {
   init()
 }
 
+class OtherClass {
+  var klass: Klass
+}
+
 struct Two {
   var a: Klass
   var b: Klass
@@ -1544,3 +1548,35 @@ bb1(%6 : @guaranteed $Klass):
   return %res : $()
 }
 
+// Lexical alloc_stacks can still be eliminated if their sources are guaranteed
+// function arguments (because those are lexical and have a live range that
+// contains the live range of the value stored into the alloc stack).
+// CHECK-LABEL: sil [ossa] @copy_addr__lexical_alloc_stack_temporary__guaranteed_function_arg : {{.*}} {
+// CHECK-NOT:     alloc_stack
+// CHECK-LABEL: } // end sil function 'copy_addr__lexical_alloc_stack_temporary__guaranteed_function_arg'
+sil [ossa] @copy_addr__lexical_alloc_stack_temporary__guaranteed_function_arg : $@convention(thin) (@guaranteed OtherClass) -> () {
+entry(%instance : @guaranteed $OtherClass):
+  %field_ptr = ref_element_addr %instance : $OtherClass, #OtherClass.klass
+  %addr = alloc_stack [lexical] $Klass
+  copy_addr %field_ptr to [initialization] %addr : $*Klass
+  destroy_addr %addr : $*Klass
+  dealloc_stack %addr : $*Klass
+  %retval = tuple ()
+  return %retval : $()
+}
+
+// Lexical alloc_stacks can't be eliminated even if their source is a function
+// argument if that function argument is inout.
+// CHECK-LABEL: sil [ossa] @copy_addr__lexical_alloc_stack_temporary__inout_function_arg : {{.*}} {
+// CHECK:         alloc_stack [lexical]
+// CHECK-LABEL: } // end sil function 'copy_addr__lexical_alloc_stack_temporary__inout_function_arg'
+sil [ossa] @copy_addr__lexical_alloc_stack_temporary__inout_function_arg : $@convention(thin) (@inout NonTrivialStruct) -> () {
+entry(%instance : $*NonTrivialStruct):
+  %field_ptr = struct_element_addr %instance : $*NonTrivialStruct, #NonTrivialStruct.val
+  %addr = alloc_stack [lexical] $Klass
+  copy_addr %field_ptr to [initialization] %addr : $*Klass
+  destroy_addr %addr : $*Klass
+  dealloc_stack %addr : $*Klass
+  %retval = tuple ()
+  return %retval : $()
+}