[ownership] Loosen restrictions around what we specialize and add generic specialization tests behind a flag.

The idea is that this will let me remove these assertions that were in place to
make sure we were really conservative around specializing ownership code. For me
to remove that I need to be able to actually test out this code (since I think
there are some code paths where this will trigger in other parts of the compiler
now).

So to work out the kinks, I added a flag that allows for the generic specializer
to process ownership code and translated most of the .sil test cases/fixed any
bugs that I found. This hopefully will expose anything that is missing.

NOTE: I have not enabled the generic specializer running in ownership in the
pipeline. This is just a step in that direction by adding tests/etc.

Author: gottesmm

include/swift/SIL/ApplySite.h

@@ -21,9 +21,11 @@
 #ifndef SWIFT_SIL_APPLYSITE_H
 #define SWIFT_SIL_APPLYSITE_H
 
+#include "swift/SIL/SILArgument.h"
 #include "swift/SIL/SILBasicBlock.h"
-#include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/SILFunction.h"
+#include "swift/SIL/SILInstruction.h"
+#include "llvm/ADT/ArrayRef.h"
 
 namespace swift {
 
@@ -502,6 +504,34 @@ class FullApplySite : public ApplySite {
     return getSubstCalleeConv().hasIndirectSILResults();
   }
 
+  /// If our apply site has a single direct result SILValue, return that
+  /// SILValue. Return SILValue() otherwise.
+  ///
+  /// This means that:
+  ///
+  /// 1. If we have an ApplyInst, we just visit the apply.
+  /// 2. If we have a TryApplyInst, we visit the first argument of the normal
+  ///    block.
+  /// 3. If we have a BeginApplyInst, we return SILValue() since the begin_apply
+  ///    yields values instead of returning them. A returned value should only
+  ///    be valid after a full apply site has completely finished executing.
+  SILValue getSingleDirectResult() const {
+    switch (getKind()) {
+    case FullApplySiteKind::ApplyInst:
+      return SILValue(cast<ApplyInst>(getInstruction()));
+    case FullApplySiteKind::BeginApplyInst: {
+      return SILValue();
+    }
+    case FullApplySiteKind::TryApplyInst: {
+      auto *normalBlock = cast<TryApplyInst>(getInstruction())->getNormalBB();
+      assert(normalBlock->getNumArguments() == 1 &&
+             "Expected try apply to have a single result");
+      return normalBlock->getArgument(0);
+    }
+    }
+    llvm_unreachable("Covered switch isn't covered?!");
+  }
+
   unsigned getNumIndirectSILResults() const {
     return getSubstCalleeConv().getNumIndirectSILResults();
   }

include/swift/SIL/SILBuilder.h

@@ -726,7 +726,8 @@ class SILBuilder {
   }
 
   LoadBorrowInst *createLoadBorrow(SILLocation Loc, SILValue LV) {
-    assert(isLoadableOrOpaque(LV->getType()));
+    assert(isLoadableOrOpaque(LV->getType()) &&
+           !LV->getType().isTrivial(getFunction()));
     return insert(new (getModule())
                       LoadBorrowInst(getSILDebugLocation(Loc), LV));
   }
@@ -737,11 +738,19 @@ class SILBuilder {
                       BeginBorrowInst(getSILDebugLocation(Loc), LV));
   }
 
+  /// Convenience function for creating a load_borrow on non-trivial values and
+  /// load [trivial] on trivial values. Becomes load unqualified in non-ossa
+  /// functions.
   SILValue emitLoadBorrowOperation(SILLocation loc, SILValue v) {
     if (!hasOwnership()) {
       return emitLoadValueOperation(loc, v,
                                     LoadOwnershipQualifier::Unqualified);
     }
+
+    if (v->getType().isTrivial(getFunction())) {
+      return emitLoadValueOperation(loc, v, LoadOwnershipQualifier::Trivial);
+    }
+
     return createLoadBorrow(loc, v);
   }
 
@@ -877,6 +886,33 @@ class SILBuilder {
                       StoreBorrowInst(getSILDebugLocation(Loc), Src, DestAddr));
   }
 
