Merge pull request swiftlang#39126 from atrick/verify-ossa-term

Enforce OSSA block terminator invariants

Author: atrick

include/swift/SIL/SILBuilder.h

@@ -2221,6 +2221,11 @@ class SILBuilder {
   BranchInst *createBranch(SILLocation Loc, SILBasicBlock *TargetBlock,
                            OperandValueArrayRef Args);
 
+  // This only creates the terminator, not the results. Create the results with
+  // OwnershipForwardingTermInst::createResult() and
+  // SwitchEnumInst::createDefaultResult() to ensure that the result ownership
+  // is correct (it must be consistent with the switch_enum's forwarding
+  // ownership, which may differ from \p Operand's ownership).
   SwitchValueInst *
   createSwitchValue(SILLocation Loc, SILValue Operand, SILBasicBlock *DefaultBB,
                     ArrayRef<std::pair<SILValue, SILBasicBlock *>> CaseBBs) {
@@ -2232,16 +2237,14 @@ class SILBuilder {
       SILLocation Loc, SILValue Operand, SILBasicBlock *DefaultBB,
       ArrayRef<std::pair<EnumElementDecl *, SILBasicBlock *>> CaseBBs,
       Optional<ArrayRef<ProfileCounter>> CaseCounts = None,
-      ProfileCounter DefaultCount = ProfileCounter()) {
-    return createSwitchEnum(Loc, Operand, DefaultBB, CaseBBs, CaseCounts,
-                            DefaultCount, Operand.getOwnershipKind());
-  }
+      ProfileCounter DefaultCount = ProfileCounter());
 
   SwitchEnumInst *createSwitchEnum(
       SILLocation Loc, SILValue Operand, SILBasicBlock *DefaultBB,
       ArrayRef<std::pair<EnumElementDecl *, SILBasicBlock *>> CaseBBs,
       Optional<ArrayRef<ProfileCounter>> CaseCounts,
-      ProfileCounter DefaultCount, ValueOwnershipKind forwardingOwnershipKind) {
+      ProfileCounter DefaultCount,
+      ValueOwnershipKind forwardingOwnershipKind) {
     return insertTerminator(SwitchEnumInst::create(
         getSILDebugLocation(Loc), Operand, DefaultBB, CaseBBs, getFunction(),
         CaseCounts, DefaultCount, forwardingOwnershipKind));

include/swift/SIL/SILInstruction.h

@@ -1056,6 +1056,7 @@ class OwnershipForwardingMixin {
                            ValueOwnershipKind ownershipKind)
       : ownershipKind(ownershipKind) {
     assert(isa(kind) && "Invalid subclass?!");
+    assert(ownershipKind && "invalid forwarding ownership");
   }
 
 public:
@@ -7877,6 +7878,7 @@ class TermInst : public NonValueInstruction {
   }
 };
 
+// Forwards the first operand to a result in each successor block.
 class OwnershipForwardingTermInst : public TermInst,
                                     public OwnershipForwardingMixin {
 protected:
@@ -7903,6 +7905,11 @@ class OwnershipForwardingTermInst : public TermInst,
     return kind == SILInstructionKind::SwitchEnumInst ||
            kind == SILInstructionKind::CheckedCastBranchInst;
   }
+
+  SILValue getOperand() const { return getAllOperands()[0].get(); }
+
+  /// Create a result for this terminator on the given successor block.
+  SILPhiArgument *createResult(SILBasicBlock *succ, SILType resultTy);
 };
 
 /// UnreachableInst - Position in the code which would be undefined to reach.
@@ -8642,6 +8649,14 @@ class SwitchEnumInst
          SILFunction &F, Optional<ArrayRef<ProfileCounter>> CaseCounts,
          ProfileCounter DefaultCount,
          ValueOwnershipKind forwardingOwnershipKind);
+
+public:
+  /// Create the default result for a partially built switch_enum.
+  /// Returns nullptr if no default argument is needed.
+  SILPhiArgument *createDefaultResult();
+
+  /// Create the .some result for an optional switch_enum.
+  SILPhiArgument *createOptionalSomeResult();
 };
 /// A switch on an enum's discriminator in memory.
 class SwitchEnumAddrInst

include/swift/SIL/SILInstructionWorklist.h

