[Mem2Reg] Omit lexical moves for lexical values.

lib/SILOptimizer/Transforms/SILMem2Reg.cpp

@@ -57,6 +57,7 @@ STATISTIC(NumInstRemoved,        "Number of Instructions removed");
 
 static bool lexicalLifetimeEnsured(AllocStackInst *asi);
 static bool isGuaranteedLexicalValue(SILValue src);
+static bool isOwnedLexicalValue(SILValue src);
 
 namespace {
 
@@ -96,9 +97,11 @@ class LiveValues {
       if (!lexicalLifetimeEnsured(asi)) {
         return stored;
       }
-      // We should have created a move of the @owned stored value.
-      assert(move);
-      return move;
+      auto storedIsLexical = stored && isOwnedLexicalValue(stored);
+      // If the value was already lexical, we use it directly.  Otherwise, a new
+      // move_value [lexical] is used.
+      assert(storedIsLexical || move);
+      return storedIsLexical ? stored : move;
     }
 
     bool canEndLexicalLifetime() {
@@ -107,7 +110,8 @@ class LiveValues {
       // to end a lexical lifetime. In that case, the lifetime end will be
       // added later, when we have enough information, namely the live in
       // values, to end it.
-      return move;
+      auto storedIsLexical = stored && isOwnedLexicalValue(stored);
+      return storedIsLexical ? stored : move;
     }
   };
   struct Guaranteed {
@@ -221,6 +225,7 @@ class LiveValues {
     return guaranteed.stored;
   }
 
+  /// Whether it's possible and appropriate to end the lifetime.
   bool canEndLexicalLifetime() {
     if (auto *owned = storage.dyn_cast<Owned>()) {
       return owned->canEndLexicalLifetime();
@@ -521,16 +526,15 @@ static bool lexicalLifetimeEnsured(AllocStackInst *asi) {
          !asi->getElementType().isTrivial(*asi->getFunction());
 }
 
+static bool isOwnedLexicalValue(SILValue src) {
+  return src->getOwnershipKind() == OwnershipKind::Owned && src->isLexical();
+}
+
 static bool isGuaranteedLexicalValue(SILValue src) {
   return src->getOwnershipKind() == OwnershipKind::Guaranteed &&
          src->isLexical();
 }
 
-/// Returns true if we have enough information to end the lifetime.
-static bool canEndLexicalLifetime(LiveValues values) {
-  return values.canEndLexicalLifetime();
-}
-
 /// Begin a lexical borrow scope for the value stored into the provided
 /// StoreInst after that instruction.
 ///
@@ -546,6 +550,9 @@ beginOwnedLexicalLifetimeAfterStore(AllocStackInst *asi, StoreInst *inst) {
   SILValue stored = inst->getOperand(CopyLikeInstruction::Src);
   SILLocation loc = RegularLocation::getAutoGeneratedLocation(inst->getLoc());
 
+  if (isOwnedLexicalValue(stored)) {
+    return {LiveValues::forOwned(stored, {}), /*isStorageValid*/ true};
+  }
   MoveValueInst *mvi = nullptr;
   SILBuilderWithScope::insertAfter(inst, [&](SILBuilder &builder) {
     mvi = builder.createMoveValue(loc, stored, /*isLexical*/ true);
@@ -825,7 +832,7 @@ SILInstruction *StackAllocationPromoter::promoteAllocationInBlock(
         if (lexicalLifetimeEnsured(asi)) {
           // End the lexical lifetime at a load [take].  The storage is no
           // longer keeping the value alive.
-          if (runningVals && canEndLexicalLifetime(runningVals->value)) {
+          if (runningVals && runningVals->value.canEndLexicalLifetime()) {
             // End it right now if we have enough information.
             endOwnedLexicalLifetimeBeforeInst(asi, /*beforeInstruction=*/li,
                                               ctx,
@@ -908,7 +915,7 @@ SILInstruction *StackAllocationPromoter::promoteAllocationInBlock(
       lastStoreInst = si;
       if (lexicalLifetimeEnsured(asi)) {
         if (oldRunningVals && oldRunningVals->isStorageValid &&
-            canEndLexicalLifetime(oldRunningVals->value)) {
+            oldRunningVals->value.canEndLexicalLifetime()) {
           endOwnedLexicalLifetimeBeforeInst(asi, /*beforeInstruction=*/si, ctx,
                                             oldRunningVals->value.getOwned());
         }
@@ -965,7 +972,7 @@ SILInstruction *StackAllocationPromoter::promoteAllocationInBlock(
       }
       // Mark storage as invalid and mark end_borrow as a deinit point.
       runningVals->isStorageValid = false;
-      if (!canEndLexicalLifetime(runningVals->value)) {
+      if (!runningVals->value.canEndLexicalLifetime()) {
         continue;
       }
       endGuaranteedLexicalLifetimeBeforeInst(
@@ -1067,6 +1074,10 @@ StackAllocationPromoter::getLiveOutValues(BlockSetVector &phiBlocks,
         auto values = LiveValues::forGuaranteed(stored, borrow);
         return values;
       }
+      if (isOwnedLexicalValue(stored)) {
+        auto values = LiveValues::forOwned(stored, {});
+        return values;
+      }
       auto move = cast<MoveValueInst>(inst->getNextInstruction());
       auto values = LiveValues::forOwned(stored, move);
       return values;
@@ -1422,7 +1433,7 @@ void StackAllocationPromoter::endLexicalLifetime(BlockSetVector &phiBlocks) {
         if (isa<EndBorrowInst>(inst)) {
           // Not all store_borrows will have a begin_borrow [lexical] that needs
           // to be ended. If the source is already lexical, we don't create it.
-          if (!canEndLexicalLifetime(*values)) {
+          if (!values->canEndLexicalLifetime()) {
             continue;
           }
           endGuaranteedLexicalLifetimeBeforeInst(
@@ -1445,7 +1456,7 @@ void StackAllocationPromoter::endLexicalLifetime(BlockSetVector &phiBlocks) {
       if (terminatesInUnreachable || uniqueSuccessorLacksLiveInValues()) {
         auto values = getLiveOutValues(phiBlocks, bb);
         if (values->isGuaranteed()) {
-          if (!canEndLexicalLifetime(*values)) {
+          if (!values->canEndLexicalLifetime()) {
             continue;
           }
           endGuaranteedLexicalLifetimeBeforeInst(
@@ -1965,7 +1976,7 @@ void MemoryToRegisters::removeSingleBlockAllocation(AllocStackInst *asi) {
         continue;
       }
       runningVals->isStorageValid = false;
-      if (!canEndLexicalLifetime(runningVals->value)) {
+      if (!runningVals->value.canEndLexicalLifetime()) {
         continue;
       }
       endGuaranteedLexicalLifetimeBeforeInst(