[ownership] Change ReturnInst to have its ValueOwnershipKind stored within it rather than always recomputing from the function type.

include/swift/SIL/SILBuilder.h

@@ -1969,7 +1969,7 @@ class SILBuilder {
 
   ReturnInst *createReturn(SILLocation Loc, SILValue ReturnValue) {
     return insertTerminator(new (getModule()) ReturnInst(
-        getSILDebugLocation(Loc), ReturnValue));
+        getFunction(), getSILDebugLocation(Loc), ReturnValue));
   }
 
   ThrowInst *createThrow(SILLocation Loc, SILValue errorValue) {

include/swift/SIL/SILInstruction.h

@@ -7360,16 +7360,25 @@ class ReturnInst
 {
   friend SILBuilder;
 
+  /// We store the ownership kind in the return inst, but we do not consider the
+  /// underlying return inst to be forwarding. This is because its ownership is
+  /// tied to the function signature and thus should be static.
+  ValueOwnershipKind ownershipKind;
+
   /// Constructs a ReturnInst representing a return.
   ///
-  /// \param DebugLoc The backing AST location.
-  ///
-  /// \param ReturnValue The value to be returned.
-  ///
-  ReturnInst(SILDebugLocation DebugLoc, SILValue ReturnValue)
-      : UnaryInstructionBase(DebugLoc, ReturnValue) {}
+  /// \param func The function we are returning from. Used to compute the
+  ///             preferred ownership kind.
+  /// \param debugLoc The backing AST location.
+  /// \param returnValue The value to be returned.
+  ReturnInst(SILFunction &func, SILDebugLocation debugLoc,
+             SILValue returnValue);
 
 public:
+  /// Return the ownership kind for this instruction if we had any direct
+  /// results.
+  ValueOwnershipKind getOwnershipKind() const { return ownershipKind; }
+
   SuccessorListTy getSuccessors() {
     // No Successors.
     return SuccessorListTy();

lib/SIL/IR/OperandOwnership.cpp

@@ -485,36 +485,11 @@ OperandOwnershipKindClassifier::visitCheckedCastBranchInst(
   return map.getValue();
 }
 
-//// FIX THIS HERE
 OperandOwnershipKindMap
 OperandOwnershipKindClassifier::visitReturnInst(ReturnInst *ri) {
-  auto *f =ri->getFunction();
-
-  // If we have a trivial value, return allLive().
-  bool isTrivial = ri->getOperand()->getType().isTrivial(*f);
-  if (isTrivial) {
-    return Map::allLive();
-  }
-
-  SILFunctionConventions fnConv = f->getConventions();
-
-  auto results = fnConv.getDirectSILResults();
-  if (results.empty())
-    return Map();
-
-  auto ownershipKindRange = makeTransformRange(results,
-     [&](const SILResultInfo &info) {
-       return info.getOwnershipKind(*f, f->getLoweredFunctionType());
-     });
-
-  // Then merge all of our ownership kinds. If we fail to merge, return an empty
-  // map so we fail on all operands.
-  auto mergedBase = ValueOwnershipKind::merge(ownershipKindRange);
-  if (!mergedBase)
-    return Map();
-
-  auto base = *mergedBase;
-  return Map::compatibilityMap(base, base.getForwardingLifetimeConstraint());
+  auto kind = ri->getOwnershipKind();
+  auto lifetimeConstraint = kind.getForwardingLifetimeConstraint();
+  return Map::compatibilityMap(kind, lifetimeConstraint);
 }
 
 OperandOwnershipKindMap

lib/SIL/IR/SILInstructions.cpp

@@ -2863,3 +2863,28 @@ bool GetAsyncContinuationInstBase::throws() const {
   return getType().castTo<BoundGenericType>()->getDecl()
     == getFunction()->getASTContext().getUnsafeThrowingContinuationDecl();
 }
+
+ReturnInst::ReturnInst(SILFunction &func, SILDebugLocation debugLoc,
+                       SILValue returnValue)
+    : UnaryInstructionBase(debugLoc, returnValue),
+      ownershipKind(ValueOwnershipKind::None) {
+  // If we have a trivial value, leave our ownership kind as none.
+  if (returnValue->getType().isTrivial(func))
+    return;
+
+  SILFunctionConventions fnConv = func.getConventions();
+
+  // If we do not have any direct SIL results, we should accept a tuple
+  // argument, meaning that we should have a none ownership kind.
+  auto results = fnConv.getDirectSILResults();
+  if (results.empty())
+    return;
+
+  auto ownershipKindRange =
+      makeTransformRange(results, [&](const SILResultInfo &info) {
+        return info.getOwnershipKind(func, func.getLoweredFunctionType());
+      });
+
+  // Then merge all of our ownership kinds. Assert if we fail to merge.
+  ownershipKind = *ValueOwnershipKind::merge(ownershipKindRange);
+}