Introduce return_borrow and unchecked_ownership instructions for use with borrow accessors

SwiftCompilerSources/Sources/SIL/Instruction.swift

@@ -1447,6 +1447,8 @@ final public class CopyableToMoveOnlyWrapperAddrInst
 final public class MoveOnlyWrapperToCopyableAddrInst
   : SingleValueInstruction, UnaryInstruction {}
 
+final public class UncheckedOwnershipInst: SingleValueInstruction, UnaryInstruction {}
+
 final public class ObjectInst : SingleValueInstruction {
   public var baseOperands: OperandArray {
     operands[0..<bridged.ObjectInst_getNumBaseElements()]
@@ -1857,6 +1859,18 @@ final public class ReturnInst : TermInst, UnaryInstruction {
   public override var isFunctionExiting: Bool { true }
 }
 
+final public class ReturnBorrowInst : TermInst {
+  public var returnValue: Value { operands[0].value }
+  public var enclosingOperands: OperandArray {
+    let ops = operands
+    return ops[1..<ops.count]
+  }
+  public var enclosingValues: LazyMapSequence<LazySequence<OperandArray>.Elements, Value> {
+    enclosingOperands.values
+  }
+  public override var isFunctionExiting: Bool { true }
+}
+
 final public class ThrowInst : TermInst, UnaryInstruction {
   public var thrownValue: Value { operand.value }
   public override var isFunctionExiting: Bool { true }

SwiftCompilerSources/Sources/SIL/Registration.swift

@@ -120,6 +120,7 @@ private func registerSILClasses() {
   register(MoveOnlyWrapperToCopyableBoxInst.self)
   register(CopyableToMoveOnlyWrapperAddrInst.self)
   register(MoveOnlyWrapperToCopyableAddrInst.self)
+  register(UncheckedOwnershipInst.self) 
   register(ObjectInst.self)
   register(VectorInst.self)
   register(VectorBaseAddrInst.self)
@@ -247,6 +248,7 @@ private func registerSILClasses() {
 
   register(UnreachableInst.self)
   register(ReturnInst.self)
+  register(ReturnBorrowInst.self)
   register(ThrowInst.self)
   register(ThrowAddrInst.self)
   register(YieldInst.self)

docs/SIL/Instructions.md

@@ -1199,6 +1199,17 @@ the `set` case is passed `self`.
 
 This is only valid in Raw SIL.
 
+### unchecked_ownership
+
+```
+sil-instruction ::= 'unchecked_ownership' sil-operand
+
+unchecked_ownership %1 : $T
+```
+
+unchecked_ownership disables the ownership verification of it's operand. This used in cases 
+we cannot resolve ownership until a mandatory pass runs. This is only valid in Raw SIL. 
+
 ### copy_addr
 
 ```
@@ -5003,6 +5014,20 @@ does not apply in the `raw` SIL stage.
 
 A function must not contain more than one `return` instruction.
 
+### return_borrow
+
+```
+sil-terminator ::= 'return_borrow' sil-operand 'from_scopes' '(' (sil-operand (',' sil-operand)*)? ')'
+
+return_borrow %0 : $T from_scopes (%1, %2 ...)
+// %0 must be a @guaranteed value
+// %1, %2, ... must be borrow introducers for %0, like `load_borrow`
+// $T must be the return type of the current function
+```
+
+return_borrow instruction is valid only for functions @guaranteed results.
+It is used to a return a @guaranteed value that maybe produced within borrow scopes local to the function.
+
 ### throw
 
 ```

include/swift/SIL/AddressWalker.h

@@ -185,6 +185,7 @@ TransitiveAddressWalker<Impl>::walk(SILValue projectedAddress) {
         llvm_unreachable("Never takes an address");
       // Point uses.
       case TermKind::ReturnInst:
+      case TermKind::ReturnBorrowInst:
       case TermKind::ThrowInst:
       case TermKind::YieldInst:
       case TermKind::TryApplyInst:

include/swift/SIL/SILBuilder.h

@@ -1555,6 +1555,12 @@ class SILBuilder {
         MoveOnlyWrapperToCopyableValueInst::Guaranteed));
   }
 
+  UncheckedOwnershipInst *createUncheckedOwnership(SILLocation Loc,
+                                                   SILValue Operand) {
+    return insert(new (getModule()) UncheckedOwnershipInst(
+        getSILDebugLocation(Loc), Operand, Operand->getOwnershipKind()));
+  }
+
   UnconditionalCheckedCastInst *
   createUnconditionalCheckedCast(SILLocation Loc,
                                  CheckedCastInstOptions options,
@@ -2658,6 +2664,12 @@ class SILBuilder {
         getFunction(), getSILDebugLocation(Loc), ReturnValue));
   }
 
+  ReturnBorrowInst *createReturnBorrow(SILLocation Loc, SILValue returnValue,
+                                       ArrayRef<SILValue> enclosingValues) {
+    return insertTerminator(ReturnBorrowInst::create(
+        getSILDebugLocation(Loc), returnValue, enclosingValues, getModule()));
+  }
+
   ThrowInst *createThrow(SILLocation Loc, SILValue errorValue) {
     return insertTerminator(
         new (getModule()) ThrowInst(getSILDebugLocation(Loc), errorValue));

include/swift/SIL/SILCloner.h

@@ -2247,6 +2247,19 @@ void SILCloner<ImplClass>::visitCopyableToMoveOnlyWrapperValueInst(
   recordClonedInstruction(inst, cvt);
 }
 
+template <typename ImplClass>
+void SILCloner<ImplClass>::visitUncheckedOwnershipInst(
+    UncheckedOwnershipInst *uoi) {
+  getBuilder().setCurrentDebugScope(getOpScope(uoi->getDebugScope()));
+  if (!getBuilder().hasOwnership()) {
+    return recordFoldedValue(uoi, getOpValue(uoi->getOperand()));
+  }
+
+  recordClonedInstruction(
+      uoi, getBuilder().createUncheckedOwnership(
+               getOpLocation(uoi->getLoc()), getOpValue(uoi->getOperand())));
+}
+
 template <typename ImplClass>
 void SILCloner<ImplClass>::visitReleaseValueInst(ReleaseValueInst *Inst) {
   getBuilder().setCurrentDebugScope(getOpScope(Inst->getDebugScope()));
@@ -3417,6 +3430,22 @@ SILCloner<ImplClass>::visitReturnInst(ReturnInst *Inst) {
                                       getOpValue(Inst->getOperand())));
 }
 
+template <typename ImplClass>
+void SILCloner<ImplClass>::visitReturnBorrowInst(ReturnBorrowInst *rbi) {
+  getBuilder().setCurrentDebugScope(getOpScope(rbi->getDebugScope()));
+  if (!getBuilder().hasOwnership()) {
+    return recordClonedInstruction(
+        rbi, getBuilder().createReturn(getOpLocation(rbi->getLoc()),
+                                       getOpValue(rbi->getReturnValue())));
+  }
+
+  auto enclosingValues = getOpValueArray<8>(rbi->getEnclosingValues());
+  recordClonedInstruction(
+      rbi, getBuilder().createReturnBorrow(getOpLocation(rbi->getLoc()),
+                                           getOpValue(rbi->getReturnValue()),
+                                           enclosingValues));
+}
+
 template<typename ImplClass>
 void
 SILCloner<ImplClass>::visitThrowInst(ThrowInst *Inst) {

include/swift/SIL/SILInstruction.h

@@ -9553,6 +9553,17 @@ class MoveOnlyWrapperToCopyableAddrInst
   }
 };
 
+class UncheckedOwnershipInst final
+    : public UnaryInstructionBase<SILInstructionKind::UncheckedOwnershipInst,
+                                  OwnershipForwardingSingleValueInstruction> {
+  friend class SILBuilder;
+
+  UncheckedOwnershipInst(SILDebugLocation DebugLoc, SILValue operand,
+                         ValueOwnershipKind forwardingOwnershipKind)
+      : UnaryInstructionBase(DebugLoc, operand, operand->getType(),
+                             forwardingOwnershipKind) {}
+};
+
 /// Given an object reference, return true iff it is non-nil and refers
 /// to a native swift object with strong reference count of 1.
 class IsUniqueInst
@@ -10123,6 +10134,7 @@ class TermInst : public NonValueInstruction {
     case TermKind::UnwindInst:
     case TermKind::UnreachableInst:
     case TermKind::ReturnInst:
+    case TermKind::ReturnBorrowInst:
     case TermKind::ThrowInst:
     case TermKind::ThrowAddrInst:
     case TermKind::YieldInst:
@@ -10252,6 +10264,34 @@ class ReturnInst
   }
 };
 
+class ReturnBorrowInst final
+    : public InstructionBaseWithTrailingOperands<
+          SILInstructionKind::ReturnBorrowInst, ReturnBorrowInst, TermInst> {
+  friend SILBuilder;
+
+  ReturnBorrowInst(SILDebugLocation DebugLoc, ArrayRef<SILValue> operands);
+
+  static ReturnBorrowInst *create(SILDebugLocation DebugLoc, SILValue value,
+                                  ArrayRef<SILValue> enclosingValues,
+                                  SILModule &M);
+
+public:
+  SILValue getReturnValue() const { return getAllOperands()[0].get(); }
+
+  ArrayRef<Operand> getEnclosingValueOperands() const {
+    return getAllOperands().drop_front();
+  }
+
+  OperandValueArrayRef getEnclosingValues() const {
+    return OperandValueArrayRef(getEnclosingValueOperands());
+  }
+
+  SuccessorListTy getSuccessors() {
+    // No Successors.
+    return SuccessorListTy();
+  }
+};
+
 /// ThrowInst - Throw a typed error, returning it via the direct error result.
 class ThrowInst
   : public UnaryInstructionBase<SILInstructionKind::ThrowInst, TermInst>
@@ -11667,6 +11707,7 @@ OwnershipForwardingSingleValueInstruction::classof(SILInstructionKind kind) {
   case SILInstructionKind::DropDeinitInst:
   case SILInstructionKind::BorrowedFromInst:
   case SILInstructionKind::ImplicitActorToOpaqueIsolationCastInst:
+  case SILInstructionKind::UncheckedOwnershipInst:
     return true;
   default:
     return false;

include/swift/SIL/SILNodes.def

@@ -469,7 +469,8 @@ ABSTRACT_VALUE_AND_INST(SingleValueInstruction, ValueBase, SILInstruction)
                     SingleValueInstruction, MayHaveSideEffects, DoesNotRelease)
 #include "swift/AST/ReferenceStorage.def"
   SINGLE_VALUE_INST(UncheckedOwnershipConversionInst, unchecked_ownership_conversion,
-                    SingleValueInstruction, None, MayRelease)
+                    SingleValueInstruction, None, DoesNotRelease)
+
   // A move_value is an OSSA only instruction. Its result does not have any side
   // effects relative to other OSSA values like copy_value.
   SINGLE_VALUE_INST(MoveValueInst, move_value, SingleValueInstruction, None,
@@ -520,6 +521,8 @@ ABSTRACT_VALUE_AND_INST(SingleValueInstruction, ValueBase, SILInstruction)
   SINGLE_VALUE_INST(CopyableToMoveOnlyWrapperAddrInst,
                     copyable_to_moveonlywrapper_addr, SingleValueInstruction,
                     None, DoesNotRelease)
+  SINGLE_VALUE_INST(UncheckedOwnershipInst, unchecked_ownership,
+                    SingleValueInstruction, None, DoesNotRelease)
 
   // IsUnique does not actually write to memory but should be modeled
   // as such. Its operand is a pointer to an object reference. The
@@ -731,6 +734,8 @@ ABSTRACT_INST(TermInst, SILInstruction)
              TermInst, None, DoesNotRelease)
   TERMINATOR(ReturnInst, return,
              TermInst, None, DoesNotRelease)