@@ -185,7 +185,10 @@ class SILInstructionWorklist : SILInstructionWorklistBase {
   /// Intended to be called during visitation after \p instruction has been
   /// removed from the worklist.
   ///
-  /// \p instruction the instruction whose usages will be replaced
+  /// \p instruction the instruction whose usages will be replaced. This
+  /// instruction may already be deleted to maintain valid SIL. For example, if
+  /// it was a block terminator.
+  ///
   /// \p result the instruction whose usages will replace \p instruction
   ///
   /// \return whether the instruction was deleted or modified.

lib/SIL/IR/SILBuilder.cpp

@@ -601,13 +601,74 @@ void SILBuilder::emitScopedBorrowOperation(SILLocation loc, SILValue original,
     createEndBorrow(loc, value);
 }
 
+/// Attempt to propagate ownership from \p operand to the returned forwarding
+/// ownership where the forwarded value has type \p targetType. If this fails,
+/// return Owned forwarding ownership instead.
+///
+/// Propagation only fails when \p operand is dynamically trivial, as indicated
+/// by ownership None, AND \p targetType is statically nontrivial.
+///
+/// Example:
+///
+///     %e = enum $Optional<AnyObject>, #Optional.none!enumelt
+///     switch_enum %e : $Optional<AnyObject>,
+///                      case #Optional.some!enumelt: bb2...,
+///                      forwarding: @owned
+///   bb2(%arg : @owned AnyObject):
+///
+/// Example:
+///
+///     %mt = metatype $@thick C.Type
+///     checked_cast_br %mt : $@thick C.Type to AnyObject.Type, bb1, bb2,
+///                           forwarding: @owned
+///   bb1(%arg : @owned AnyObject.Type):
+///
+/// If the forwarded value is statically known nontrivial, then the forwarding
+/// ownership cannot be None. Such a result is unreachable, but the SIL on that
+/// path must still be valid. When creating ownership out of thin air, default
+/// to Owned because that allows the value to be consumed without generating a
+/// copy. This does require the client code to handle ending the lifetime of an
+/// owned result even if the input was passed as guaranteed.
+///
+/// Note: For simplicitly, ownership None is not propagated for any statically
+/// nontrivial result, even if \p targetType may also be dynamically
+/// trivial. For example, the operand of a switch_enum could be a nested enum
+/// such that all switch cases may be dynamically trivial. Or a checked_cast_br
+/// could cast from one dynamically trivial enum to another. Figuring out
+/// whether the dynamically trivial operand value maps onto a dynamically
+/// trivial terminator result would be very complex with no practical benefit.
+static ValueOwnershipKind deriveForwardingOwnership(SILValue operand,
+                                                    SILType targetType,
+                                                    SILFunction &func) {
+  if (operand.getOwnershipKind() != OwnershipKind::None
+      || targetType.isTrivial(func)) {
+    return operand.getOwnershipKind();
+  }
+  return OwnershipKind::Owned;
+}
+
+SwitchEnumInst *SILBuilder::createSwitchEnum(
+    SILLocation Loc, SILValue Operand, SILBasicBlock *DefaultBB,
+    ArrayRef<std::pair<EnumElementDecl *, SILBasicBlock *>> CaseBBs,
+    Optional<ArrayRef<ProfileCounter>> CaseCounts,
+    ProfileCounter DefaultCount) {
+  // Consider the operand's type to be the target's type since a switch
+  // covers all cases including the default argument.
+  auto forwardingOwnership =
+      deriveForwardingOwnership(Operand, Operand->getType(), getFunction());
+  return createSwitchEnum(Loc, Operand, DefaultBB, CaseBBs, CaseCounts,
+                          DefaultCount, forwardingOwnership);
+}
+
 CheckedCastBranchInst *SILBuilder::createCheckedCastBranch(
     SILLocation Loc, bool isExact, SILValue op,
     SILType destLoweredTy, CanType destFormalTy,
     SILBasicBlock *successBB, SILBasicBlock *failureBB,
     ProfileCounter target1Count, ProfileCounter target2Count) {
+  auto forwardingOwnership =
+      deriveForwardingOwnership(op, destLoweredTy, getFunction());
   return createCheckedCastBranch(Loc, isExact, op, destLoweredTy, destFormalTy,
-                                 successBB, failureBB, op.getOwnershipKind(),
+                                 successBB, failureBB, forwardingOwnership,
                                  target1Count, target2Count);
 }
 
@@ -619,10 +680,11 @@ CheckedCastBranchInst *SILBuilder::createCheckedCastBranch(
   assert((!hasOwnership() || !failureBB->getNumArguments() ||
           failureBB->getArgument(0)->getType() == op->getType()) &&
          "failureBB's argument doesn't match incoming argument type");
+
   return insertTerminator(CheckedCastBranchInst::create(
       getSILDebugLocation(Loc), isExact, op, destLoweredTy, destFormalTy,
-      successBB, failureBB, getFunction(), target1Count,
-      target2Count, forwardingOwnershipKind));
+      successBB, failureBB, getFunction(), target1Count, target2Count,
+      forwardingOwnershipKind));
 }
 
 void SILBuilderWithScope::insertAfter(SILInstruction *inst,

lib/SIL/IR/SILInstructions.cpp

@@ -2895,3 +2895,68 @@ ReturnInst::ReturnInst(SILFunction &func, SILDebugLocation debugLoc,
          "Conflicting ownership kinds when creating term inst from function "
          "result info?!");
 }
