Merge pull request #64048 from rjmccall/variadic-pack-expansion-arguments

rjmccall · web-flow · commit b3b90a82bee5 · 2023-03-03T11:46:00.000-05:00
Implement the emission of pack expansion arguments in SILGen
diff --git a/include/swift/AST/Types.h b/include/swift/AST/Types.h
@@ -753,6 +753,15 @@ class alignas(1 << TypeAlignInBits) TypeBase
   /// Return the root generic parameter of this type parameter type.
   GenericTypeParamType *getRootGenericParam();
 
+  /// Given that this type is the result of substituting a pack parameter,
+  /// return it as a the pack type.  We want to have a representational
+  /// invariant that these substitutions always produce PackTypes; when
+  /// that's in place, we can replace all uses of this function with
+  /// `->castTo<PackType>()`.  In the meantime, it's permitted for these
+  /// substitutions to be unadorned pack parameters or archetypes, which
+  /// this function will wrap into a pack containing a singleton expansion.
+  PackType *getPackSubstitutionAsPackType();
+
   /// Determines whether this type is an lvalue. This includes both straight
   /// lvalue types as well as tuples or optionals of lvalues.
   bool hasLValueType() {
@@ -6732,6 +6741,11 @@ class PackType final : public TypeBase, public llvm::FoldingSetNode,
                        TypeArrayView<GenericTypeParamType> params,
                        ArrayRef<Type> args);
 
+  /// Given a type T, which must be a pack parameter, a member type
+  /// of a pack parameter, or a pack archetype, construct the type
+  /// Pack{repeat each T}.
+  static PackType *getSingletonPackExpansion(Type packParameter);
+
 public:
   /// Retrieves the number of elements in this pack.
   unsigned getNumElements() const { return Bits.PackType.Count; }
diff --git a/lib/AST/ASTContext.cpp b/lib/AST/ASTContext.cpp
@@ -3302,6 +3302,11 @@ PackType *PackType::getEmpty(const ASTContext &C) {
   return cast<PackType>(CanType(C.TheEmptyPackType));
 }
 
+PackType *PackType::getSingletonPackExpansion(Type param) {
+  assert(param->isParameterPack() || param->is<PackArchetypeType>());
+  return get(param->getASTContext(), {PackExpansionType::get(param, param)});
+}
+
 CanPackType CanPackType::get(const ASTContext &C, ArrayRef<CanType> elements) {
   SmallVector<Type, 8> ncElements(elements.begin(), elements.end());
   return CanPackType(PackType::get(C, ncElements));
diff --git a/lib/AST/GenericEnvironment.cpp b/lib/AST/GenericEnvironment.cpp
@@ -177,7 +177,7 @@ void GenericEnvironment::forEachPackElementBinding(
     auto *elementArchetype =
         mapTypeIntoContext(*elementIt++)->castTo<ElementArchetypeType>();
     auto *packSubstitution = maybeApplyOuterContextSubstitutions(genericParam)
-                                 ->template castTo<PackType>();
+      ->getPackSubstitutionAsPackType();
     function(elementArchetype, packSubstitution);
   });
 
diff --git a/lib/AST/ParameterPack.cpp b/lib/AST/ParameterPack.cpp
@@ -86,6 +86,14 @@ bool GenericTypeParamType::isParameterPack() const {
          GenericTypeParamType::TYPE_SEQUENCE_BIT;
 }
 
+PackType *TypeBase::getPackSubstitutionAsPackType() {
+  if (auto pack = getAs<PackType>()) {
+    return pack;
+  } else {
+    return PackType::getSingletonPackExpansion(this);
+  }
+}
+
 /// G<{X1, ..., Xn}, {Y1, ..., Yn}>... => {G<X1, Y1>, ..., G<Xn, Yn>}...
 PackExpansionType *PackExpansionType::expand() {
   auto countType = getCountType();
diff --git a/lib/AST/Type.cpp b/lib/AST/Type.cpp
@@ -729,6 +729,11 @@ static bool isLegalFormalType(CanType type) {
     return isLegalFormalType(objectType);
   }
 
+  // Expansions are legal if their pattern type is legal.
+  if (auto expansionType = dyn_cast<PackExpansionType>(type)) {
+    return isLegalFormalType(expansionType.getPatternType());
+  }
+
   return true;
 }
 
