Use in_guaranteed for let captures (#29812)

* Use in_guaranteed for let captures

With this all let values will be captured with in_guaranteed convention
by the closure. Following are the main changes :

SILGen changes:
- A new CaptureKind::Immutable is introduced, to capture let values as in_guaranteed.
- SILGen of in_guaranteed capture had to be fixed.
  in_guaranteed captures as per convention are consumed by the closure. And so SILGen should not generate a destroy_addr for an in_guaranteed capture.
  But LetValueInitialization can push Dealloc and Release states of the captured arg in the Cleanup stack, and there is no way to access the CleanupHandle and disable the emission of destroy_addr while emitting the captures in SILGenFunction::emitCaptures.
  So we now create, temporary allocation of the in_guaranteed capture iduring SILGenFunction::emitCaptures without emitting destroy_addr for it.

SILOptimizer changes:
- Handle in_guaranteed in CopyForwarding.
- Adjust dealloc_stack of in_guaranteed capture to occur after destroy_addr for on_stack closures in ClosureLifetimeFixup.

IRGen changes :
  - Since HeapLayout can be non-fixed now, make sure emitSize is used conditionally
  - Don't consider ClassPointerSource kind parameter type for fulfillments while generating code for partial apply forwarder.
    The TypeMetadata of ClassPointSource kind sources are not populated in HeapLayout's NecessaryBindings. If we have a generic parameter on the HeapLayout which can be fulfilled by a ClassPointerSource, its TypeMetaData will not be found while constructing the dtor function of the HeapLayout.
    So it is important to skip considering sources of ClassPointerSource kind, so that TypeMetadata of a dependent generic parameters gets populated in HeapLayout's NecessaryBindings.

Author: meg-gupta

include/swift/SIL/TypeLowering.h

@@ -561,6 +561,8 @@ enum class CaptureKind {
   StorageAddress,
   /// A local value captured as a constant.
   Constant,
+  /// A let constant captured as a pointer to storage
+  Immutable
 };
 
 

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -364,12 +364,19 @@ void releasePartialApplyCapturedArg(
     SILParameterInfo paramInfo,
     InstModCallbacks callbacks = InstModCallbacks());
 
+void deallocPartialApplyCapturedArg(
+    SILBuilder &builder, SILLocation loc, SILValue arg,
+    SILParameterInfo paramInfo);
+
 /// Insert destroys of captured arguments of partial_apply [stack].
 void insertDestroyOfCapturedArguments(
     PartialApplyInst *pai, SILBuilder &builder,
     llvm::function_ref<bool(SILValue)> shouldInsertDestroy =
         [](SILValue arg) -> bool { return true; });
 
+void insertDeallocOfCapturedArguments(
+    PartialApplyInst *pai, SILBuilder &builder);
+
 /// This iterator 'looks through' one level of builtin expect users exposing all
 /// users of the looked through builtin expect instruction i.e it presents a
 /// view that shows all users as if there were no builtin expect instructions

lib/IRGen/GenFunc.cpp

@@ -1319,9 +1319,25 @@ Optional<StackAddress> irgen::emitFunctionPartialApplication(
   SmallVector<SILType, 4> argValTypes;
   SmallVector<ParameterConvention, 4> argConventions;
 
+ // Go over the params and check if any of them can cause the HeapLayout to be non-fixed
+ // This is needed because we should not consider parameters of kind ClassPointer and Metadata sources
+ // If not, we may end up with missing TypeMetadata for a type dependent generic parameter
+ // while generating code for destructor of HeapLayout.
+  bool considerParameterSources = true;
+  for (auto param : params) {
+    SILType argType = IGF.IGM.silConv.getSILType(param, origType);
+    auto argLoweringTy = getArgumentLoweringType(argType.getASTType(), param);
+    auto &ti = IGF.getTypeInfoForLowered(argLoweringTy);
+
+    if (!isa<FixedTypeInfo>(ti)) {
+      considerParameterSources = false;
+      break;
+    }
+  }
+
   // Reserve space for polymorphic bindings.
   auto bindings =
-      NecessaryBindings::forPartialApplyForwarder(IGF.IGM, origType, subs);
+      NecessaryBindings::forPartialApplyForwarder(IGF.IGM, origType, subs, considerParameterSources);
 
   if (!bindings.empty()) {
     hasSingleSwiftRefcountedContext = No;
@@ -1524,7 +1540,7 @@ Optional<StackAddress> irgen::emitFunctionPartialApplication(
     HeapNonFixedOffsets offsets(IGF, layout);
     if (outType->isNoEscape()) {
       stackAddr = IGF.emitDynamicAlloca(
-          IGF.IGM.Int8Ty, layout.emitSize(IGF.IGM), Alignment(16));
+          IGF.IGM.Int8Ty, layout.isFixedLayout() ? layout.emitSize(IGF.IGM) : offsets.getSize() , Alignment(16));
       stackAddr = stackAddr->withAddress(IGF.Builder.CreateBitCast(
           stackAddr->getAddress(), IGF.IGM.OpaquePtrTy));
       data = stackAddr->getAddress().getAddress();

lib/IRGen/GenProto.cpp

@@ -112,7 +112,7 @@ class PolymorphicConvention {
   }
 
 public:
-  PolymorphicConvention(IRGenModule &IGM, CanSILFunctionType fnType);
+  PolymorphicConvention(IRGenModule &IGM, CanSILFunctionType fnType, bool considerParameterSources);
 
   ArrayRef<MetadataSource> getSources() const { return Sources; }
 
@@ -179,8 +179,9 @@ class PolymorphicConvention {
 } // end anonymous namespace
 
 PolymorphicConvention::PolymorphicConvention(IRGenModule &IGM,
-                                             CanSILFunctionType fnType)
-  : IGM(IGM), M(*IGM.getSwiftModule()), FnType(fnType) {
+                                             CanSILFunctionType fnType,
+                                             bool considerParameterSources = true)
+  : IGM(IGM), M(*IGM.getSwiftModule()), FnType(fnType){
   initGenerics();
 
   auto rep = fnType->getRepresentation();
@@ -212,16 +213,18 @@ PolymorphicConvention::PolymorphicConvention(IRGenModule &IGM,
     unsigned selfIndex = ~0U;
     auto params = fnType->getParameters();
 
-    // Consider 'self' first.
-    if (fnType->hasSelfParam()) {
-      selfIndex = params.size() - 1;
-      considerParameter(params[selfIndex], selfIndex, true);
-    }
+    if (considerParameterSources) {
+      // Consider 'self' first.
+      if (fnType->hasSelfParam()) {
+        selfIndex = params.size() - 1;
+        considerParameter(params[selfIndex], selfIndex, true);
+      }
 
-    // Now consider the rest of the parameters.
-    for (auto index : indices(params)) {
-      if (index != selfIndex)
-        considerParameter(params[index], index, false);
+      // Now consider the rest of the parameters.
+      for (auto index : indices(params)) {
+        if (index != selfIndex)
+          considerParameter(params[index], index, false);
+      }
     }
   }
 }
@@ -3000,14 +3003,16 @@ NecessaryBindings::forFunctionInvocations(IRGenModule &IGM,
 NecessaryBindings
 NecessaryBindings::forPartialApplyForwarder(IRGenModule &IGM,
                                           CanSILFunctionType origType,
-                                          SubstitutionMap subs) {
+                                          SubstitutionMap subs,
+                                          bool considerParameterSources) {
   return computeBindings(IGM, origType, subs,
-                         true /*forPartialApplyForwarder*/);
+                         true  /*forPartialApplyForwarder*/,
+                         considerParameterSources);
 }
 
 NecessaryBindings NecessaryBindings::computeBindings(
     IRGenModule &IGM, CanSILFunctionType origType, SubstitutionMap subs,
-    bool forPartialApplyForwarder) {
+    bool forPartialApplyForwarder, bool considerParameterSources) {
 
   NecessaryBindings bindings;
   bindings.forPartialApply = forPartialApplyForwarder;
@@ -3017,7 +3022,7 @@ NecessaryBindings NecessaryBindings::computeBindings(
     return bindings;
 
   // Figure out what we're actually required to pass:
-  PolymorphicConvention convention(IGM, origType);
+  PolymorphicConvention convention(IGM, origType, considerParameterSources);
 
   //  - unfulfilled requirements
   convention.enumerateUnfulfilledRequirements(

lib/IRGen/NecessaryBindings.h

@@ -59,7 +59,8 @@ class NecessaryBindings {
                                                   SubstitutionMap subs);
   static NecessaryBindings forPartialApplyForwarder(IRGenModule &IGM,
                                                     CanSILFunctionType origType,
-                                                    SubstitutionMap subs);
+                                                    SubstitutionMap subs,
+                                                    bool considerParameterSources = true);
 
   /// Add whatever information is necessary to reconstruct type metadata
   /// for the given type.
@@ -98,7 +99,8 @@ class NecessaryBindings {
   static NecessaryBindings computeBindings(IRGenModule &IGM,
                                            CanSILFunctionType origType,
                                            SubstitutionMap subs,
-                                           bool forPartialApplyForwarder);
+                                           bool forPartialApplyForwarder,
+                                           bool considerParameterSources = true);
 };
 
 } // end namespace irgen

lib/SIL/SILFunctionType.cpp

@@ -1460,9 +1460,8 @@ lowerCaptureContextParameters(TypeConverter &TC, SILDeclRef function,
     case CaptureKind::Constant: {
       // Constants are captured by value.
       ParameterConvention convention;
-      if (loweredTL.isAddressOnly()) {
-        convention = ParameterConvention::Indirect_In_Guaranteed;
-      } else if (loweredTL.isTrivial()) {
+      assert (!loweredTL.isAddressOnly());
+      if (loweredTL.isTrivial()) {
         convention = ParameterConvention::Direct_Unowned;
       } else {
         convention = ParameterConvention::Direct_Guaranteed;
@@ -1495,6 +1494,16 @@ lowerCaptureContextParameters(TypeConverter &TC, SILDeclRef function,
       inputs.push_back(param);
       break;
     }
+    case CaptureKind::Immutable: {
+      // 'let' constants that are address-only are captured as the address of the value
+      //  and will be consumed by the closure.
+      SILType ty = loweredTy.getAddressType();
+      auto param =
+          SILParameterInfo(ty.getASTType(),
+                           ParameterConvention::Indirect_In_Guaranteed);
+      inputs.push_back(param);
+      break;
+    }
     }
   }
 }

lib/SIL/SILVerifier.cpp

@@ -1117,7 +1117,7 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
       if (VarInfo)
         if (unsigned ArgNo = VarInfo->ArgNo) {
           // It is a function argument.
-          if (ArgNo < DebugVars.size() && !DebugVars[ArgNo].empty()) {
+          if (ArgNo < DebugVars.size() && !DebugVars[ArgNo].empty() && !VarInfo->Name.empty()) {
             require(
                 DebugVars[ArgNo] == VarInfo->Name,
                 "Scope contains conflicting debug variables for one function "

lib/SIL/TypeLowering.cpp

@@ -126,6 +126,11 @@ CaptureKind TypeConverter::getDeclCaptureKind(CapturedValue capture,
     return CaptureKind::Box;
   }
 
+  // For 'let' constants
+  if (!var->supportsMutation()) {
+    return CaptureKind::Immutable;
+  }
+
   // If we're capturing into a non-escaping closure, we can generally just
   // capture the address of the value as no-escape.
   return (capture.isNoEscape()

lib/SILGen/SILGenFunction.cpp

@@ -262,6 +262,7 @@ void SILGenFunction::emitCaptures(SILLocation loc,
       case CaptureKind::Constant:
         capturedArgs.push_back(emitUndef(getLoweredType(type)));
         break;
+      case CaptureKind::Immutable:
       case CaptureKind::StorageAddress:
         capturedArgs.push_back(emitUndef(getLoweredType(type).getAddressType()));
         break;
@@ -333,7 +334,23 @@ void SILGenFunction::emitCaptures(SILLocation loc,
       capturedArgs.push_back(emitManagedRValueWithCleanup(Val));
       break;
     }
-
+    case CaptureKind::Immutable: {
+      if (canGuarantee) {
+        auto entryValue = getAddressValue(Entry.value);
+        // No-escaping stored declarations are captured as the
+        // address of the value.
+        assert(entryValue->getType().isAddress() && "no address for captured var!");
+        capturedArgs.push_back(ManagedValue::forLValue(entryValue));
+      }
+      else {
+        auto entryValue = getAddressValue(Entry.value);
+        auto addr = emitTemporaryAllocation(vd, entryValue->getType());
+        auto val = B.createCopyAddr(loc, entryValue, addr, IsNotTake,
+                         IsInitialization);
+        capturedArgs.push_back(ManagedValue::forLValue(addr));
+      }
+      break;
+    }
     case CaptureKind::StorageAddress: {
       auto entryValue = getAddressValue(Entry.value);
       // No-escaping stored declarations are captured as the

lib/SILGen/SILGenProlog.cpp

@@ -420,6 +420,7 @@ static void emitCaptureArguments(SILGenFunction &SGF,
     SGF.B.createDebugValueAddr(Loc, addr, DbgVar);
     break;
   }
+  case CaptureKind::Immutable:
   case CaptureKind::StorageAddress: {
     // Non-escaping stored decls are captured as the address of the value.
     auto type = getVarTypeInCaptureContext();

lib/SILOptimizer/Mandatory/ClosureLifetimeFixup.cpp

@@ -459,6 +459,27 @@ static bool tryRewriteToPartialApplyStack(
     saveDeleteInst(convertOrPartialApply);
   saveDeleteInst(origPA);
 
+  ApplySite site(newPA);
+  SILFunctionConventions calleeConv(site.getSubstCalleeType(),
+                                      newPA->getModule());
+
+  // Since we create temporary allocation for in_guaranteed captures during SILGen,
+  // the dealloc_stack of it can occur before the apply due to conversion scopes.
+  // When we insert destroy_addr of the in_guaranteed capture after the apply,
+  // we may end up with a situation when the dealloc_stack occurs before the destroy_addr.
+  // The code below proactively removes the dealloc_stack of in_guaranteed capture,
+  // so that it can be reinserted at the correct place after the destroy_addr below.
+  for (auto &arg : newPA->getArgumentOperands()) {
+    unsigned calleeArgumentIndex = site.getCalleeArgIndex(arg);
+    assert(calleeArgumentIndex >= calleeConv.getSILArgIndexOfFirstParam());
+    auto paramInfo = calleeConv.getParamInfoForSILArg(calleeArgumentIndex);
+    if (paramInfo.getConvention() == ParameterConvention::Indirect_In_Guaranteed)
+      // go over all the dealloc_stack, remove it
+      for (auto *use : arg.get()->getUses())
+        if (auto *deallocInst = dyn_cast<DeallocStackInst>(use->getUser()))
+          deallocInst->eraseFromParent();
+  }
+
   // Insert destroys of arguments after the apply and the dealloc_stack.
   if (auto *apply = dyn_cast<ApplyInst>(singleApplyUser)) {
     auto insertPt = std::next(SILBasicBlock::iterator(apply));
@@ -468,11 +489,15 @@ static bool tryRewriteToPartialApplyStack(
     SILBuilderWithScope b3(insertPt);
     b3.createDeallocStack(loc, newPA);
     insertDestroyOfCapturedArguments(newPA, b3);
+    // dealloc_stack of the in_guaranteed capture is inserted
+    insertDeallocOfCapturedArguments(newPA, b3);
   } else if (auto *tai = dyn_cast<TryApplyInst>(singleApplyUser)) {
     for (auto *succBB : tai->getSuccessorBlocks()) {
       SILBuilderWithScope b3(succBB->begin());
       b3.createDeallocStack(loc, newPA);
       insertDestroyOfCapturedArguments(newPA, b3);
+      // dealloc_stack of the in_guaranteed capture is inserted
+      insertDeallocOfCapturedArguments(newPA, b3);
     }
   } else {
     llvm_unreachable("Unknown FullApplySite instruction kind");

lib/SILOptimizer/Transforms/TempRValueElimination.cpp

@@ -147,11 +147,23 @@ bool TempRValueOptPass::collectLoads(
     LLVM_DEBUG(llvm::dbgs()
                << "  Temp use may write/destroy its source" << *user);
     return false;
-
   case SILInstructionKind::BeginAccessInst:
     return cast<BeginAccessInst>(user)->getAccessKind() == SILAccessKind::Read;
 
+  case SILInstructionKind::MarkDependenceInst: {
+    auto mdi = cast<MarkDependenceInst>(user);
+    if (mdi->getBase() == address) {
+      return true;
+    }
+    for (auto *mdiUseOper : mdi->getUses()) {
+      if (!collectLoads(mdiUseOper, mdiUseOper->getUser(), mdi, srcAddr,
+                        loadInsts))
+        return false;
+    }
+    return true;
+  }
   case SILInstructionKind::ApplyInst:
+  case SILInstructionKind::PartialApplyInst:
   case SILInstructionKind::TryApplyInst: {
     ApplySite apply(user);
 
@@ -650,7 +662,14 @@ TempRValueOptPass::tryOptimizeStoreIntoTemp(StoreInst *si) {
       toDelete.push_back(fli);
       break;
     }
-
+    case SILInstructionKind::MarkDependenceInst: {
+      SILBuilderWithScope builder(user);
+      auto mdi = cast<MarkDependenceInst>(user);
+      auto newInst = builder.createMarkDependence(user->getLoc(), mdi->getValue(), si->getSrc());
+      mdi->replaceAllUsesWith(newInst);
+      toDelete.push_back(user);
+      break;
+    }
     // ASSUMPTION: no operations that may be handled by this default clause can
     // destroy tempObj. This includes operations that load the value from memory
     // and release it.

lib/SILOptimizer/Utils/InstOptUtils.cpp

@@ -1372,6 +1372,15 @@ void swift::releasePartialApplyCapturedArg(SILBuilder &builder, SILLocation loc,
   callbacks.createdNewInst(u.get<ReleaseValueInst *>());
 }
 
+void swift::deallocPartialApplyCapturedArg(SILBuilder &builder, SILLocation loc,
+                                           SILValue arg,
+                                           SILParameterInfo paramInfo) {
+  if (!paramInfo.isIndirectInGuaranteed())
+    return;
+
+  builder.createDeallocStack(loc, arg);
+}
+
 static bool
 deadMarkDependenceUser(SILInstruction *inst,
                        SmallVectorImpl<SILInstruction *> &deleteInsts) {
@@ -1937,6 +1946,22 @@ void swift::insertDestroyOfCapturedArguments(
   }
 }
 
+void swift::insertDeallocOfCapturedArguments(
+    PartialApplyInst *pai, SILBuilder &builder) {
+  assert(pai->isOnStack());
+
+  ApplySite site(pai);
+  SILFunctionConventions calleeConv(site.getSubstCalleeType(),
+                                    pai->getModule());
+  auto loc = RegularLocation::getAutoGeneratedLocation();
+  for (auto &arg : pai->getArgumentOperands()) {
+    unsigned calleeArgumentIndex = site.getCalleeArgIndex(arg);
+    assert(calleeArgumentIndex >= calleeConv.getSILArgIndexOfFirstParam());
+    auto paramInfo = calleeConv.getParamInfoForSILArg(calleeArgumentIndex);
+    deallocPartialApplyCapturedArg(builder, loc, arg.get(), paramInfo);
+  }
+}
+
 AbstractFunctionDecl *swift::getBaseMethod(AbstractFunctionDecl *FD) {
   while (FD->getOverriddenDecl()) {
     FD = FD->getOverriddenDecl();