Merge pull request #62140 from hborla/variadic-generics-locators

hborla · web-flow · commit 72e296e32b44 · 2022-11-16T19:22:27.000-08:00
[ConstraintSystem] Add locator path elements for pack expansion patterns and pack shapes.
diff --git a/include/swift/Sema/ConstraintLocatorPathElts.def b/include/swift/Sema/ConstraintLocatorPathElts.def
@@ -195,6 +195,12 @@ CUSTOM_LOCATOR_PATH_ELT(PackType)
 /// An element of a pack type - the T in <T, U, V, ...>
 CUSTOM_LOCATOR_PATH_ELT(PackElement)
 
+/// The shape of a parameter pack.
+SIMPLE_LOCATOR_PATH_ELT(PackShape)
+
+/// The pattern  of a pack expansion.
+SIMPLE_LOCATOR_PATH_ELT(PackExpansionPattern)
+
 /// An unresolved member.
 SIMPLE_LOCATOR_PATH_ELT(UnresolvedMember)
 
diff --git a/lib/AST/ASTPrinter.cpp b/lib/AST/ASTPrinter.cpp
@@ -6411,8 +6411,6 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
   void visitElementArchetypeType(ElementArchetypeType *T) {
     if (Options.PrintForSIL) {
       Printer << "@element(\"" << T->getOpenedElementID() << ") ";
-      visit(T->getGenericEnvironment()->getOpenedExistentialType());
-      Printer << ") ";
 
       auto interfaceTy = T->getInterfaceType();
       visit(interfaceTy);
diff --git a/lib/AST/ParameterPack.cpp b/lib/AST/ParameterPack.cpp
@@ -39,8 +39,7 @@ struct PackTypeParameterCollector: TypeWalker {
       if (paramTy->isParameterPack())
         typeParams.insert(paramTy);
     } else if (auto *archetypeTy = t->getAs<PackArchetypeType>()) {
-      if (archetypeTy->isRoot())
-        typeParams.insert(paramTy);
+      typeParams.insert(archetypeTy->getRoot());
     }
 
     return Action::Continue;
@@ -315,14 +314,20 @@ CanPackType PackType::getReducedShape() {
 
 CanType TypeBase::getReducedShape() {
   if (auto *packArchetype = getAs<PackArchetypeType>())
-    return packArchetype->getReducedShape();
+     return packArchetype->getReducedShape();
 
   if (auto *packType = getAs<PackType>())
     return packType->getReducedShape();
 
   if (auto *expansionType = getAs<PackExpansionType>())
     return expansionType->getReducedShape();
 
+  SmallVector<Type, 2> rootParameterPacks;
+  getTypeParameterPacks(rootParameterPacks);
+
+  if (!rootParameterPacks.empty())
+    return rootParameterPacks.front()->getReducedShape();
+
   assert(!isTypeVariableOrMember());
   assert(!hasTypeParameter());
 
diff --git a/lib/Sema/CSGen.cpp b/lib/Sema/CSGen.cpp
@@ -2918,27 +2918,24 @@ namespace {
         CS.setType(binding, type);
       }
 
-      auto elementResultType = CS.getType(expr->getPatternExpr());
-      auto patternTy = CS.createTypeVariable(CS.getConstraintLocator(expr),
+      auto *patternLoc =
+          CS.getConstraintLocator(expr, ConstraintLocator::PackExpansionPattern);
+      auto patternTy = CS.createTypeVariable(patternLoc,
                                              TVO_CanBindToPack |
                                              TVO_CanBindToHole);
+      auto elementResultType = CS.getType(expr->getPatternExpr());
       CS.addConstraint(ConstraintKind::PackElementOf, elementResultType,
                        patternTy, CS.getConstraintLocator(expr));
 
-      // FIXME: Use a ShapeOf constraint here.
-      Type shapeType;
-      auto *binding = expr->getBindings().front();
-      auto type = CS.simplifyType(CS.getType(binding));
-      type.visit([&](Type type) {
-        if (shapeType)
-          return;
-
-        if (auto archetype = type->getAs<PackArchetypeType>()) {
-          shapeType = archetype->getReducedShape();
-        }
-      });
+      auto *shapeLoc =
+          CS.getConstraintLocator(expr, ConstraintLocator::PackShape);
+      auto *shapeTypeVar = CS.createTypeVariable(shapeLoc,
+                                                 TVO_CanBindToPack |
+                                                 TVO_CanBindToHole);
+      CS.addConstraint(ConstraintKind::ShapeOf, patternTy, shapeTypeVar,
+                       CS.getConstraintLocator(expr));
 
-      return PackExpansionType::get(patternTy, shapeType);
+      return PackExpansionType::get(patternTy, shapeTypeVar);
     }
 
     Type visitDynamicTypeExpr(DynamicTypeExpr *expr) {
diff --git a/lib/Sema/CSSimplify.cpp b/lib/Sema/CSSimplify.cpp
@@ -2423,14 +2423,17 @@ static PackType *replaceTypeVariablesWithFreshPacks(ConstraintSystem &cs,
       // wrapping a pack type variable. Otherwise, create a new scalar
       // type variable.
       //
-      // FIXME: Locator for diagnostics
       // FIXME: Other TVO_* flags for type variables?
+      auto elementLoc = cs.getConstraintLocator(loc,
+          LocatorPathElt::PackElement(freshTypeVars.size()));
       if (packExpansionElt != nullptr) {
-        auto *freshTypeVar = cs.createTypeVariable(loc, TVO_CanBindToPack);
+        auto *freshTypeVar =
+            cs.createTypeVariable(elementLoc, TVO_CanBindToPack);
         freshTypeVars.push_back(PackExpansionType::get(
             freshTypeVar, packExpansionElt->getCountType()));
       } else {
-        freshTypeVars.push_back(cs.createTypeVariable(loc, /*options=*/0));
+        freshTypeVars.push_back(
+            cs.createTypeVariable(elementLoc, /*options=*/0));
       }
     }
   }
@@ -2483,7 +2486,6 @@ static PackType *replaceTypeVariablesWithFreshPacks(ConstraintSystem &cs,
   // Bind each pack type variable occurring in the pattern type to its
   // binding pack that was constructed above.
   for (const auto &pair : typeVars) {
-    // FIXME: Locator for diagnostics
     cs.addConstraint(ConstraintKind::Bind,
                      pair.first, PackType::get(ctx, pair.second), locator);
   }
@@ -2498,10 +2500,9 @@ ConstraintSystem::matchPackExpansionTypes(PackExpansionType *expansion1,
                                           ConstraintKind kind, TypeMatchOptions flags,
                                           ConstraintLocatorBuilder locator) {
   // The count types of two pack expansion types must have the same shape.
-  //
-  // FIXME: Locator for diagnostics.
-  auto *loc = getConstraintLocator(locator);
-  auto *shapeTypeVar = createTypeVariable(loc, TVO_CanBindToPack);
+  auto *shapeLoc = getConstraintLocator(
+      locator.withPathElement(ConstraintLocator::PackShape));
+  auto *shapeTypeVar = createTypeVariable(shapeLoc, TVO_CanBindToPack);
   addConstraint(ConstraintKind::ShapeOf,
                 expansion1->getCountType(), shapeTypeVar, locator);
   addConstraint(ConstraintKind::ShapeOf,
@@ -2522,7 +2523,6 @@ ConstraintSystem::matchPackExpansionTypes(PackExpansionType *expansion1,
     if (auto *pack2 = pattern2->getAs<PackType>()) {
       if (auto *pack1 = replaceTypeVariablesWithFreshPacks(
              *this, pattern1, pack2, locator)) {
-        // FIXME: Locator for diagnostics.
         addConstraint(kind, pack1, pack2, locator);
         return getTypeMatchSuccess();
       }
@@ -2538,7 +2538,6 @@ ConstraintSystem::matchPackExpansionTypes(PackExpansionType *expansion1,
     if (auto *pack1 = pattern1->getAs<PackType>()) {
       if (auto *pack2 = replaceTypeVariablesWithFreshPacks(
               *this, pattern2, pack1, locator)) {
-        // FIXME: Locator for diagnostics.
         addConstraint(kind, pack1, pack2, locator);
         return getTypeMatchSuccess();
       }
@@ -6852,8 +6851,6 @@ ConstraintSystem::matchTypes(Type type1, Type type2, ConstraintKind kind,
                             kind, subflags, packLoc);
     }
     case TypeKind::PackExpansion: {
-      // FIXME: Need a new locator element
-
       auto expansion1 = cast<PackExpansionType>(desugar1);
       auto expansion2 = cast<PackExpansionType>(desugar2);
 
@@ -7562,9 +7559,10 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifySubclassOfConstraint(
     for (unsigned i = 0, e = packType->getNumElements(); i < e; ++i) {
       auto eltType = packType->getElementType(i);
       if (auto *packExpansionType = eltType->getAs<PackExpansionType>()) {
-        // FIXME: Locator element for pack expansion pattern
+        auto patternLoc =
+            locator.withPathElement(ConstraintLocator::PackExpansionPattern);
         addConstraint(ConstraintKind::SubclassOf, packExpansionType->getPatternType(),
-                      classType, locator.withPathElement(LocatorPathElt::PackElement(i)));
+                      classType, patternLoc);
       } else {
         addConstraint(ConstraintKind::SubclassOf, eltType,
                       classType, locator.withPathElement(LocatorPathElt::PackElement(i)));
@@ -7680,11 +7678,12 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyConformsToConstraint(
     for (unsigned i = 0, e = packType->getNumElements(); i < e; ++i) {
       auto eltType = packType->getElementType(i);
       if (auto *packExpansionType = eltType->getAs<PackExpansionType>()) {
-        // FIXME: Locator element for pack expansion pattern
+        auto patternLoc =
+            locator.withPathElement(ConstraintLocator::PackExpansionPattern);
         addConstraint(ConstraintKind::ConformsTo,
                       packExpansionType->getPatternType(),
                       protocol->getDeclaredInterfaceType(),
-                      locator.withPathElement(LocatorPathElt::PackElement(i)));
+                      patternLoc);
       } else {
         addConstraint(ConstraintKind::ConformsTo, eltType,
                       protocol->getDeclaredInterfaceType(),
@@ -14158,8 +14157,9 @@ void ConstraintSystem::addConstraint(Requirement req,
     auto type1 = req.getFirstType();
     auto type2 = req.getSecondType();
 
-    // FIXME: Locator for diagnostics
-    auto typeVar = createTypeVariable(getConstraintLocator(locator),
+    auto *shapeLoc = getConstraintLocator(
+        locator.withPathElement(ConstraintLocator::PackShape));
+    auto typeVar = createTypeVariable(shapeLoc,
                                       TVO_CanBindToPack);
 
     addConstraint(ConstraintKind::ShapeOf, type1, typeVar, locator);
diff --git a/lib/Sema/ConstraintLocator.cpp b/lib/Sema/ConstraintLocator.cpp
@@ -97,6 +97,8 @@ unsigned LocatorPathElt::getNewSummaryFlags() const {
   case ConstraintLocator::SyntacticElement:
   case ConstraintLocator::PackType:
   case ConstraintLocator::PackElement:
+  case ConstraintLocator::PackShape:
+  case ConstraintLocator::PackExpansionPattern:
   case ConstraintLocator::PatternBindingElement:
   case ConstraintLocator::NamedPatternDecl:
   case ConstraintLocator::AnyPatternDecl:
@@ -434,16 +436,28 @@ void LocatorPathElt::dump(raw_ostream &out) const {
     break;
   }
 
-  case ConstraintLocator::ConstraintLocator::PackType:
-    out << "pack type";
+  case ConstraintLocator::ConstraintLocator::PackType: {
+    auto packElt = elt.castTo<LocatorPathElt::PackType>();
+    out << "pack type '" << packElt.getType()->getString(PO) << "'";
     break;
+  }
 
   case ConstraintLocator::PackElement: {
     auto packElt = elt.castTo<LocatorPathElt::PackElement>();
     out << "pack element #" << llvm::utostr(packElt.getIndex());
     break;
   }
 
+  case ConstraintLocator::PackShape: {
+    out << "pack shape";
+    break;
+  }
+
+  case ConstraintLocator::PackExpansionPattern: {
+    out << "pack expansion pattern";
+    break;
+  }
+
   case ConstraintLocator::PatternBindingElement: {
     auto patternBindingElt =
         elt.castTo<LocatorPathElt::PatternBindingElement>();
diff --git a/lib/Sema/ConstraintSystem.cpp b/lib/Sema/ConstraintSystem.cpp
@@ -5327,8 +5327,17 @@ void constraints::simplifyLocator(ASTNode &anchor,
       break;
 
     case ConstraintLocator::PackElement:
+    case ConstraintLocator::PackShape:
       break;
 
+    case ConstraintLocator::PackExpansionPattern: {
+      if (auto *expansion = getAsExpr<PackExpansionExpr>(anchor))
+        anchor = expansion->getPatternExpr();
+
+      path = path.slice(1);
+      break;
+    }
+
     case ConstraintLocator::PatternBindingElement: {
       auto pattern = path[0].castTo<LocatorPathElt::PatternBindingElement>();
       auto *patternBinding = cast<PatternBindingDecl>(anchor.get<Decl *>());
