Merge pull request #62592 from gottesmm/pr-cbfd912fd4d58ff272d3f1fe4ceb8cf02589da9a

[sil] Add a new attribute called @closureCaptured to SILFunctionArguments that are closure capture arguments.

Author: gottesmm

include/swift/SIL/SILArgument.h

@@ -332,18 +332,21 @@ class SILPhiArgument : public SILArgument {
 class SILFunctionArgument : public SILArgument {
   friend class SILBasicBlock;
 
-  bool noImplicitCopy = false;
-  LifetimeAnnotation lifetimeAnnotation = LifetimeAnnotation::None;
+  USE_SHARED_UINT32;
 
   SILFunctionArgument(
       SILBasicBlock *parentBlock, SILType type,
       ValueOwnershipKind ownershipKind, const ValueDecl *decl = nullptr,
       bool isNoImplicitCopy = false,
-      LifetimeAnnotation lifetimeAnnotation = LifetimeAnnotation::None)
+      LifetimeAnnotation lifetimeAnnotation = LifetimeAnnotation::None,
+      bool isCapture = false)
       : SILArgument(ValueKind::SILFunctionArgument, parentBlock, type,
-                    ownershipKind, decl),
-        noImplicitCopy(isNoImplicitCopy),
-        lifetimeAnnotation(lifetimeAnnotation) {}
+                    ownershipKind, decl) {
+    sharedUInt32().SILFunctionArgument.noImplicitCopy = isNoImplicitCopy;
+    sharedUInt32().SILFunctionArgument.lifetimeAnnotation = lifetimeAnnotation;
+    sharedUInt32().SILFunctionArgument.closureCapture = isCapture;
+  }
+
   // A special constructor, only intended for use in
   // SILBasicBlock::replaceFunctionArg.
   explicit SILFunctionArgument(SILType type, ValueOwnershipKind ownershipKind,
@@ -352,16 +355,28 @@ class SILFunctionArgument : public SILArgument {
   }
 
 public:
-  bool isNoImplicitCopy() const { return noImplicitCopy; }
+  bool isNoImplicitCopy() const {
+    return sharedUInt32().SILFunctionArgument.noImplicitCopy;
+  }
+
+  void setNoImplicitCopy(bool newValue) {
+    sharedUInt32().SILFunctionArgument.noImplicitCopy = newValue;
+  }
 
-  void setNoImplicitCopy(bool newValue) { noImplicitCopy = newValue; }
+  bool isClosureCapture() const {
+    return sharedUInt32().SILFunctionArgument.closureCapture;
+  }
+  void setClosureCapture(bool newValue) {
+    sharedUInt32().SILFunctionArgument.closureCapture = newValue;
+  }
 
   LifetimeAnnotation getLifetimeAnnotation() const {
-    return lifetimeAnnotation;
+    return LifetimeAnnotation::Case(
+        sharedUInt32().SILFunctionArgument.lifetimeAnnotation);
   }
 
   void setLifetimeAnnotation(LifetimeAnnotation newValue) {
-    lifetimeAnnotation = newValue;
+    sharedUInt32().SILFunctionArgument.lifetimeAnnotation = newValue;
   }
 
   Lifetime getLifetime() const {
@@ -390,6 +405,16 @@ class SILFunctionArgument : public SILArgument {
     return getArgumentConvention() == convention;
   }
 
+  /// Copy all flags stored in this->sharedUInt32() into arg.
+  ///
+  /// By using this API, cloners can be sure they are updated for the addition
+  /// of further flags.
+  void copyFlags(SILFunctionArgument *arg) {
+    setNoImplicitCopy(arg->isNoImplicitCopy());
+    setLifetimeAnnotation(arg->getLifetimeAnnotation());
+    setClosureCapture(arg->isClosureCapture());
+  }
+
   static bool classof(const SILInstruction *) = delete;
   static bool classof(const SILUndef *) = delete;
   static bool classof(SILNodePointer node) {

include/swift/SIL/SILNode.h

@@ -267,8 +267,10 @@ class alignas(8) SILNode :
     SHARED_FIELD(FloatLiteralInst, uint32_t numBits);
     SHARED_FIELD(StringLiteralInst, uint32_t length);
     SHARED_FIELD(PointerToAddressInst, uint32_t alignment);
+    SHARED_FIELD(SILFunctionArgument, uint32_t noImplicitCopy : 1,
+                 lifetimeAnnotation : 2, closureCapture : 1);
 
-  // Do not use `_sharedUInt32_private` outside of SILNode.
+    // Do not use `_sharedUInt32_private` outside of SILNode.
   } _sharedUInt32_private;
 
   static_assert(sizeof(SharedUInt32Fields) == sizeof(uint32_t),

lib/SIL/IR/SILBasicBlock.cpp

@@ -136,8 +136,7 @@ void SILBasicBlock::cloneArgumentList(SILBasicBlock *Other) {
     for (auto *FuncArg : Other->getSILFunctionArguments()) {
       auto *NewArg =
           createFunctionArgument(FuncArg->getType(), FuncArg->getDecl());
-      NewArg->setNoImplicitCopy(FuncArg->isNoImplicitCopy());
-      NewArg->setLifetimeAnnotation(FuncArg->getLifetimeAnnotation());
+      NewArg->copyFlags(FuncArg);
     }
     return;
   }

lib/SIL/IR/SILPrinter.cpp

@@ -160,6 +160,7 @@ struct SILValuePrinterInfo {
   Optional<ValueOwnershipKind> OwnershipKind;
   bool IsNoImplicitCopy = false;
   LifetimeAnnotation Lifetime = LifetimeAnnotation::None;
+  bool IsCapture = false;
 
   SILValuePrinterInfo(ID ValueID) : ValueID(ValueID), Type(), OwnershipKind() {}
   SILValuePrinterInfo(ID ValueID, SILType Type)
@@ -169,13 +170,15 @@ struct SILValuePrinterInfo {
       : ValueID(ValueID), Type(Type), OwnershipKind(OwnershipKind) {}
   SILValuePrinterInfo(ID ValueID, SILType Type,
                       ValueOwnershipKind OwnershipKind, bool IsNoImplicitCopy,
-                      LifetimeAnnotation Lifetime)
+                      LifetimeAnnotation Lifetime, bool IsCapture)
       : ValueID(ValueID), Type(Type), OwnershipKind(OwnershipKind),
-        IsNoImplicitCopy(IsNoImplicitCopy), Lifetime(Lifetime) {}
+        IsNoImplicitCopy(IsNoImplicitCopy), Lifetime(Lifetime),
+        IsCapture(IsCapture) {}
   SILValuePrinterInfo(ID ValueID, SILType Type, bool IsNoImplicitCopy,
-                      LifetimeAnnotation Lifetime)
+                      LifetimeAnnotation Lifetime, bool IsCapture)
       : ValueID(ValueID), Type(Type), OwnershipKind(),
-        IsNoImplicitCopy(IsNoImplicitCopy), Lifetime(Lifetime) {}
+        IsNoImplicitCopy(IsNoImplicitCopy), Lifetime(Lifetime),
+        IsCapture(IsCapture) {}
 };
 
 /// Return the fully qualified dotted path for DeclContext.
@@ -659,6 +662,8 @@ class SILPrinter : public SILInstructionVisitor<SILPrinter> {
       *this << "@_lexical ";
       break;
     }
+    if (i.IsCapture)
+      *this << "@closureCapture ";
     if (i.OwnershipKind && *i.OwnershipKind != OwnershipKind::None) {
       *this << "@" << i.OwnershipKind.value() << " ";
     }
@@ -693,14 +698,18 @@ class SILPrinter : public SILInstructionVisitor<SILPrinter> {
   }
   SILValuePrinterInfo getIDAndType(SILFunctionArgument *arg) {
     return {Ctx.getID(arg), arg->getType(), arg->isNoImplicitCopy(),
-            arg->getLifetimeAnnotation()};
+            arg->getLifetimeAnnotation(), arg->isClosureCapture()};
   }
   SILValuePrinterInfo getIDAndTypeAndOwnership(SILValue V) {
     return {Ctx.getID(V), V ? V->getType() : SILType(), V->getOwnershipKind()};
   }
   SILValuePrinterInfo getIDAndTypeAndOwnership(SILFunctionArgument *arg) {
-    return {Ctx.getID(arg), arg->getType(), arg->getOwnershipKind(),
-            arg->isNoImplicitCopy(), arg->getLifetimeAnnotation()};
+    return {Ctx.getID(arg),
+            arg->getType(),
+            arg->getOwnershipKind(),
+            arg->isNoImplicitCopy(),
+            arg->getLifetimeAnnotation(),
+            arg->isClosureCapture()};
   }
 
   //===--------------------------------------------------------------------===//

lib/SIL/Parser/ParseSIL.cpp

@@ -6495,16 +6495,20 @@ bool SILParser::parseSILBasicBlock(SILBuilder &B) {
           return true;
 
         bool foundNoImplicitCopy = false;
+        bool foundClosureCapture = false;
         bool foundLexical = false;
         bool foundEagerMove = false;
-        while (auto attributeName = parseOptionalAttribute(
-                   {"noImplicitCopy", "_lexical", "_eagerMove"})) {
+        while (auto attributeName =
+                   parseOptionalAttribute({"noImplicitCopy", "_lexical",
+                                           "_eagerMove", "closureCapture"})) {
           if (*attributeName == "noImplicitCopy")
             foundNoImplicitCopy = true;
           else if (*attributeName == "_lexical")
             foundLexical = true;
           else if (*attributeName == "_eagerMove")
             foundEagerMove = true;
+          else if (*attributeName == "closureCapture")
+            foundClosureCapture = true;
           else {
             llvm_unreachable("Unexpected attribute!");
           }
@@ -6533,6 +6537,7 @@ bool SILParser::parseSILBasicBlock(SILBuilder &B) {
         if (IsEntry) {
           auto *fArg = BB->createFunctionArgument(Ty);
           fArg->setNoImplicitCopy(foundNoImplicitCopy);
+          fArg->setClosureCapture(foundClosureCapture);
           fArg->setLifetimeAnnotation(lifetime);
           Arg = fArg;
 

lib/SILGen/SILGenBuilder.cpp

@@ -436,13 +436,15 @@ ManagedValue SILGenBuilder::createLoadCopy(SILLocation loc, ManagedValue v,
 static ManagedValue createInputFunctionArgument(
     SILGenBuilder &B, SILType type, SILLocation loc, ValueDecl *decl = nullptr,
     bool isNoImplicitCopy = false,
-    LifetimeAnnotation lifetimeAnnotation = LifetimeAnnotation::None) {
+    LifetimeAnnotation lifetimeAnnotation = LifetimeAnnotation::None,
+    bool isClosureCapture = false) {
   auto &SGF = B.getSILGenFunction();
   SILFunction &F = B.getFunction();
   assert((F.isBare() || decl) &&
          "Function arguments of non-bare functions must have a decl");
   auto *arg = F.begin()->createFunctionArgument(type, decl);
   arg->setNoImplicitCopy(isNoImplicitCopy);
+  arg->setClosureCapture(isClosureCapture);
   arg->setLifetimeAnnotation(lifetimeAnnotation);
   switch (arg->getArgumentConvention()) {
   case SILArgumentConvention::Indirect_In_Guaranteed:
@@ -477,9 +479,10 @@ static ManagedValue createInputFunctionArgument(
 
 ManagedValue SILGenBuilder::createInputFunctionArgument(
     SILType type, ValueDecl *decl, bool isNoImplicitCopy,
-    LifetimeAnnotation lifetimeAnnotation) {
+    LifetimeAnnotation lifetimeAnnotation, bool isClosureCapture) {
   return ::createInputFunctionArgument(*this, type, SILLocation(decl), decl,
-                                       isNoImplicitCopy, lifetimeAnnotation);
+                                       isNoImplicitCopy, lifetimeAnnotation,
+                                       isClosureCapture);
 }
 
 ManagedValue

lib/SILGen/SILGenBuilder.h

@@ -228,7 +228,8 @@ class SILGenBuilder : public SILBuilder {
   /// function argument for an out parameter.
   ManagedValue createInputFunctionArgument(
       SILType type, ValueDecl *decl, bool isNoImplicitCopy = false,
-      LifetimeAnnotation lifetimeAnnotation = LifetimeAnnotation::None);
+      LifetimeAnnotation lifetimeAnnotation = LifetimeAnnotation::None,
+      bool isClosureCapture = false);
 
   /// Create a SILArgument for an input parameter. Uses \p loc to create any
   /// copies necessary. Asserts if used to create a function argument for an out

lib/SILGen/SILGenProlog.cpp

@@ -520,8 +520,10 @@ static void emitCaptureArguments(SILGenFunction &SGF,
     auto &lowering = SGF.getTypeLowering(type);
     // Constant decls are captured by value.
     SILType ty = lowering.getLoweredType();
-    ManagedValue val = ManagedValue::forUnmanaged(
-        SGF.F.begin()->createFunctionArgument(ty, VD));
+    auto *arg = SGF.F.begin()->createFunctionArgument(ty, VD);
+    arg->setClosureCapture(true);
+
+    ManagedValue val = ManagedValue::forUnmanaged(arg);
 
     // If the original variable was settable, then Sema will have treated the
     // VarDecl as an lvalue, even in the closure's use.  As such, we need to
@@ -575,8 +577,9 @@ static void emitCaptureArguments(SILGenFunction &SGF,
         SGF.SGM.Types.getLoweredRValueType(TypeExpansionContext::minimal(),
                                            type),
         SGF.F.getGenericEnvironment(), /*mutable*/ true);
-    SILValue box = SGF.F.begin()->createFunctionArgument(
+    auto *box = SGF.F.begin()->createFunctionArgument(
         SILType::getPrimitiveObjectType(boxTy), VD);
+    box->setClosureCapture(true);
     SILValue addr = SGF.B.createProjectBox(VD, box, 0);
     SGF.VarLocs[VD] = SILGenFunction::VarLoc::get(addr, box);
     SILDebugVariable DbgVar(VD->isLet(), ArgNo);
@@ -595,7 +598,9 @@ static void emitCaptureArguments(SILGenFunction &SGF,
     if (isInOut || SGF.SGM.M.useLoweredAddresses()) {
       ty = ty.getAddressType();
     }
-    SILValue arg = SGF.F.begin()->createFunctionArgument(ty, VD);
+    auto *fArg = SGF.F.begin()->createFunctionArgument(ty, VD);
+    fArg->setClosureCapture(true);
+    SILValue arg = SILValue(fArg);
     if (isInOut && (ty.isMoveOnly() && !ty.isMoveOnlyWrapped())) {
       arg = SGF.B.createMarkMustCheckInst(
           Loc, arg, MarkMustCheckInst::CheckKind::NoImplicitCopy);

lib/SILOptimizer/FunctionSignatureTransforms/ArgumentExplosionTransform.cpp

@@ -393,8 +393,7 @@ void FunctionSignatureTransform::ArgumentExplosionFinalizeOptimizedFunction() {
       auto *Argument =
           BB->insertFunctionArgument(ArgOffset, Node->getType(), OwnershipKind,
                                      BB->getArgument(OldArgIndex)->getDecl());
-      Argument->setNoImplicitCopy(AD.Arg->isNoImplicitCopy());
-      Argument->setLifetimeAnnotation(AD.Arg->getLifetimeAnnotation());
+      Argument->copyFlags(AD.Arg);
       LeafValues.push_back(Argument);
       TransformDescriptor.AIM[TotalArgIndex - 1] = AD.Index;
       ++ArgOffset;

lib/SILOptimizer/FunctionSignatureTransforms/ExistentialTransform.cpp

@@ -165,10 +165,7 @@ void ExistentialSpecializerCloner::cloneArguments(
           LoweredTy.getCategoryType(ArgDesc.Arg->getType().getCategory());
       auto *NewArg =
           ClonedEntryBB->createFunctionArgument(MappedTy, ArgDesc.Decl);
-      NewArg->setNoImplicitCopy(ArgDesc.Arg->isNoImplicitCopy());
-      NewArg->setLifetimeAnnotation(ArgDesc.Arg->getLifetimeAnnotation());
-      NewArg->setOwnershipKind(ValueOwnershipKind(
-          NewF, MappedTy, ArgDesc.Arg->getArgumentConvention()));
+      NewArg->copyFlags(ArgDesc.Arg);
       entryArgs.push_back(NewArg);
       continue;
     }
@@ -182,8 +179,7 @@ void ExistentialSpecializerCloner::cloneArguments(
         GenericSILType, ArgDesc.Decl,
         ValueOwnershipKind(NewF, GenericSILType,
                            ArgDesc.Arg->getArgumentConvention()));
-    NewArg->setNoImplicitCopy(ArgDesc.Arg->isNoImplicitCopy());
-    NewArg->setLifetimeAnnotation(ArgDesc.Arg->getLifetimeAnnotation());
+    NewArg->copyFlags(ArgDesc.Arg);
     // Determine the Conformances.
     SILType ExistentialType = ArgDesc.Arg->getType().getObjectType();
     CanType OpenedType = NewArg->getType().getASTType();
@@ -408,8 +404,7 @@ void ExistentialTransform::populateThunkBody() {
     auto argumentType = ArgDesc.Arg->getType();
     auto *NewArg =
         ThunkBody->createFunctionArgument(argumentType, ArgDesc.Decl);
-    NewArg->setNoImplicitCopy(ArgDesc.Arg->isNoImplicitCopy());
-    NewArg->setLifetimeAnnotation(ArgDesc.Arg->getLifetimeAnnotation());
+    NewArg->copyFlags(ArgDesc.Arg);
   }
 
   /// Builder to add new instructions in the Thunk.