[AllocBoxToStack] Lexical borrowees get lexical alloc_stacks.

Previously, whenever an alloc_box that was promoted to an alloc_stack,
the new alloc_stack would be lexical.  The result was that alloc_boxes
which didn't need (or explicitly didn't want, in the case of eager move
vars) received lexical alloc_stacks.

Here, only add the lexical flag to the new alloc_stack instruction if
any of the box's uses are `begin_borrow [lexical]`.  That way,
alloc_boxes end up having lexical alloc_stacks only if they were
"lexical alloc_boxes".

Author: nate-chandler

lib/SILOptimizer/Transforms/AllocBoxToStack.cpp

@@ -13,12 +13,13 @@
 #define DEBUG_TYPE "allocbox-to-stack"
 #include "swift/AST/DiagnosticsSIL.h"
 #include "swift/Basic/BlotMapVector.h"
+#include "swift/Basic/GraphNodeWorklist.h"
 #include "swift/SIL/ApplySite.h"
+#include "swift/SIL/BasicBlockDatastructures.h"
 #include "swift/SIL/Dominance.h"
 #include "swift/SIL/SILArgument.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILCloner.h"
-#include "swift/SIL/BasicBlockDatastructures.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
@@ -27,6 +28,7 @@
 #include "swift/SILOptimizer/Utils/StackNesting.h"
 #include "swift/SILOptimizer/Utils/ValueLifetime.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
@@ -538,13 +540,35 @@ static bool rewriteAllocBoxAsAllocStack(AllocBoxInst *ABI) {
   SILBuilderWithScope Builder(ABI);
   assert(ABI->getBoxType()->getLayout()->getFields().size() == 1
          && "rewriting multi-field box not implemented");
-  auto &mod = ABI->getFunction()->getModule();
-  bool isLexical = mod.getASTContext().SILOpts.supportsLexicalLifetimes(mod);
-  auto *ASI = Builder.createAllocStack(
-      ABI->getLoc(),
-      getSILBoxFieldType(TypeExpansionContext(*ABI->getFunction()),
-                         ABI->getBoxType(), ABI->getModule().Types, 0),
-      ABI->getVarInfo(), ABI->hasDynamicLifetime(), isLexical);
+  auto ty = getSILBoxFieldType(TypeExpansionContext(*ABI->getFunction()),
+                               ABI->getBoxType(), ABI->getModule().Types, 0);
+  auto isLexical = [&]() -> bool {
+    auto &mod = ABI->getFunction()->getModule();
+    bool lexicalLifetimesEnabled =
+        mod.getASTContext().SILOpts.supportsLexicalLifetimes(mod);
+    if (!lexicalLifetimesEnabled)
+      return false;
+    // Look for lexical borrows of the alloc_box.
+    GraphNodeWorklist<Operand *, 4> worklist;
+    worklist.initializeRange(ABI->getUses());
+    while (auto *use = worklist.pop()) {
+      // See through mark_uninitialized and non-lexical begin_borrow
+      // instructions.  It's verified that lexical begin_borrows of SILBoxType
+      // values originate either from AllocBoxInsts or SILFunctionArguments.
+      if (auto *mui = dyn_cast<MarkUninitializedInst>(use->getUser())) {
+        for (auto *use : mui->getUses())
+          worklist.insert(use);
+      } else if (auto *bbi = dyn_cast<BeginBorrowInst>(use->getUser())) {
+        if (bbi->isLexical())
+          return true;
+        for (auto *use : bbi->getUses())
+          worklist.insert(use);
+      }
+    }
+    return false;
+  };
+  auto *ASI = Builder.createAllocStack(ABI->getLoc(), ty, ABI->getVarInfo(),
+                                       ABI->hasDynamicLifetime(), isLexical());
 
   // Transfer a mark_uninitialized if we have one.
   SILValue StackBox = ASI;

test/SILOptimizer/allocbox_to_stack_lifetime.sil

@@ -14,8 +14,13 @@ struct Bool {
 
 protocol Error {}
 
-// CHECK-LABEL: sil [ossa] @simple_promotion
-sil [ossa] @simple_promotion : $@convention(thin) (Int) -> Int {
+// CHECK-LABEL: sil [ossa] @promote_nonlexical
+// CHECK:         alloc_stack $
+// CHECK-NOT:     alloc_box
+// CHECK-NOT:     destroy_value
+// CHECK:         return
+// CHECK-LABEL: } // end sil function 'promote_nonlexical'
+sil [ossa] @promote_nonlexical : $@convention(thin) (Int) -> Int {
 bb0(%0 : $Int):
   %1 = alloc_box ${ var Int }
   %1a = project_box %1 : ${ var Int }, 0
@@ -24,9 +29,24 @@ bb0(%0 : $Int):
   %3 = load [trivial] %1a : $*Int
   destroy_value %1 : ${ var Int }
   return %3 : $Int
-// CHECK: alloc_stack [lexical]
-// CHECK-NOT: alloc_box
-// CHECK-NOT: destroy_value
-// CHECK: return
+}
+
+// CHECK-LABEL: sil [ossa] @promote_lexical
+// CHECK:         alloc_stack [lexical]
+// CHECK-NOT:     alloc_box
+// CHECK-NOT:     destroy_value
+// CHECK:         return
+// CHECK-LABEL: } // end sil function 'promote_lexical'
+sil [ossa] @promote_lexical : $@convention(thin) (Int) -> Int {
+bb0(%0 : $Int):
+  %1 = alloc_box ${ var Int }
+  %b = begin_borrow [lexical] %1 : ${ var Int }
+  %1a = project_box %b : ${ var Int }, 0
+  store %0 to [trivial] %1a : $*Int
+
+  %3 = load [trivial] %1a : $*Int
+  end_borrow %b : ${ var Int }
+  destroy_value %1 : ${ var Int }
+  return %3 : $Int
 }