+  /// A helper function for emitting store_borrow in operations where one must
+  /// handle both ossa and non-ossa code.
+  ///
+  /// In words:
+  ///
+  /// * If the function does not have ownership, this just emits an unqualified
+  ///   store.
+  ///
+  /// * If the function has ownership, but the type is trivial, use store
+  ///   [trivial].
+  ///
+  /// * Otherwise, emit an actual store_borrow.
+  void emitStoreBorrowOperation(SILLocation loc, SILValue src,
+                                SILValue destAddr) {
+    if (!hasOwnership()) {
+      return emitStoreValueOperation(loc, src, destAddr,
+                                     StoreOwnershipQualifier::Unqualified);
+    }
+
+    if (src->getType().isTrivial(getFunction())) {
+      return emitStoreValueOperation(loc, src, destAddr,
+                                     StoreOwnershipQualifier::Trivial);
+    }
+
+    createStoreBorrow(loc, src, destAddr);
+  }
+
   MarkUninitializedInst *
   createMarkUninitialized(SILLocation Loc, SILValue src,
                           MarkUninitializedInst::Kind k) {

lib/SILOptimizer/Transforms/GenericSpecializer.cpp

@@ -28,6 +28,10 @@
 
 using namespace swift;
 
+// For testing during bring up.
+static llvm::cl::opt<bool> EnableGenericSpecializerWithOwnership(
+    "sil-generic-specializer-enable-ownership", llvm::cl::init(false));
+
 namespace {
 
 class GenericSpecializer : public SILFunctionTransform {
@@ -39,7 +43,7 @@ class GenericSpecializer : public SILFunctionTransform {
     SILFunction &F = *getFunction();
 
     // TODO: We should be able to handle ownership.
-    if (F.hasOwnership())
+    if (F.hasOwnership() && !EnableGenericSpecializerWithOwnership)
       return;
 
     LLVM_DEBUG(llvm::dbgs() << "***** GenericSpecializer on function:"

lib/SILOptimizer/Utils/GenericCloner.cpp

@@ -79,8 +79,9 @@ void GenericCloner::populateCloned() {
 
     auto createAllocStack = [&]() {
       // We need an alloc_stack as a replacement for the indirect parameter.
-      assert(mappedType.isAddress());
-      mappedType = mappedType.getObjectType();
+      if (mappedType.isAddress()) {
+        mappedType = mappedType.getObjectType();
+      }
       auto AllocStackLoc = RegularLocation::getAutoGeneratedLocation();
       ASI = getBuilder().createAllocStack(AllocStackLoc, mappedType);
       AllocStacks.push_back(ASI);
@@ -106,24 +107,36 @@ void GenericCloner::populateCloned() {
         // Handle arguments for formal parameters.
         unsigned paramIdx = ArgIdx - origConv.getSILArgIndexOfFirstParam();
         if (ReInfo.isParamConverted(paramIdx)) {
-          // Store the new direct parameter to the alloc_stack.
-          createAllocStack();
+          assert(mappedType.isAddress());
+          mappedType = mappedType.getObjectType();
           auto *NewArg = ClonedEntryBB->createFunctionArgument(
               mappedType, OrigArg->getDecl());
-          getBuilder().createStore(Loc, NewArg, ASI,
-                                   StoreOwnershipQualifier::Unqualified);
 
-          // Try to create a new debug_value from an existing debug_value_addr.
+          // Try to create a new debug_value from an existing debug_value_addr
+          // for the argument. We do this before storing to ensure that when we
+          // are cloning code in ossa the argument has not been consumed by the
+          // store below.
           for (Operand *ArgUse : OrigArg->getUses()) {
             if (auto *DVAI = dyn_cast<DebugValueAddrInst>(ArgUse->getUser())) {
+              auto *oldScope = getBuilder().getCurrentDebugScope();
               getBuilder().setCurrentDebugScope(
                   remapScope(DVAI->getDebugScope()));
               getBuilder().createDebugValue(DVAI->getLoc(), NewArg,
                                             *DVAI->getVarInfo());
-              getBuilder().setCurrentDebugScope(nullptr);
+              getBuilder().setCurrentDebugScope(oldScope);
               break;
             }
           }
+
+          // Store the new direct parameter to an alloc_stack.
+          createAllocStack();
+          if (!NewArg->getArgumentConvention().isGuaranteedConvention()) {
+            getBuilder().emitStoreValueOperation(Loc, NewArg, ASI,
+                                                 StoreOwnershipQualifier::Init);
+          } else {
+            getBuilder().emitStoreBorrowOperation(Loc, NewArg, ASI);
+          }
+
           entryArgs.push_back(ASI);
           return true;
         }
@@ -150,9 +163,9 @@ void GenericCloner::visitTerminator(SILBasicBlock *BB) {
     if (ReturnValueAddr) {
       // The result is converted from indirect to direct. We have to load the
       // returned value from the alloc_stack.
-      ReturnValue =
-        getBuilder().createLoad(ReturnValueAddr->getLoc(), ReturnValueAddr,
-                                LoadOwnershipQualifier::Unqualified);
+      ReturnValue = getBuilder().emitLoadValueOperation(
+          ReturnValueAddr->getLoc(), ReturnValueAddr,
+          LoadOwnershipQualifier::Take);
     }
     for (AllocStackInst *ASI : reverse(AllocStacks)) {
       getBuilder().createDeallocStack(ASI->getLoc(), ASI);