+  TERMINATOR(ReturnBorrowInst, return_borrow,
+             TermInst, None, DoesNotRelease)
   TERMINATOR(ThrowInst, throw,
              TermInst, None, DoesNotRelease)
   TERMINATOR(ThrowAddrInst, throw_addr,

include/swift/SILOptimizer/Utils/SCCVisitor.h

@@ -130,6 +130,7 @@ class SCCVisitor {
 
     case TermKind::UnreachableInst:
     case TermKind::ReturnInst:
+    case TermKind::ReturnBorrowInst:
     case TermKind::SwitchValueInst:
     case TermKind::ThrowInst:
     case TermKind::ThrowAddrInst:

lib/IRGen/IRGenSIL.cpp

@@ -1302,7 +1302,9 @@ class IRGenSILFunction :
     auto e = getLoweredExplosion(i->getOperand());
     setLoweredExplosion(i, e);
   }
-
+  void visitUncheckedOwnershipInst(UncheckedOwnershipInst *i) {
+    llvm_unreachable("unimplemented");
+  }
   void visitMergeIsolationRegionInst(MergeIsolationRegionInst *i) {
     llvm_unreachable("Valid only when ownership is enabled");
   }
@@ -1504,6 +1506,9 @@ class IRGenSILFunction :
   void visitBranchInst(BranchInst *i);
   void visitCondBranchInst(CondBranchInst *i);
   void visitReturnInst(ReturnInst *i);
+  void visitReturnBorrowInst(ReturnBorrowInst *i) {
+    llvm_unreachable("unimplemented");
+  }
   void visitThrowInst(ThrowInst *i);
   void visitThrowAddrInst(ThrowAddrInst *i);
   void visitUnwindInst(UnwindInst *i);

lib/SIL/IR/OperandOwnership.cpp

@@ -421,6 +421,7 @@ AGGREGATE_OWNERSHIP(DestructureTuple)
 AGGREGATE_OWNERSHIP(Enum)
 AGGREGATE_OWNERSHIP(UncheckedEnumData)
 AGGREGATE_OWNERSHIP(SwitchEnum)
+AGGREGATE_OWNERSHIP(UncheckedOwnership)
 #undef AGGREGATE_OWNERSHIP
 
 // A begin_borrow is conditionally nested.
@@ -636,6 +637,12 @@ OperandOwnership OperandOwnershipClassifier::visitReturnInst(ReturnInst *i) {
   llvm_unreachable("covered switch");
 }
 
+OperandOwnership
+OperandOwnershipClassifier::visitReturnBorrowInst(ReturnBorrowInst *rbi) {
+  return getOperandIndex() == 0 ? OperandOwnership::GuaranteedForwarding
+                                : OperandOwnership::EndBorrow;
+}
+
 OperandOwnership OperandOwnershipClassifier::visitAssignInst(AssignInst *i) {
   if (getValue() != i->getSrc()) {
     return OperandOwnership::TrivialUse;

lib/SIL/IR/SILArgument.cpp

@@ -315,6 +315,7 @@ getSingleTerminatorOperandForPred(const SILBasicBlock *parentBlock,
   switch (predTermInst->getTermKind()) {
   case TermKind::UnreachableInst:
   case TermKind::ReturnInst:
+  case TermKind::ReturnBorrowInst:
   case TermKind::ThrowInst:
   case TermKind::ThrowAddrInst:
   case TermKind::UnwindInst:

lib/SIL/IR/SILInstructions.cpp

@@ -1814,6 +1814,7 @@ bool TermInst::isFunctionExiting() const {
   case TermKind::YieldInst:
     return false;
   case TermKind::ReturnInst:
+  case TermKind::ReturnBorrowInst:
   case TermKind::ThrowInst:
   case TermKind::ThrowAddrInst:
   case TermKind::UnwindInst:
@@ -1835,6 +1836,7 @@ bool TermInst::isProgramTerminating() const {
   case TermKind::CheckedCastBranchInst:
   case TermKind::CheckedCastAddrBranchInst:
   case TermKind::ReturnInst:
+  case TermKind::ReturnBorrowInst:
   case TermKind::ThrowInst:
   case TermKind::ThrowAddrInst:
   case TermKind::UnwindInst:
@@ -1861,6 +1863,7 @@ const Operand *TermInst::forwardedOperand() const {
   case TermKind::UnwindInst:
   case TermKind::UnreachableInst:
   case TermKind::ReturnInst:
+  case TermKind::ReturnBorrowInst:
   case TermKind::ThrowInst:
   case TermKind::ThrowAddrInst:
   case TermKind::YieldInst:
@@ -3428,6 +3431,26 @@ ReturnInst::ReturnInst(SILFunction &func, SILDebugLocation debugLoc,
          "result info?!");
 }
 
+ReturnBorrowInst *ReturnBorrowInst::create(SILDebugLocation DebugLoc,
+                                           SILValue returnValue,
+                                           ArrayRef<SILValue> enclosingValues,
+                                           SILModule &M) {
+  auto Size = totalSizeToAlloc<swift::Operand>(enclosingValues.size() + 1);
+  auto Buffer = M.allocateInst(Size, alignof(ReturnBorrowInst));
+  SmallVector<SILValue, 8> operands;
+  operands.push_back(returnValue);
+  for (SILValue ev : enclosingValues) {
+    operands.push_back(ev);
+  }
+  return ::new (Buffer) ReturnBorrowInst(DebugLoc, operands);
+}
+
+ReturnBorrowInst::ReturnBorrowInst(SILDebugLocation DebugLoc,
+                                   ArrayRef<SILValue> operands)
+    : InstructionBaseWithTrailingOperands(operands, DebugLoc) {
+  assert(operands[0]->getOwnershipKind() == OwnershipKind::Guaranteed);
+}
+
 // 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.

lib/SIL/IR/SILPrinter.cpp

@@ -2367,6 +2367,10 @@ class SILPrinter : public SILInstructionVisitor<SILPrinter> {
     *this << getIDAndType(I->getOperand());
   }
 
+  void visitUncheckedOwnershipInst(UncheckedOwnershipInst *I) {
+    *this << getIDAndType(I->getOperand());
+  }
+
   void visitUnownedCopyValueInst(UnownedCopyValueInst *I) {
     *this << getIDAndType(I->getOperand());
   }
@@ -2907,6 +2911,21 @@ class SILPrinter : public SILInstructionVisitor<SILPrinter> {
     *this << getIDAndType(RI->getOperand());
   }
 
+  void visitReturnBorrowInst(ReturnBorrowInst *rbi) {
+    *this << getIDAndType(rbi->getReturnValue());
+
+    *this << " from_scopes (";
+    bool first = true;
+    for (SILValue ev : rbi->getEnclosingValues()) {
+      if (!first) {
+        *this << ", ";
+      }
+      first = false;
+      *this << getIDAndType(ev);
+    }
+    *this << ")";
+  }
+
   void visitSpecifyTestInst(SpecifyTestInst *TSI) {
     *this << QuotedString(TSI->getArgumentsSpecification());
   }