[NFC] OwnedValueCan: Add support for end markers.

Enhance the utility with the ability to end lifetimes of lexical values
at indicated instructions, overriding the usual behavior of maintaining
such lifetimes' previous endpoints (modulo non-deinit-barrier
instructions).

Author: nate-chandler

include/swift/SILOptimizer/Utils/CanonicalizeOSSALifetime.h

@@ -256,6 +256,9 @@ class CanonicalizeOSSALifetime final {
   /// The SILValue to canonicalize.
   SILValue currentDef;
 
+  /// Instructions beyond which liveness is not extended by destroy uses.
+  ArrayRef<SILInstruction *> currentLexicalLifetimeEnds;
+
   /// Original points in the CFG where the current value's lifetime is consumed
   /// or destroyed.  Each block either contains a consuming instruction (e.g.
   /// `destroy_value`) or is on the availability boundary of the value in a
@@ -316,9 +319,10 @@ class CanonicalizeOSSALifetime final {
   struct LivenessState {
     BitfieldRef<SSAPrunedLiveness>::StackState state;
 
-    LivenessState(CanonicalizeOSSALifetime &parent, SILValue def)
+    LivenessState(CanonicalizeOSSALifetime &parent, SILValue def,
+                  ArrayRef<SILInstruction *> lexicalLifetimeEnds)
         : state(parent.liveness, def->getFunction()) {
-      parent.initializeLiveness(def);
+      parent.initializeLiveness(def, lexicalLifetimeEnds);
     }
   };
 
@@ -335,7 +339,8 @@ class CanonicalizeOSSALifetime final {
 
   SILValue getCurrentDef() const { return currentDef; }
 
-  void initializeLiveness(SILValue def) {
+  void initializeLiveness(SILValue def,
+                          ArrayRef<SILInstruction *> lexicalLifetimeEnds) {
     assert(consumingBlocks.empty() && debugValues.empty());
     // Clear the cached analysis pointer just in case the client invalidates the
     // analysis, freeing its memory.
@@ -344,6 +349,7 @@ class CanonicalizeOSSALifetime final {
     destroys.clear();
 
     currentDef = def;
+    currentLexicalLifetimeEnds = lexicalLifetimeEnds;
 
     if (maximizeLifetime || respectsDeinitBarriers()) {
       liveness->initializeDiscoveredBlocks(&discoveredBlocks);
@@ -358,8 +364,18 @@ class CanonicalizeOSSALifetime final {
   }
 
   /// Top-Level API: rewrites copies and destroys within \p def's extended
-  /// lifetime. \p lifetime caches transient analysis state across multiple
-  /// calls.
+  /// lifetime.
+  ///
+  /// For lexical values, canonicalization respects deinit barriers, introducing
+  /// copies as needed to maintain lifetimes beyond final consuming uses to the
+  /// original lexical lifetime end.  When a lexical value is explicitly
+  /// consumed (via the `consume` keyword), however, the lifetime does not
+  /// extend to original destroys beyond that consume--the value must be dead
+  /// after the corresponding marker instructions (`move_value`); to support
+  /// this shortening, the marker instructions must be provided as \p
+  /// lexicalLifetimeEnds. When provided, deinit barriers will be respected
+  /// except to the extent doing so would result in the value being live after
+  /// the marker instructions.
   ///
   /// Return true if any change was made to \p def's extended lifetime. \p def
   /// itself will not be deleted and no instructions outside of \p def's
@@ -369,7 +385,8 @@ class CanonicalizeOSSALifetime final {
   /// This only deletes instructions within \p def's extended lifetime. Use
   /// InstructionDeleter::cleanUpDeadInstructions() to recursively delete dead
   /// operands.
-  bool canonicalizeValueLifetime(SILValue def);
+  bool canonicalizeValueLifetime(
+      SILValue def, ArrayRef<SILInstruction *> lexicalLifetimeEnds = {});
 
   /// Compute the liveness information for \p def. But do not do any rewriting
   /// or computation of boundaries.

lib/SILOptimizer/Mandatory/MoveOnlyObjectCheckerUtils.h

@@ -64,7 +64,7 @@ struct OSSACanonicalizer {
     CanonicalizeOSSALifetime::LivenessState canonicalizerState;
 
     LivenessState(OSSACanonicalizer &parent, SILValue def)
-        : parent(parent), canonicalizerState(parent.canonicalizer, def) {}
+        : parent(parent), canonicalizerState(parent.canonicalizer, def, {}) {}
 
     ~LivenessState() { parent.clear(); }
   };

lib/SILOptimizer/Utils/CanonicalizeOSSALifetime.cpp

@@ -278,14 +278,29 @@ void CanonicalizeOSSALifetime::extendLivenessToDeinitBarriers() {
         });
       });
 
+  ArrayRef<SILInstruction *> ends = {};
+  SmallVector<SILInstruction *, 8> lexicalEnds;
+  if (currentLexicalLifetimeEnds.size() > 0) {
+    visitExtendedUnconsumedBoundary(
+        currentLexicalLifetimeEnds,
+        [&lexicalEnds](auto *instruction, auto lifetimeEnding) {
+          instruction->visitSubsequentInstructions([&](auto *next) {
+            lexicalEnds.push_back(next);
+            return true;
+          });
+        });
+    ends = lexicalEnds;
+  } else {
+    ends = outsideDestroys;
+  }
+
   auto *def = getCurrentDef()->getDefiningInstruction();
   using InitialBlocks = ArrayRef<SILBasicBlock *>;
   auto *defBlock = getCurrentDef()->getParentBlock();
   auto initialBlocks = defBlock ? InitialBlocks(defBlock) : InitialBlocks();
   ReachableBarriers barriers;
-  findBarriersBackward(outsideDestroys, initialBlocks,
-                       *getCurrentDef()->getFunction(), barriers,
-                       [&](auto *inst) {
+  findBarriersBackward(ends, initialBlocks, *getCurrentDef()->getFunction(),
+                       barriers, [&](auto *inst) {
                          if (inst == def)
                            return true;
                          if (!isDeinitBarrier(inst, calleeAnalysis))
@@ -569,11 +584,17 @@ void CanonicalizeOSSALifetime::findOriginalBoundary(
 ///   extent.
 ///   [Extend liveness down to the boundary between green blocks and uncolored.]
 void CanonicalizeOSSALifetime::visitExtendedUnconsumedBoundary(
-    ArrayRef<SILInstruction *> ends,
+    ArrayRef<SILInstruction *> consumes,
     llvm::function_ref<void(SILInstruction *, PrunedLiveness::LifetimeEnding)>
         visitor) {
   auto currentDef = getCurrentDef();
 
+#ifndef NDEBUG
+  for (auto *consume : consumes) {
+    assert(!liveness->isWithinBoundary(consume));
+  }
+#endif
+
   // First, collect the blocks that were _originally_ live.  We can't use
   // liveness here because it doesn't include blocks that occur before a
   // destroy_value.
@@ -619,7 +640,7 @@ void CanonicalizeOSSALifetime::visitExtendedUnconsumedBoundary(
     // consumes. These are just the instructions on the boundary which aren't
     // destroys.
     BasicBlockWorklist worklist(currentDef->getFunction());
-    for (auto *instruction : ends) {
+    for (auto *instruction : consumes) {
       if (destroys.contains(instruction))
         continue;
       if (liveness->isInterestingUser(instruction)
@@ -1223,8 +1244,9 @@ void CanonicalizeOSSALifetime::rewriteLifetimes() {
 }
 
 /// Canonicalize a single extended owned lifetime.
-bool CanonicalizeOSSALifetime::canonicalizeValueLifetime(SILValue def) {
-  LivenessState livenessState(*this, def);
+bool CanonicalizeOSSALifetime::canonicalizeValueLifetime(
+    SILValue def, ArrayRef<SILInstruction *> lexicalLifetimeEnds) {
+  LivenessState livenessState(*this, def, lexicalLifetimeEnds);
 
   // Don't canonicalize the lifetimes of values of move-only type.  According to
   // language rules, they are fixed.
@@ -1252,6 +1274,7 @@ namespace swift::test {
 //         access scopes which they previously enclosed but can't be hoisted
 //         before
 // - SILValue: value to canonicalize
+// - [SILInstruction]: the lexicalLifetimeEnds to recognize
 // Dumps:
 // - function after value canonicalization
 static FunctionTest CanonicalizeOSSALifetimeTest(
@@ -1271,7 +1294,11 @@ static FunctionTest CanonicalizeOSSALifetimeTest(
           respectAccessScopes ? accessBlockAnalysis : nullptr, domTree,
           calleeAnalysis, deleter);
       auto value = arguments.takeValue();
-      canonicalizer.canonicalizeValueLifetime(value);
+      SmallVector<SILInstruction *, 4> lexicalLifetimeEnds;
+      while (arguments.hasUntaken()) {
+        lexicalLifetimeEnds.push_back(arguments.takeInstruction());
+      }
+      canonicalizer.canonicalizeValueLifetime(value, lexicalLifetimeEnds);
       function.print(llvm::outs());
     });
 } // end namespace swift::test