BorrowToDestructureUtils: handle borrowing switches.

jckarter · jckarter · commit 03e2e8f476e9 · 2024-01-31T20:38:22.000-08:00
Instead of trying to convert all switches into consumes, allow a switch of a borrow
to remain as such.
diff --git a/lib/SIL/Utils/FieldSensitivePrunedLiveness.cpp b/lib/SIL/Utils/FieldSensitivePrunedLiveness.cpp
@@ -304,11 +304,23 @@ SubElementOffset::computeForValue(SILValue projectionDerivedFromRoot,
     // So our payload is always going to start at the current field number since
     // we are the left most child of our parent enum. So we just need to look
     // through to our parent enum.
+    //
+    // Enum projections can happen either directly via an unchecked instruction…
     if (auto *enumData =
             dyn_cast<UncheckedEnumDataInst>(projectionDerivedFromRoot)) {
       projectionDerivedFromRoot = enumData->getOperand();
       continue;
     }
+    
+    // …or via the bb arg of a `switch_enum` successor.
+    if (auto bbArg = dyn_cast<SILArgument>(projectionDerivedFromRoot)) {
+      if (auto pred = bbArg->getParent()->getSinglePredecessorBlock()) {
+        if (auto switchEnum = dyn_cast<SwitchEnumInst>(pred->getTerminator())) {
+          projectionDerivedFromRoot = switchEnum->getOperand();
+          continue;
+        }
+      }
+    }
 
     // If we do not know how to handle this case, just return None.
     //
diff --git a/lib/SILGen/SILGenPattern.cpp b/lib/SILGen/SILGenPattern.cpp
@@ -523,7 +523,8 @@ class PatternMatchEmission {
                              const FailureHandler &failure);
 
   void emitGuardBranch(SILLocation loc, Expr *guard,
-                       const FailureHandler &failure);
+                       const FailureHandler &failure,
+                       const ClauseRow &row, ArgArray args);
 
   // Bind copyable variable bindings as independent variables.
   void bindIrrefutablePatterns(const ClauseRow &row, ArgArray args,
@@ -1183,25 +1184,21 @@ void PatternMatchEmission::emitWildcardDispatch(ClauseMatrix &clauses,
 
   if (auto ownership = getNoncopyableOwnership()) {
     // A noncopyable pattern match always happens over a borrow first.
-    // If there is a guard expression, then we also have to match pattern
-    // variables as borrows while executing the guard. This ensures
-    // the guard can't consume or modify the value without us committing to this
-    // case branch. If the final pattern match is only borrowing as well,
+    // If the final pattern match is only borrowing as well,
     // we can bind the variables immediately here too.
-    if (hasGuard || *ownership <= ValueOwnership::Shared) {
+    if (*ownership <= ValueOwnership::Shared) {
       bindIrrefutableBorrows(clauses[row], args);
     }
     
     if (hasGuard) {
-      this->emitGuardBranch(guardExpr, guardExpr, failure);
+      // The guard will bind borrows locally if necessary.
+      this->emitGuardBranch(guardExpr, guardExpr, failure,
+                            clauses[row], args);
     }
     
     if (*ownership > ValueOwnership::Shared) {
-      // Unbind the variables and unborrow the subject so we can consume it
-      // later.
       unbindAndEndBorrows(clauses[row], args);
     }
-    
   } else {
     // Bind the rest of the patterns.
     // For noncopyable bindings, this will bind them as borrows initially if there
@@ -1211,7 +1208,8 @@ void PatternMatchEmission::emitWildcardDispatch(ClauseMatrix &clauses,
 
     // Emit the guard branch, if it exists.
     if (guardExpr) {
-      this->emitGuardBranch(guardExpr, guardExpr, failure);
+      this->emitGuardBranch(guardExpr, guardExpr, failure,
+                            clauses[row], args);
     }
   }
 
@@ -1244,7 +1242,8 @@ bindRefutablePatterns(const ClauseRow &row, ArgArray args,
       FullExpr scope(SGF.Cleanups, CleanupLocation(pattern));
       bindVariable(pattern, exprPattern->getMatchVar(), args[i],
                    /*isForSuccess*/ false, /* hasMultipleItems */ false);
-      emitGuardBranch(pattern, exprPattern->getMatchExpr(), failure);
+      emitGuardBranch(pattern, exprPattern->getMatchExpr(), failure,
+                      row, args);
       break;
     }
     default:
@@ -1391,21 +1390,38 @@ void PatternMatchEmission::bindBorrow(Pattern *pattern, VarDecl *var,
                                       ConsumableManagedValue value) {
   assert(value.getFinalConsumption() == CastConsumptionKind::BorrowAlways);
   
-  SGF.VarLocs[var] = SILGenFunction::VarLoc::get(
-    value.asBorrowedOperand2(SGF, pattern).getValue());
+  auto bindValue = value.asBorrowedOperand2(SGF, pattern).getFinalManagedValue();
+  if (bindValue.getType().isObject()) {
+    // Create a notional copy for the borrow checker to use.
+    bindValue = bindValue.copy(SGF, pattern);
+  }
+  bindValue = SGF.B.createMarkUnresolvedNonCopyableValueInst(pattern, bindValue,
+              MarkUnresolvedNonCopyableValueInst::CheckKind::NoConsumeOrAssign);
+  
+  SGF.VarLocs[var] = SILGenFunction::VarLoc::get(bindValue.getValue());
 }
 
 /// Evaluate a guard expression and, if it returns false, branch to
 /// the given destination.
 void PatternMatchEmission::emitGuardBranch(SILLocation loc, Expr *guard,
-                                           const FailureHandler &failure) {
+                                           const FailureHandler &failure,
+                                           const ClauseRow &row, ArgArray args){
   SILBasicBlock *falseBB = SGF.B.splitBlockForFallthrough();
   SILBasicBlock *trueBB = SGF.B.splitBlockForFallthrough();
 
   // Emit the match test.
   SILValue testBool;
   {
     FullExpr scope(SGF.Cleanups, CleanupLocation(guard));
+    
+    // If the final pattern match is destructive, then set up borrow bindings
+    // to evaluate the guard expression without allowing it to destruct the
+    // subject yet.
+    if (auto ownership = getNoncopyableOwnership()) {
+      if (*ownership > ValueOwnership::Shared) {
+        bindIrrefutableBorrows(row, args);
+      }
+    }
     testBool = SGF.emitRValueAsSingleValue(guard).getUnmanagedValue();
   }
 
@@ -2720,33 +2736,28 @@ void PatternMatchEmission::emitDestructiveCaseBlocks() {
                                       CleanupState::PersistentlyActive);
     }
     
-    {
-      // Create a scope to break down the subject value.
-      Scope caseScope(SGF, pattern);
-      
-      // Clone the original subject's cleanup state so that it will be reliably
-      // consumed in this scope, while leaving the original for other case
-      // blocks to re-consume.
-      ManagedValue subject = SGF.emitManagedRValueWithCleanup(
-                                       NoncopyableConsumableValue.forward(SGF));
-
-      // TODO: handle fallthroughs and multiple cases bindings
-      // In those cases we'd need to forward bindings through the shared case
-      // destination blocks.
-      assert(!stmt->hasFallthroughDest()
-             && stmt->getCaseLabelItems().size() == 1);
-
-      // Bind variables from the pattern.
-      if (stmt->hasCaseBodyVariables()) {
-        ConsumingPatternBindingVisitor(*this, stmt)
-          .visit(pattern, subject);
-      }
-      
-      // By this point, whatever parts of the value we bound are forwarded into
-      // variables, so we can pop the scope and destroy the remainder before
-      // entering the case body (TODO: or common block).
-    }
+    // Create a scope to break down the subject value.
+    Scope caseScope(SGF, pattern);
     
+    // Clone the original subject's cleanup state so that it will be reliably
+    // consumed in this scope, while leaving the original for other case
+    // blocks to re-consume.
+    ManagedValue subject = SGF.emitManagedRValueWithCleanup(
+                                     NoncopyableConsumableValue.forward(SGF));
+
+    // TODO: handle fallthroughs and multiple cases bindings
+    // In those cases we'd need to forward bindings through the shared case
+    // destination blocks.
+    assert(!stmt->hasFallthroughDest()
+           && stmt->getCaseLabelItems().size() == 1);
+
+    // Bind variables from the pattern.
+    if (stmt->hasCaseBodyVariables()) {
+      ConsumingPatternBindingVisitor(*this, stmt)
+        .visit(pattern, subject);
+    }
+
+    // Emit the case body.
     emitCaseBody(stmt);
   }
 }
diff --git a/lib/SILOptimizer/Mandatory/MoveOnlyBorrowToDestructureUtils.cpp b/lib/SILOptimizer/Mandatory/MoveOnlyBorrowToDestructureUtils.cpp
@@ -372,6 +372,29 @@ bool Implementation::gatherUses(SILValue value) {
     }
 
     case OperandOwnership::GuaranteedForwarding: {
+      // Always treat switches as full liveness uses of the enum being switched
+      // since the control flow is significant, and we can't destructure through
+      // the switch dispatch. If the final pattern match ends up destructuring
+      // the value, then SILGen emits that as a separate access.
+      if (auto switchEnum = dyn_cast<SwitchEnumInst>(nextUse->getUser())) {
+        auto leafRange = TypeTreeLeafTypeRange::get(switchEnum->getOperand(),
+                                                    getRootValue());
+        if (!leafRange) {
+          LLVM_DEBUG(llvm::dbgs() << "        Failed to compute leaf range?!\n");
+          return false;
+        }
+
+        LLVM_DEBUG(llvm::dbgs() << "        Found non lifetime ending use!\n");
+        blocksToUses.insert(nextUse->getParentBlock(),
+                            {nextUse,
+                             {liveness.getNumSubElements(), *leafRange,
+                              false /*is lifetime ending*/}});
+        liveness.updateForUse(nextUse->getUser(), *leafRange,
+                              false /*is lifetime ending*/);
+        instToInterestingOperandIndexMap.insert(nextUse->getUser(), nextUse);
+        continue;
+      }
+    
       // Look through guaranteed forwarding if we have at least one non-trivial
       // value. If we have all non-trivial values, treat this as a liveness use.
       SmallVector<SILValue, 8> forwardedValues;
@@ -1044,6 +1067,23 @@ static void insertEndBorrowsForNonConsumingUse(Operand *op,
                                  borrow);
       return true;
     });
+  } else if (auto swi = dyn_cast<SwitchEnumInst>(op->getUser())) {
+    LLVM_DEBUG(llvm::dbgs() << "    -- Ending borrow for switch:\n"
+                               "    ";
+               swi->print(llvm::dbgs()));
+    // End the borrow where the original borrow of the subject was ended.
+    // TODO: handle if the switch isn't directly on a borrow?
+    auto beginBorrow = cast<BeginBorrowInst>(swi->getOperand());
+    BorrowingOperand(&beginBorrow->getOperandRef())
+      .visitScopeEndingUses([&](Operand *endScope) -> bool {
+        auto *endScopeInst = endScope->getUser();
+        LLVM_DEBUG(llvm::dbgs() << "       ";
+                   endScopeInst->print(llvm::dbgs()));
+        SILBuilderWithScope endBuilder(endScopeInst);
+        endBuilder.createEndBorrow(getSafeLoc(endScopeInst),
+                                   borrow);
+        return true;
+      });
   } else {
     auto *nextInst = op->getUser()->getNextInstruction();
     LLVM_DEBUG(llvm::dbgs() << "    -- Ending borrow after momentary use at: ";
@@ -1521,9 +1561,17 @@ static bool gatherBorrows(SILValue rootValue,
 static bool
 gatherSwitchEnum(SILValue value,
                  SmallVectorImpl<SwitchEnumInst *> &switchEnumWorklist) {
+  auto *fn = value->getFunction();
+  auto &C = fn->getASTContext();
+  // TODO: For borrowing switches, we want to leave the switch as a borrowing
+  // operation. This phase of the pass should become moot once borrowing
+  // switch is on by default.
+  if (C.LangOpts.hasFeature(Feature::BorrowingSwitch)) {
+    return true;
+  }
+
   LLVM_DEBUG(llvm::dbgs() << "Gathering switch enums for value: " << *value);
 
-  auto *fn = value->getFunction();
   StackList<Operand *> useWorklist(fn);
   for (auto *use : value->getUses()) {
     useWorklist.push_back(use);
@@ -1639,8 +1687,10 @@ gatherSwitchEnum(SILValue value,
 //===----------------------------------------------------------------------===//
 
 bool BorrowToDestructureTransform::transform() {
-  LLVM_DEBUG(llvm::dbgs() << "Performing Borrow To Destructure Tranform!\n");
   auto *fn = mmci->getFunction();
+  LLVM_DEBUG(llvm::dbgs() << "Performing Borrow To Destructure Transform!\n";
+             fn->print(llvm::dbgs()));
+  auto &C = fn->getASTContext();
   StackList<BeginBorrowInst *> borrowWorklist(mmci->getFunction());
 
   // If we failed to gather borrows due to the transform not understanding part
diff --git a/lib/SILOptimizer/Mandatory/MoveOnlyObjectCheckerUtils.cpp b/lib/SILOptimizer/Mandatory/MoveOnlyObjectCheckerUtils.cpp
@@ -488,10 +488,10 @@ void MoveOnlyObjectCheckerPImpl::check(DominanceInfo *domTree,
 
           // Handle:
           //
-          // bb0(%0 : @guaranteed $Type):
-          //   %1 = copy_value %0
-          //   %2 = mark_unresolved_non_copyable_value [no_consume_or_assign] %1
-          if (auto *arg = dyn_cast<SILFunctionArgument>(i->getOperand(0))) {
+          // bb(%arg : @guaranteed $Type):
+          //   %copy = copy_value %arg
+          //   %mark = mark_unresolved_non_copyable_value [no_consume_or_assign] %copy
+          if (auto *arg = dyn_cast<SILArgument>(i->getOperand(0))) {
             if (arg->getOwnershipKind() == OwnershipKind::Guaranteed) {
               for (auto *use : markedInst->getConsumingUses()) {
                 destroys.push_back(cast<DestroyValueInst>(use->getUser()));
