Implement TypeBase::isTypeSequenceParameter

CodaFi · CodaFi · commit d4b8f1df961b · 2021-12-16T01:16:45.000-08:00
Behaves much like isTypeParameter but specifically checks for the type sequence bits. Also, add the type sequence bits as a recursive type property.
diff --git a/include/swift/AST/Types.h b/include/swift/AST/Types.h
@@ -159,7 +159,11 @@ class RecursiveTypeProperties {
     /// This type contains a type placeholder.
     HasPlaceholder       = 0x800,
 
-    Last_Property = HasPlaceholder
+    /// This type contains a generic type parameter that is declared as a
+    /// type sequence
+    HasTypeSequence = 0x1000,
+
+    Last_Property = HasTypeSequence
   };
   enum { BitWidth = countBitsUsed(Property::Last_Property) };
 
@@ -217,6 +221,8 @@ class RecursiveTypeProperties {
   /// Does a type with these properties structurally contain a placeholder?
   bool hasPlaceholder() const { return Bits & HasPlaceholder; }
 
+  bool hasTypeSequence() const { return Bits & HasTypeSequence; }
+
   /// Returns the set of properties present in either set.
   friend RecursiveTypeProperties operator|(Property lhs, Property rhs) {
     return RecursiveTypeProperties(unsigned(lhs) | unsigned(rhs));
@@ -605,6 +611,10 @@ class alignas(1 << TypeAlignInBits) TypeBase
     return getRecursiveProperties().hasOpenedExistential();
   }
 
+  bool hasTypeSequence() const {
+    return getRecursiveProperties().hasTypeSequence();
+  }
+
   /// Determine whether the type involves the given opened existential
   /// archetype.
   bool hasOpenedExistential(OpenedArchetypeType *opened);
@@ -655,6 +665,13 @@ class alignas(1 << TypeAlignInBits) TypeBase
   /// whether a type parameter exists at any position.
   bool isTypeParameter();
 
+  /// Determine whether this type is a type sequence parameter, which is
+  /// either a GenericTypeParamType or a DependentMemberType.
+  ///
+  /// Like \c isTypeParameter, this routine will return \c false for types that
+  /// include type parameters in nested positions e.g. \c X<T...>.
+  bool isTypeSequenceParameter();
+
   /// Determine whether this type can dynamically be an optional type.
   ///
   /// \param includeExistential Whether an existential type should be considered
@@ -3431,6 +3448,7 @@ struct ParameterListInfo {
   SmallBitVector propertyWrappers;
   SmallBitVector implicitSelfCapture;
   SmallBitVector inheritActorContext;
+  SmallBitVector variadicGenerics;
 
 public:
   ParameterListInfo() { }
@@ -3460,6 +3478,8 @@ struct ParameterListInfo {
   /// Whether there is any contextual information set on this parameter list.
   bool anyContextualInfo() const;
 
+  bool isVariadicGenericParameter(unsigned paramIdx) const;
+
   /// Retrieve the number of non-defaulted parameters.
   unsigned numNonDefaultedParameters() const {
     return defaultArguments.count();
@@ -5883,15 +5903,15 @@ class GenericTypeParamType : public SubstitutableType {
 private:
   friend class GenericTypeParamDecl;
 
-  explicit GenericTypeParamType(GenericTypeParamDecl *param)
-    : SubstitutableType(TypeKind::GenericTypeParam, nullptr,
-                        RecursiveTypeProperties::HasTypeParameter),
+  explicit GenericTypeParamType(GenericTypeParamDecl *param,
+                                RecursiveTypeProperties props)
+    : SubstitutableType(TypeKind::GenericTypeParam, nullptr, props),
       ParamOrDepthIndex(param) { }
 
   explicit GenericTypeParamType(bool isTypeSequence, unsigned depth,
-                                unsigned index, const ASTContext &ctx)
-      : SubstitutableType(TypeKind::GenericTypeParam, &ctx,
-                          RecursiveTypeProperties::HasTypeParameter),
+                                unsigned index, RecursiveTypeProperties props,
+                                const ASTContext &ctx)
+      : SubstitutableType(TypeKind::GenericTypeParam, &ctx, props),
         ParamOrDepthIndex(depth << 16 | index |
                           ((isTypeSequence ? 1 : 0) << 30)) {}
 };
@@ -6300,7 +6320,14 @@ inline bool TypeBase::isTypeParameter() {
   return t->is<GenericTypeParamType>();
 }
 
-  return false;
+inline bool TypeBase::isTypeSequenceParameter() {
+  Type t(this);
+
+  while (auto *memberTy = t->getAs<DependentMemberType>())
+    t = memberTy->getBase();
+
+  return t->is<GenericTypeParamType>() &&
+         t->castTo<GenericTypeParamType>()->isTypeSequence();
 }
 
 inline bool TypeBase::isMaterializable() {
diff --git a/include/swift/Sema/ConstraintSystem.h b/include/swift/Sema/ConstraintSystem.h
@@ -363,6 +363,8 @@ class TypeVariableType::Implementation {
   /// Determine whether this type variable represents a subscript result type.
   bool isSubscriptResultType() const;
 
+  bool isTypeSequence() const;
+
   /// Retrieve the representative of the equivalence class to which this
   /// type variable belongs.
   ///
diff --git a/lib/AST/ASTContext.cpp b/lib/AST/ASTContext.cpp
@@ -3495,7 +3495,7 @@ isAnyFunctionTypeCanonical(ArrayRef<AnyFunctionType::Param> params,
 static RecursiveTypeProperties
 getGenericFunctionRecursiveProperties(ArrayRef<AnyFunctionType::Param> params,
                                       Type result) {
-  static_assert(RecursiveTypeProperties::BitWidth == 12,
+  static_assert(RecursiveTypeProperties::BitWidth == 13,
                 "revisit this if you add new recursive type properties");
   RecursiveTypeProperties properties;
 
@@ -3548,6 +3548,8 @@ Type AnyFunctionType::Param::getParameterType(bool forCanonical,
     auto arrayDecl = ctx->getArrayDecl();
     if (!arrayDecl)
       type = ErrorType::get(*ctx);
+    else if (type->is<PackType>())
+      return type;
     else if (forCanonical)
       type = BoundGenericType::get(arrayDecl, Type(), {type});
     else
@@ -3791,8 +3793,12 @@ GenericTypeParamType *GenericTypeParamType::get(bool isTypeSequence,
   if (known != ctx.getImpl().GenericParamTypes.end())
     return known->second;
 
+  RecursiveTypeProperties props = RecursiveTypeProperties::HasTypeParameter;
+  if (isTypeSequence)
+    props |= RecursiveTypeProperties::HasTypeSequence;
+
   auto result = new (ctx, AllocationArena::Permanent)
-      GenericTypeParamType(isTypeSequence, depth, index, ctx);
+      GenericTypeParamType(isTypeSequence, depth, index, props, ctx);
   ctx.getImpl().GenericParamTypes[{depthKey, index}] = result;
   return result;
 }
@@ -4078,7 +4084,7 @@ CanSILFunctionType SILFunctionType::get(
   void *mem = ctx.Allocate(bytes, alignof(SILFunctionType));
 
   RecursiveTypeProperties properties;
-  static_assert(RecursiveTypeProperties::BitWidth == 12,
+  static_assert(RecursiveTypeProperties::BitWidth == 13,
                 "revisit this if you add new recursive type properties");
   for (auto &param : params)
     properties |= param.getInterfaceType()->getRecursiveProperties();
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -4107,7 +4107,10 @@ GenericTypeParamDecl::GenericTypeParamDecl(DeclContext *dc, Identifier name,
   assert(Bits.GenericTypeParamDecl.Index == index && "Truncation");
   Bits.GenericTypeParamDecl.TypeSequence = isTypeSequence;
   auto &ctx = dc->getASTContext();
-  auto type = new (ctx, AllocationArena::Permanent) GenericTypeParamType(this);
+  RecursiveTypeProperties props = RecursiveTypeProperties::HasTypeParameter;
+  if (this->isTypeSequence())
+    props |= RecursiveTypeProperties::HasTypeSequence;
+  auto type = new (ctx, AllocationArena::Permanent) GenericTypeParamType(this, props);
   setInterfaceType(MetatypeType::get(type, ctx));
 }
 
diff --git a/lib/AST/Type.cpp b/lib/AST/Type.cpp
@@ -1037,6 +1037,7 @@ ParameterListInfo::ParameterListInfo(
   propertyWrappers.resize(params.size());
   implicitSelfCapture.resize(params.size());
   inheritActorContext.resize(params.size());
+  variadicGenerics.resize(params.size());
 
   // No parameter owner means no parameter list means no default arguments
   // - hand back the zeroed bitvector.
@@ -1096,6 +1097,11 @@ ParameterListInfo::ParameterListInfo(
     if (param->getAttrs().hasAttribute<InheritActorContextAttr>()) {
       inheritActorContext.set(i);
     }
+
+    if (param->isVariadic() &&
+        param->getVarargBaseTy()->hasTypeSequence()) {
+      variadicGenerics.set(i);
+    }
   }
 }
 
@@ -1130,6 +1136,13 @@ bool ParameterListInfo::anyContextualInfo() const {
   return implicitSelfCapture.any() || inheritActorContext.any();
 }
 
+bool ParameterListInfo::isVariadicGenericParameter(unsigned paramIdx) const {
+  return paramIdx < variadicGenerics.size()
+      ? variadicGenerics[paramIdx]
+      : false;
+}
+
+
 /// Turn a param list into a symbolic and printable representation that does not
 /// include the types, something like (_:, b:, c:)
 std::string swift::getParamListAsString(ArrayRef<AnyFunctionType::Param> params) {
@@ -1360,6 +1373,13 @@ CanType TypeBase::computeCanonicalType() {
     break;
   }
 
+  case TypeKind::PackExpansion: {
+    auto *expansion = cast<PackExpansionType>(this);
+    auto pattern = expansion->getPatternType()->getCanonicalType();
+    Result = PackExpansionType::get(pattern);
+    break;
+  }
+
   case TypeKind::Tuple: {
     TupleType *TT = cast<TupleType>(this);
     assert(TT->getNumElements() != 0 && "Empty tuples are always canonical");
@@ -5061,18 +5081,152 @@ case TypeKind::Id:
     return ParenType::get(Ptr->getASTContext(), underlying->getInOutObjectType(), otherFlags);
   }
 
+  case TypeKind::Pack: {
+    auto pack = cast<PackType>(base);
+    bool anyChanged = false;
+    SmallVector<Type, 4> elements;
+    unsigned Index = 0;
+    for (Type eltTy : pack->getElementTypes()) {
+      Type transformedEltTy = eltTy.transformRec(fn);
+      if (!transformedEltTy)
+        return Type();
+
+      // If nothing has changed, just keep going.
+      if (!anyChanged &&
+          transformedEltTy.getPointer() == eltTy.getPointer()) {
+        ++Index;
+        continue;
+      }
+
+      // If this is the first change we've seen, copy all of the previous
+      // elements.
+      if (!anyChanged) {
+        // Copy all of the previous elements.
+        elements.append(pack->getElementTypes().begin(),
+                        pack->getElementTypes().begin() + Index);
+        anyChanged = true;
+      }
+
+      elements.push_back(transformedEltTy);
+      ++Index;
+    }
+
+    if (!anyChanged)
+      return *this;
+
+    return PackType::get(Ptr->getASTContext(), elements);
+  }
+
+  case TypeKind::PackExpansion: {
+    auto expand = cast<PackExpansionType>(base);
+    struct ExpansionGatherer {
+      llvm::function_ref<Optional<Type>(TypeBase *)> baselineFn;
+      llvm::DenseMap<TypeBase *, PackType *> cache;
+      unsigned maxArity;
+
+    public:
+      ExpansionGatherer(
+          llvm::function_ref<Optional<Type>(TypeBase *)> baselineFn)
+          : baselineFn(baselineFn), maxArity(0) {}
+
+      Optional<Type> operator()(TypeBase *input) {
+        auto remap = baselineFn(input);
+        if (!remap) {
+          return remap;
+        }
+
+        if (input->is<TypeVariableType>()) {
+          if (auto *PT = (*remap)->getAs<PackType>()) {
+            maxArity = std::max(maxArity, PT->getNumElements());
+            cache.insert({input, PT});
+          }
+        } else if (input->isTypeSequenceParameter()) {
+          if (auto *PT = (*remap)->getAs<PackType>()) {
+            maxArity = std::max(maxArity, PT->getNumElements());
+            cache.insert({input, PT});
+          }
+        }
+        return remap;
+      }
+
+      std::pair<llvm::DenseMap<TypeBase *, PackType *>, unsigned>
+      intoExpansions() && {
+        return std::make_pair(cache, maxArity);
+      }
+    };
+
+    // First, substitute down the pattern type to gather the mapping from
+    // contained substitutable types to packs.
+    auto gather = ExpansionGatherer{fn};
+    Type transformedPat = expand->getPatternType().transformRec(gather);
+    if (!transformedPat)
+      return Type();
+
+    if (transformedPat.getPointer() == expand->getPatternType().getPointer())
+      return *this;
+
+    llvm::DenseMap<TypeBase *, PackType *> expansions;
+    unsigned arity;
+    std::tie(expansions, arity) = std::move(gather).intoExpansions();
+    if (expansions.empty()) {
+      // If we didn't find any expansions, either the caller wasn't interested
+      // in expanding this pack, or something has gone wrong. Leave off the
+      // expansion and return the transformed type.
+      return PackExpansionType::get(transformedPat);
+    }
+
+    SmallVector<Type, 8> elts;
+    elts.reserve(arity);
+    // Perform the expansion element-wise according to the maximum arity we
+    // picked up during the gather step above.
+    //
+    // For a pack expansion (F<... T..., U..., ...>) and mapping
+    //
+    //   T... -> <X, Y, Z>
+    //   U... -> <A, B, C>
+    //
+    // The expected expansion is
+    //
+    // <F<... X, A, ...>, F<... Y, B, ...>, F<... Z, C, ...> ...>
+    for (unsigned i = 0; i < arity; ++i) {
+      struct ElementExpander {
+        const llvm::DenseMap<TypeBase *, PackType *> &expansions;
+        llvm::function_ref<Optional<Type>(TypeBase *)> outerFn;
+        unsigned index;
+
+      public:
+        Optional<Type> operator()(TypeBase *input) {
+          // FIXME: Does this need to do bounds checking?
+          if (PackType *element = expansions.lookup(input))
+            return element->getElementType(index);
+          return outerFn(input);
+        }
+      };
+
+      auto expandedElt = expand->getPatternType().transformRec(
+          ElementExpander{expansions, fn, i});
+      if (!expandedElt)
+        return Type();
+
+      elts.push_back(expandedElt);
+    }
+    return PackType::get(base->getASTContext(), elts);
+  }
+
   case TypeKind::Tuple: {
     auto tuple = cast<TupleType>(base);
     bool anyChanged = false;
     SmallVector<TupleTypeElt, 4> elements;
     unsigned Index = 0;
     for (const auto &elt : tuple->getElements()) {
-      Type eltTy = elt.getType().transformRec(fn);
-      if (!eltTy)
+      Type eltTy = elt.getType();
+      Type transformedEltTy = eltTy.transformRec(fn);
+      if (!transformedEltTy)
         return Type();
 
       // If nothing has changed, just keep going.
-      if (!anyChanged && eltTy.getPointer() == elt.getType().getPointer()) {
+      if (!anyChanged &&
+          transformedEltTy.getPointer() == elt.getType().getPointer()) {
         ++Index;
         continue;
       }
@@ -5086,9 +5240,18 @@ case TypeKind::Id:
         anyChanged = true;
       }
 
-      // Add the new tuple element, with the new type, no initializer,
-      elements.push_back(elt.getWithType(eltTy));
-      ++Index;
+      if (eltTy->isTypeSequenceParameter() &&
+          transformedEltTy->is<PackType>()) {
+        assert(anyChanged);
+        // Splat the tuple in by copying in all of the transformed elements.
+        auto tuple = dyn_cast<PackType>(transformedEltTy.getPointer());
+        elements.append(tuple->getElementTypes().begin(),
+                        tuple->getElementTypes().end());
+      } else {
+        // Add the new tuple element, with the transformed type.
+        elements.push_back(elt.getWithType(transformedEltTy));
+        ++Index;
+      }
     }
 
     if (!anyChanged)
@@ -5505,6 +5668,8 @@ ReferenceCounting TypeBase::getReferenceCounting() {
   case TypeKind::SILToken:
   case TypeKind::GenericTypeParam:
   case TypeKind::DependentMember:
+  case TypeKind::Pack:
+  case TypeKind::PackExpansion:
 #define REF_STORAGE(Name, ...) \
   case TypeKind::Name##Storage:
 #include "swift/AST/ReferenceStorage.def"
diff --git a/lib/Sema/TypeCheckConstraints.cpp b/lib/Sema/TypeCheckConstraints.cpp