+
+// This may be called in an invalid SIL state. SILCombine creates new
+// terminators in non-terminator position and defers deleting the original
+// terminator until after all modification.
+SILPhiArgument *OwnershipForwardingTermInst::createResult(SILBasicBlock *succ,
+                                                          SILType resultTy) {
+  // The forwarding instruction declares a forwarding ownership kind that
+  // determines the ownership of its results.
+  auto resultOwnership = getForwardingOwnershipKind();
+
+  // Trivial results have no ownership. Although it is valid for a trivially
+  // typed value to have ownership, it is never necessary and less efficient.
+  if (resultTy.isTrivial(*getFunction())) {
+    resultOwnership = OwnershipKind::None;
+
+  } else if (resultOwnership == OwnershipKind::None) {
+    // switch_enum strangely allows results to acquire ownership out of thin
+    // air whenever the operand has no ownership and result is nontrivial:
+    //     %e = enum $Optional<AnyObject>, #Optional.none!enumelt
+    //     switch_enum %e : $Optional<AnyObject>,
+    //                 case #Optional.some!enumelt: bb2...
+    //   bb2(%arg : @guaranteed T):
+    //
+    // We can either use None or Guaranteed. None would correctly propagate
+    // ownership and would maintain the invariant that guaranteed values are
+    // always within a borrow scope. However it would result in a nontrivial
+    // type without ownership. The lifetime verifier does not like that.
+    resultOwnership = OwnershipKind::Guaranteed;
+  }
+  return succ->createPhiArgument(resultTy, resultOwnership);
+}
+
+SILPhiArgument *SwitchEnumInst::createDefaultResult() {
+  auto *f = getFunction();
+  if (!f->hasOwnership())
+    return nullptr;
+
+  if (!hasDefault())
+    return nullptr;
+
+  assert(getDefaultBB()->getNumArguments() == 0 && "precondition");
+
+  auto enumTy = getOperand()->getType();
+  NullablePtr<EnumElementDecl> uniqueCase = getUniqueCaseForDefault();
+
+  // Without a unique default case, the OSSA result simply forwards the
+  // switch_enum operand.
+  if (!uniqueCase)
+    return createResult(getDefaultBB(), enumTy);
+
+  // With a unique default case, the result is materialized exactly the same way
+  // as a matched result. It has a value iff the unique case has a payload.
+  if (!uniqueCase.get()->hasAssociatedValues())
+    return nullptr;
+
+  auto resultTy = enumTy.getEnumElementType(uniqueCase.get(), f->getModule(),
+                                            f->getTypeExpansionContext());
+  return createResult(getDefaultBB(), resultTy);
+}
+
+SILPhiArgument *SwitchEnumInst::createOptionalSomeResult() {
+  auto someDecl = getModule().getASTContext().getOptionalSomeDecl();
+  auto someBB = getCaseDestination(someDecl);
+  return createResult(someBB, getOperand()->getType().unwrapOptionalType());
+}