diff --git a/lib/SIL/IR/TypeLowering.cpp b/lib/SIL/IR/TypeLowering.cpp
@@ -331,11 +331,13 @@ namespace {
     RetTy visitPackExpansionType(CanPackExpansionType type,
                                  AbstractionPattern origType,
                                  IsTypeExpansionSensitive_t isSensitive) {
-      return asImpl().handleAddressOnly(type, {IsNotTrivial, IsFixedABI,
-                                               IsAddressOnly, IsNotResilient,
-                                               isSensitive,
-                                               DoesNotHaveRawPointer,
-                                               IsLexical});
+      RecursiveProperties props;
+      props.setAddressOnly();
+      props.addSubobject(classifyType(origType.getPackExpansionPatternType(),
+                                      type.getPatternType(),
+                                      TC, Expansion));
+      props = mergeIsTypeExpansionSensitive(isSensitive, props);
+      return asImpl().handleAddressOnly(type, props);
     }
 
     RetTy visitBuiltinRawPointerType(CanBuiltinRawPointerType type,
@@ -2193,18 +2195,20 @@ namespace {
     TypeLowering *visitPackType(CanPackType packType,
                                 AbstractionPattern origType,
                                 IsTypeExpansionSensitive_t isSensitive) {
-      RecursiveProperties properties;
-      properties.setAddressOnly();
-      properties = mergeIsTypeExpansionSensitive(isSensitive, properties);
-
-      return handleAddressOnly(packType, properties);
+      llvm_unreachable("shouldn't get here with an unlowered type");
     }
 
     TypeLowering *visitSILPackType(CanSILPackType packType,
                                    AbstractionPattern origType,
                                    IsTypeExpansionSensitive_t isSensitive) {
       RecursiveProperties properties;
       properties.setAddressOnly();
+      for (auto i : indices(packType.getElementTypes())) {
+        auto &eltLowering =
+          TC.getTypeLowering(packType->getSILElementType(i),
+                             Expansion);
+        properties.addSubobject(eltLowering.getRecursiveProperties());
+      }
       properties = mergeIsTypeExpansionSensitive(isSensitive, properties);
 
       return handleAddressOnly(packType, properties);
@@ -2215,6 +2219,11 @@ namespace {
                                          IsTypeExpansionSensitive_t isSensitive) {
       RecursiveProperties properties;
       properties.setAddressOnly();
+      auto &patternLowering =
+        TC.getTypeLowering(origType.getPackExpansionPatternType(),
+                           packExpansionType.getPatternType(),
+                           Expansion);
+      properties.addSubobject(patternLowering.getRecursiveProperties());
       properties = mergeIsTypeExpansionSensitive(isSensitive, properties);
 
       return handleAddressOnly(packExpansionType, properties);
diff --git a/lib/SILGen/SILGenApply.cpp b/lib/SILGen/SILGenApply.cpp
@@ -3757,7 +3757,7 @@ class ArgEmitter {
                              expectedEltTy, consumed,
                              packAddr, formalPackType, i);
       }
-      if (cleanup.isValid()) eltCleanups.push_back(cleanup);
+      if (consumed && cleanup.isValid()) eltCleanups.push_back(cleanup);
     }
 
     if (!consumed) {
@@ -3788,9 +3788,6 @@ class ArgEmitter {
                                           SILValue packAddr,
                                           CanPackType formalPackType,
                                           unsigned packComponentIndex) {
-    auto expansionLoweredType =
-      expectedParamType.castTo<PackExpansionType>();
-
     // TODO: we'll need to handle already-emitted packs for things like
     // subscripts
     assert(arg.isExpr() && "emitting a non-expression pack expansion");
@@ -3830,20 +3827,26 @@ class ArgEmitter {
       // elements in this slice of the pack (then pop it before we exit
       // this scope).
 
-      // Turn pack archetypes in the lowered pack expansion type into
-      // opened element archetypes.  This should work fine on SIL types
-      // since we're not changing any interesting structure.
-      auto expectedLoweredElementType =
-        openedElementEnv->mapPackTypeIntoElementContext(
-          expansionLoweredType.getPatternType())->getCanonicalType();
+      // Turn pack archetypes in the pattern type of the lowered pack
+      // expansion type into opened element archetypes.  These AST-level
+      // manipulations should work fine on SIL types since we're not
+      // changing any interesting structure.
+      SILType expectedElementType = [&] {
+        auto loweredPatternType =
+          expectedParamType.castTo<PackExpansionType>().getPatternType();
+        auto loweredElementType =
+          openedElementEnv->mapPackTypeIntoElementContext(
+            loweredPatternType->mapTypeOutOfContext())->getCanonicalType();
+        return SILType::getPrimitiveAddressType(loweredElementType);
+      }();
 
       // Project the tuple element.  This projection uses the
       // pack expansion index because the tuple is only for the
       // projection elements.
       auto eltAddr =
         SGF.B.createTuplePackElementAddr(expansionExpr,
                                          packExpansionIndex, tupleAddr,
-          SILType::getPrimitiveAddressType(expectedLoweredElementType));
+                                         expectedElementType);
       auto &eltTL = SGF.getTypeLowering(eltAddr->getType());
 
       // Evaluate the pattern expression into that address.
diff --git a/lib/SILGen/SILGenBuilder.cpp b/lib/SILGen/SILGenBuilder.cpp
@@ -475,9 +475,11 @@ static ManagedValue createInputFunctionArgument(
     return ManagedValue::forUnmanaged(arg).copy(SGF, loc);
 
   case SILArgumentConvention::Direct_Owned:
-  case SILArgumentConvention::Pack_Owned:
     return SGF.emitManagedRValueWithCleanup(arg);
 
+  case SILArgumentConvention::Pack_Owned:
+    return SGF.emitManagedPackWithCleanup(arg);
+
   case SILArgumentConvention::Indirect_In:
     if (SGF.silConv.useLoweredAddresses())
       return SGF.emitManagedBufferWithCleanup(arg);
diff --git a/lib/SILGen/SILGenExpr.cpp b/lib/SILGen/SILGenExpr.cpp
@@ -1531,7 +1531,14 @@ RValueEmitter::visitPackExpansionExpr(PackExpansionExpr *E,
 
 RValue
 RValueEmitter::visitPackElementExpr(PackElementExpr *E, SGFContext C) {
-  llvm_unreachable("not implemented for PackElementExpr");
+  FormalEvaluationScope scope(SGF);
+
+  LValue lv = SGF.emitLValue(E, SGFAccessKind::OwnedObjectRead);
+
+  // Otherwise, we can't load at +0 without further analysis, since the formal
+  // access into the lvalue will end immediately.
+  return SGF.emitLoadOfLValue(E, std::move(lv),
+                              C.withFollowingSideEffects());
 }
 
 RValue
diff --git a/lib/SILGen/SILGenFunction.h b/lib/SILGen/SILGenFunction.h
@@ -559,6 +559,14 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   /// points.
   SILValue ExpectedExecutor;
 
+  /// The pack index value for the innermost pack expansion.
+  SILValue InnermostPackIndex;
+
+  SILValue getInnermostPackIndex() const {
+    assert(InnermostPackIndex && "not inside a pack expansion!");
+    return InnermostPackIndex;
+  }
+
   /// True if 'return' without an operand or falling off the end of the current
   /// function is valid.
   bool allowsVoidReturn() const { return ReturnDest.getBlock()->args_empty(); }
@@ -1666,6 +1674,10 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   ManagedValue emitManagedBufferWithCleanup(SILValue addr,
                                             const TypeLowering &lowering);
 
+  ManagedValue emitManagedPackWithCleanup(SILValue addr,
+                                          CanPackType formalPackType
+                                            = CanPackType());
+
   ManagedValue emitFormalAccessManagedRValueWithCleanup(SILLocation loc,
                                                         SILValue value);
   ManagedValue emitFormalAccessManagedBufferWithCleanup(SILLocation loc,
@@ -2490,14 +2502,18 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   ///   starting at index limitWithinComponent - 1
   /// \param emitBody - a function that will be called to emit the body of
   ///   the loop. It's okay if this has paths that exit the body of the loop,
-  ///   but it should leave the insertion point set at the end. Will be
-  ///   called within a cleanups scope. The first parameter is the current
-  ///   index within the expansion component, a value of type Builtin.Word.
-  ///   The second parameter is that index as a pack indexing instruction
-  ///   that indexes into packs with the shape of the pack expasion.
-  ///   The third parameter is the current pack index within the overall
-  ///   pack, a pack indexing instruction that indexes into packs with the
-  ///   shape of formalPackType.
+  ///   but it should leave the insertion point set at the end.
+  ///
+  ///   The first parameter is the current index within the expansion
+  ///   component, a value of type Builtin.Word.  The second parameter is
+  ///   that index as a pack indexing instruction that indexes into packs
+  ///   with the shape of the pack expasion.  The third parameter is the
+  ///   current pack index within the overall pack, a pack indexing instruction
+  ///   that indexes into packs with the shape of formalPackType.
+  ///
+  ///   This function will be called within a cleanups scope and with
+  ///   InnermostPackIndex set to the pack expansion index value (the
+  ///   second parameter).
   void emitDynamicPackLoop(SILLocation loc,
                            CanPackType formalPackType,
                            unsigned componentIndex,
diff --git a/lib/SILGen/SILGenLValue.cpp b/lib/SILGen/SILGenLValue.cpp
@@ -309,6 +309,8 @@ class LLVM_LIBRARY_VISIBILITY SILGenLValue
                           LValueOptions options);
   LValue visitOpaqueValueExpr(OpaqueValueExpr *e, SGFAccessKind accessKind,
                               LValueOptions options);
+  LValue visitPackElementExpr(PackElementExpr *e, SGFAccessKind accessKind,
+                              LValueOptions options);
 
   // Nodes that make up components of lvalue paths
   
@@ -3313,6 +3315,42 @@ LValue SILGenLValue::visitOpaqueValueExpr(OpaqueValueExpr *e,
   return lv;
 }
 
+LValue SILGenLValue::visitPackElementExpr(PackElementExpr *e,
+                                          SGFAccessKind accessKind,
+                                          LValueOptions options) {
+  auto refExpr = e->getPackRefExpr();
+
+  auto substFormalType = e->getType()->getCanonicalType();
+
+  if (auto declRefExpr = dyn_cast<DeclRefExpr>(refExpr)) {
+    if (auto var = dyn_cast<VarDecl>(declRefExpr->getDecl())) {
+      auto origFormalType = SGF.SGM.Types.getAbstractionPattern(var);
+      auto elementTy =
+        SGF.getLoweredType(origFormalType, substFormalType).getAddressType();
+      SILValue packAddr =
+        SGF.emitAddressOfLocalVarDecl(declRefExpr, var,
+                                      substFormalType,
+                                      accessKind)
+           .getLValueAddress();
+      auto packIndex = SGF.getInnermostPackIndex();
+      auto elementAddr =
+        SGF.B.createPackElementGet(e, packIndex, packAddr, elementTy);
+      return LValue::forAddress(accessKind,
+                                ManagedValue::forLValue(elementAddr),
+                                /*access enforcement*/ None,
+                                origFormalType, substFormalType);
+    }
+  }
+
+  SGF.SGM.diagnose(refExpr, diag::not_implemented,
+                   "emission of 'each' for this kind of expression");
+  auto loweredTy = SGF.getLoweredType(substFormalType).getAddressType();
+  auto fakeAddr = ManagedValue::forLValue(SILUndef::get(loweredTy, SGF.F));
+  return LValue::forAddress(accessKind, fakeAddr, /*access enforcement*/ None,
+                            AbstractionPattern(substFormalType),
+                            substFormalType);
+}
+
 LValue SILGenLValue::visitDotSyntaxBaseIgnoredExpr(DotSyntaxBaseIgnoredExpr *e,
                                                    SGFAccessKind accessKind,
                                                    LValueOptions options) {
diff --git a/lib/SILGen/SILGenPack.cpp b/lib/SILGen/SILGenPack.cpp
diff --git a/test/SILGen/variadic-generics.swift b/test/SILGen/variadic-generics.swift

Original file line number	Diff line number	Diff line change
`@@ -729,6 +729,11 @@ static bool isLegalFormalType(CanType type) {`
`729`	`729`	`return isLegalFormalType(objectType);`
`730`	`730`	`}`
`731`	`731`
	`732`	`+ // Expansions are legal if their pattern type is legal.`
	`733`	`+ if (auto expansionType = dyn_cast<PackExpansionType>(type)) {`
	`734`	`+ return isLegalFormalType(expansionType.getPatternType());`
	`735`	`+ }`
	`736`	`+`
`732`	`737`	`return true;`
`733`	`738`	`}`
`734`	`739`