[Requirement Machine] Implement same-element requirements.

hborla · simanerush · commit 2ea45865803a · 2024-07-15T10:19:32.000-07:00
diff --git a/lib/AST/ASTPrinter.cpp b/lib/AST/ASTPrinter.cpp
@@ -2022,6 +2022,9 @@ void PrintAST::printRequirement(const Requirement &req) {
   SmallVector<Type, 2> rootParameterPacks;
   getTransformedType(req.getFirstType())
       ->getTypeParameterPacks(rootParameterPacks);
+  if (req.getKind() != RequirementKind::Layout)
+    getTransformedType(req.getSecondType())
+        ->getTypeParameterPacks(rootParameterPacks);
   bool isPackRequirement = !rootParameterPacks.empty();
 
   switch (req.getKind()) {
diff --git a/lib/AST/GenericEnvironment.cpp b/lib/AST/GenericEnvironment.cpp
@@ -181,16 +181,16 @@ struct FindElementArchetypeForOpenedPackParam {
     : findElementParam(env, openedPacks), getElementArchetype(env) {}
 
 
-  ElementArchetypeType *operator()(Type interfaceType) {
+  Type operator()(Type interfaceType) {
     assert(interfaceType->isTypeParameter());
     if (auto member = interfaceType->getAs<DependentMemberType>()) {
-      auto baseArchetype = (*this)(member->getBase());
+      auto baseArchetype = (*this)(member->getBase())
+             ->castTo<ElementArchetypeType>();
       return baseArchetype->getNestedType(member->getAssocType())
                ->castTo<ElementArchetypeType>();
     }
     assert(interfaceType->is<GenericTypeParamType>());
-    return getElementArchetype(findElementParam(interfaceType))
-             ->castTo<ElementArchetypeType>();
+    return getElementArchetype(findElementParam(interfaceType));
   }
 };
 
diff --git a/lib/AST/GenericSignature.cpp b/lib/AST/GenericSignature.cpp
@@ -782,6 +782,10 @@ int swift::compareDependentTypes(Type type1, Type type2) {
   // Fast-path check for equality.
   if (type1->isEqual(type2)) return 0;
 
+  // Packs are always ordered after scalar type parameters.
+  if (type1->isParameterPack() != type2->isParameterPack())
+    return type2->isParameterPack() ? -1 : +1;
+
   // Ordering is as follows:
   // - Generic params
   auto gp1 = type1->getAs<GenericTypeParamType>();
diff --git a/lib/AST/RequirementMachine/GenericSignatureQueries.cpp b/lib/AST/RequirementMachine/GenericSignatureQueries.cpp
@@ -281,6 +281,7 @@ RequirementMachine::getLongestValidPrefix(const MutableTerm &term) const {
     case Symbol::Kind::ConcreteType:
     case Symbol::Kind::ConcreteConformance:
     case Symbol::Kind::Shape:
+    case Symbol::Kind::PackElement:
       llvm::errs() <<"Invalid symbol in a type term: " << term << "\n";
       abort();
     }
@@ -310,6 +311,20 @@ bool RequirementMachine::isReducedType(Type type) const {
       if (!component->hasTypeParameter())
         return Action::SkipNode;
 
+      if (auto *expansion = component->getAs<PackExpansionType>()) {
+        auto pattern = expansion->getPatternType();
+        auto shape = expansion->getCountType();
+        if (!Self.isReducedType(pattern))
+          return Action::Stop;
+
+        auto reducedShape =
+            Self.getReducedShape(shape, Self.getGenericParams());
+        if (reducedShape->getCanonicalType() != CanType(shape))
+          return Action::Stop;
+
+        return Action::SkipChildren;
+      }
+
       if (!component->isTypeParameter())
         return Action::Continue;
 
@@ -804,6 +819,7 @@ void RequirementMachine::verify(const MutableTerm &term) const {
       switch (symbol.getKind()) {
       case Symbol::Kind::Protocol:
       case Symbol::Kind::GenericParam:
+      case Symbol::Kind::PackElement:
         erased.add(symbol);
         continue;
 
@@ -843,6 +859,7 @@ void RequirementMachine::verify(const MutableTerm &term) const {
     case Symbol::Kind::Superclass:
     case Symbol::Kind::ConcreteType:
     case Symbol::Kind::ConcreteConformance:
+    case Symbol::Kind::PackElement:
       llvm::errs() << "Bad interior symbol " << symbol << " in " << term << "\n";
       abort();
       break;
diff --git a/lib/AST/RequirementMachine/HomotopyReduction.cpp b/lib/AST/RequirementMachine/HomotopyReduction.cpp
@@ -691,7 +691,8 @@ RewriteSystem::getMinimizedGenericSignatureRules() const {
       continue;
     }
 
-    if (rule.getLHS()[0].getKind() != Symbol::Kind::GenericParam)
+    if (rule.getLHS()[0].getKind() != Symbol::Kind::PackElement &&
+        rule.getLHS()[0].getKind() != Symbol::Kind::GenericParam)
       continue;
 
     rules.push_back(ruleID);
diff --git a/lib/AST/RequirementMachine/InterfaceType.cpp b/lib/AST/RequirementMachine/InterfaceType.cpp
@@ -298,6 +298,9 @@ getTypeForSymbolRange(const Symbol *begin, const Symbol *end,
         // member type rooted at Self; handle the associated type below.
         break;
 
+      case Symbol::Kind::PackElement:
+        continue;
+
       case Symbol::Kind::Name:
       case Symbol::Kind::Layout:
       case Symbol::Kind::Superclass:
diff --git a/lib/AST/RequirementMachine/MinimalConformances.cpp b/lib/AST/RequirementMachine/MinimalConformances.cpp
@@ -287,6 +287,7 @@ static const ProtocolDecl *getParentConformanceForTerm(Term lhs) {
   case Symbol::Kind::ConcreteType:
   case Symbol::Kind::ConcreteConformance:
   case Symbol::Kind::Shape:
+  case Symbol::Kind::PackElement:
     break;
   }
 
@@ -553,6 +554,7 @@ void RewriteSystem::computeCandidateConformancePaths(
         //
         // where Y is the simplified form of X.W.
         } else if (rhs.isAnyConformanceRule() &&
+                   !lhs.isSameElementRule() &&
                    (unsigned)(lhs.getLHS().end() - from) < rhs.getLHS().size()) {
           if (Debug.contains(DebugFlags::MinimalConformancesDetail)) {
             llvm::dbgs() << "Case 2: same-type suffix\n";
diff --git a/lib/AST/RequirementMachine/PropertyUnification.cpp b/lib/AST/RequirementMachine/PropertyUnification.cpp
@@ -669,6 +669,7 @@ void PropertyMap::addProperty(
   case Symbol::Kind::GenericParam:
   case Symbol::Kind::AssociatedType:
   case Symbol::Kind::Shape:
+  case Symbol::Kind::PackElement:
     break;
   }
 
diff --git a/lib/AST/RequirementMachine/RequirementBuilder.cpp b/lib/AST/RequirementMachine/RequirementBuilder.cpp
@@ -229,6 +229,7 @@ void RequirementBuilder::addRequirementRules(ArrayRef<unsigned> rules) {
       case Symbol::Kind::AssociatedType:
       case Symbol::Kind::GenericParam:
       case Symbol::Kind::Shape:
+      case Symbol::Kind::PackElement:
         break;
       }
 
diff --git a/lib/AST/RequirementMachine/RequirementLowering.cpp b/lib/AST/RequirementMachine/RequirementLowering.cpp
@@ -220,14 +220,6 @@ static void desugarSameTypeRequirement(
         break;
       }
 
-      // If one side is a parameter pack, this is a same-element requirement, which
-      // is not yet supported.
-      if (firstType->isParameterPack() != secondType->isParameterPack()) {
-        errors.push_back(RequirementError::forSameElement(
-            {kind, sugaredFirstType, secondType}, loc));
-        return true;
-      }
-
       if (firstType->isTypeParameter() && secondType->isTypeParameter()) {
         result.emplace_back(kind, sugaredFirstType, secondType);
         return true;
diff --git a/lib/AST/RequirementMachine/RequirementMachine.cpp b/lib/AST/RequirementMachine/RequirementMachine.cpp
@@ -551,7 +551,7 @@ void RequirementMachine::dump(llvm::raw_ostream &out) const {
     for (auto paramTy : Params) {
       out << " " << Type(paramTy);
       if (paramTy->isParameterPack())
-        out << "…";
+        out << " " << paramTy;
     }
     out << " >";
   }
diff --git a/lib/AST/RequirementMachine/RequirementMachineRequests.cpp b/lib/AST/RequirementMachine/RequirementMachineRequests.cpp
@@ -986,13 +986,13 @@ InferredGenericSignatureRequest::evaluate(
         if (reduced->hasError() || reduced->isEqual(genericParam))
           continue;
 
-        if (reduced->isTypeParameter()) {
-          // If one side is a parameter pack and the other is not, this is a
-          // same-element requirement that cannot be expressed with only one
-          // type parameter.
-          if (genericParam->isParameterPack() != reduced->isParameterPack())
-            continue;
+        // If one side is a parameter pack and the other is not, this is a
+        // same-element requirement that cannot be expressed with only one
+        // type parameter.
+        if (genericParam->isParameterPack() != reduced->isParameterPack())
+          continue;
 
+        if (reduced->isTypeParameter()) {
           ctx.Diags.diagnose(loc, diag::requires_generic_params_made_equal,
                              genericParam, result->getSugaredType(reduced))
             .warnUntilSwiftVersion(6);
diff --git a/lib/AST/RequirementMachine/RewriteContext.h b/lib/AST/RequirementMachine/RewriteContext.h
@@ -48,6 +48,9 @@ class RewriteContext final {
   /// The singleton storage for shape symbols.
   Symbol::Storage *TheShapeSymbol;
 
+  /// The singleton storage for pack element symbols.
+  Symbol::Storage *ThePackElementSymbol;
+
   /// Folding set for uniquing terms.
   llvm::FoldingSet<Term::Storage> Terms;
 
diff --git a/lib/AST/RequirementMachine/RewriteSystem.cpp b/lib/AST/RequirementMachine/RewriteSystem.cpp
@@ -555,7 +555,8 @@ void RewriteSystem::verifyRewriteRules(ValidityPolicy policy) const {
       }
 
       if (index != 0) {
-        ASSERT_RULE(symbol.getKind() != Symbol::Kind::GenericParam);
+        ASSERT_RULE(symbol.getKind() != Symbol::Kind::GenericParam ||
+                    lhs[index - 1].getKind() == Symbol::Kind::PackElement);
       }
 
       if (!rule.isLHSSimplified() &&
diff --git a/lib/AST/RequirementMachine/Rule.cpp b/lib/AST/RequirementMachine/Rule.cpp
@@ -74,6 +74,7 @@ const ProtocolDecl *Rule::isAnyConformanceRule() const {
     case Symbol::Kind::AssociatedType:
     case Symbol::Kind::GenericParam:
     case Symbol::Kind::Shape:
+    case Symbol::Kind::PackElement:
       break;
     }
 
@@ -145,6 +146,11 @@ bool Rule::isCircularConformanceRule() const {
   return true;
 }
 
+/// Returns \c true if this rule is prefixed with the \c [element] symbol.
+bool Rule::isSameElementRule() const {
+  return LHS[0].getKind() == Symbol::Kind::PackElement;
+}
+
 /// A protocol typealias rule takes one of the following two forms,
 /// where T is a name symbol:
 ///
diff --git a/lib/AST/RequirementMachine/Rule.h b/lib/AST/RequirementMachine/Rule.h
@@ -123,6 +123,8 @@ class Rule final {
 
   bool isCircularConformanceRule() const;
 
+  bool isSameElementRule() const;
+
   /// See above for an explanation of these predicates.
   bool isPermanent() const {
     return Permanent;
diff --git a/lib/AST/RequirementMachine/RuleBuilder.cpp b/lib/AST/RequirementMachine/RuleBuilder.cpp
@@ -390,6 +390,21 @@ void RuleBuilder::addRequirement(const Requirement &req,
                             otherType, *substitutions)
                       : Context.getMutableTermForType(
                             otherType, proto));
+
+    if (subjectType->isParameterPack() != otherType->isParameterPack()) {
+      // This is a same-element requirement.
+      auto elementSymbol = Symbol::forPackElement(Context);
+      llvm::SmallVector<Symbol, 3> symbols{elementSymbol};
+
+      if (subjectType->isParameterPack()) {
+        symbols.append(subjectTerm.begin(), subjectTerm.end());
+        subjectTerm = MutableTerm(std::move(symbols));
+      } else {
+        symbols.append(constraintTerm.begin(), constraintTerm.end());
+        constraintTerm = MutableTerm(std::move(symbols));
+      }
+    }
+
     break;
   }
   }
diff --git a/lib/AST/RequirementMachine/Symbol.cpp b/lib/AST/RequirementMachine/Symbol.cpp
@@ -82,6 +82,14 @@ struct Symbol::Storage final
     Kind = Kind::Shape;
   }
 
+  /// A dummy type for overload resolution of the
+  /// 'pack element' constructor for Storage.
+  struct ForPackElement {};
+
+  explicit Storage(ForPackElement shape) {
+    Kind = Kind::PackElement;
+  }
+
   Storage(const ProtocolDecl *proto, Identifier name) {
     Kind = Symbol::Kind::AssociatedType;
     Proto = proto;
@@ -318,6 +326,16 @@ Symbol Symbol::forShape(RewriteContext &ctx) {
   return (ctx.TheShapeSymbol = symbol);
 }
 
+Symbol Symbol::forPackElement(RewriteContext &ctx) {
+  if (auto *symbol = ctx.ThePackElementSymbol)
+    return symbol;
+
+  unsigned size = Storage::totalSizeToAlloc<unsigned, Term>(0, 0);
+  void *mem = ctx.Allocator.Allocate(size, alignof(Storage));
+  auto *symbol = new (mem) Storage(Storage::ForPackElement());
+  return (ctx.ThePackElementSymbol = symbol);
+}
+
 /// Creates a layout symbol, representing a layout constraint.
 Symbol Symbol::forLayout(LayoutConstraint layout,
                          RewriteContext &ctx) {
@@ -457,6 +475,7 @@ const ProtocolDecl *Symbol::getRootProtocol() const {
     return getProtocol();
 
   case Symbol::Kind::GenericParam:
+  case Symbol::Kind::PackElement:
     return nullptr;
 
   case Symbol::Kind::Name:
@@ -481,6 +500,8 @@ const ProtocolDecl *Symbol::getRootProtocol() const {
 /// - AssociatedType
 /// - GenericParam
 /// - Name
+/// - Shape
+/// - PackElement
 /// - Layout
 /// - Superclass
 /// - ConcreteType
@@ -537,6 +558,7 @@ std::optional<int> Symbol::compare(Symbol other, RewriteContext &ctx) const {
   }
 
   case Kind::Shape:
+  case Kind::PackElement:
   case Kind::GenericParam: {
     auto *param = getGenericParam();
     auto *otherParam = other.getGenericParam();
@@ -619,6 +641,7 @@ Symbol Symbol::withConcreteSubstitutions(
   case Kind::Protocol:
   case Kind::AssociatedType:
   case Kind::Shape:
+  case Kind::PackElement:
   case Kind::Layout:
     break;
   }
@@ -736,6 +759,11 @@ void Symbol::dump(llvm::raw_ostream &out) const {
     out << "[shape]";
     return;
   }
+
+  case Kind::PackElement: {
+    out << "[element]";
+    return;
+  }
   }
 
   llvm_unreachable("Bad symbol kind");
@@ -758,6 +786,7 @@ void Symbol::Storage::Profile(llvm::FoldingSetNodeID &id) const {
     return;
 
   case Symbol::Kind::GenericParam:
+  case Symbol::Kind::PackElement:
     id.AddPointer(GenericParam);
     return;
 
diff --git a/lib/AST/RequirementMachine/Symbol.h b/lib/AST/RequirementMachine/Symbol.h
@@ -117,6 +117,10 @@ class Symbol final {
     /// generic parameter.
     Shape,
 
+    /// A pack element [element].(each T) where 'each T' is a type
+    /// parameter pack.
+    PackElement,
+
     //////
     ////// "Property-like" symbol kinds:
     //////
@@ -206,6 +210,8 @@ class Symbol final {
 
   static Symbol forShape(RewriteContext &ctx);
 
+  static Symbol forPackElement(RewriteContext &Ctx);
+
   static Symbol forLayout(LayoutConstraint layout,
                           RewriteContext &ctx);
 
diff --git a/lib/AST/RequirementMachine/Term.cpp b/lib/AST/RequirementMachine/Term.cpp
@@ -134,19 +134,32 @@ static std::optional<int> shortlexCompare(const Symbol *lhsBegin,
                                           const Symbol *rhsBegin,
                                           const Symbol *rhsEnd,
                                           RewriteContext &ctx) {
-  // First, compare the number of name symbols.
+  // First, compare the number of name and pack element symbols.
   unsigned lhsNameCount = 0;
+  unsigned lhsPackElementCount = 0;
   for (auto *iter = lhsBegin; iter != lhsEnd; ++iter) {
     if (iter->getKind() == Symbol::Kind::Name)
       ++lhsNameCount;
+
+    if (iter->getKind() == Symbol::Kind::PackElement)
+      ++lhsPackElementCount;
   }
 
   unsigned rhsNameCount = 0;
+  unsigned rhsPackElementCount = 0;
   for (auto *iter = rhsBegin; iter != rhsEnd; ++iter) {
     if (iter->getKind() == Symbol::Kind::Name)
       ++rhsNameCount;
+
+    if (iter->getKind() == Symbol::Kind::PackElement)
+      ++rhsPackElementCount;
   }
 
+  // A term with more pack element symbols orders after a term with
+  // fewer pack element symbols.
+  if (lhsPackElementCount != rhsPackElementCount)
+    return lhsPackElementCount > rhsPackElementCount ? 1 : -1;
+
   // A term with more name symbols orders after a term with fewer name symbols.
   if (lhsNameCount != rhsNameCount)
     return lhsNameCount > rhsNameCount ? 1 : -1;
diff --git a/test/Constraints/pack-expansion-expressions.swift b/test/Constraints/pack-expansion-expressions.swift
diff --git a/test/Generics/parameter-pack-requirements.swift b/test/Generics/parameter-pack-requirements.swift
diff --git a/test/Generics/parameter-pack-same-element-requirements.swift b/test/Generics/parameter-pack-same-element-requirements.swift

Original file line number	Diff line number	Diff line change
`@@ -691,7 +691,8 @@ RewriteSystem::getMinimizedGenericSignatureRules() const {`
`691`	`691`	`continue;`
`692`	`692`	`}`
`693`	`693`
`694`		`- if (rule.getLHS()[0].getKind() != Symbol::Kind::GenericParam)`
	`694`	`+ if (rule.getLHS()[0].getKind() != Symbol::Kind::PackElement &&`
	`695`	`+ rule.getLHS()[0].getKind() != Symbol::Kind::GenericParam)`
`695`	`696`	`continue;`
`696`	`697`
`697`	`698`	`rules.push_back(ruleID);`
Original file line number	Diff line number	Diff line change
`@@ -669,6 +669,7 @@ void PropertyMap::addProperty(`
`669`	`669`	`case Symbol::Kind::GenericParam:`
`670`	`670`	`case Symbol::Kind::AssociatedType:`
`671`	`671`	`case Symbol::Kind::Shape:`
	`672`	`+ case Symbol::Kind::PackElement:`
`672`	`673`	`break;`
`673`	`674`	`}`
`674`	`675`
Original file line number	Diff line number	Diff line change
`@@ -229,6 +229,7 @@ void RequirementBuilder::addRequirementRules(ArrayRef<unsigned> rules) {`
`229`	`229`	`case Symbol::Kind::AssociatedType:`
`230`	`230`	`case Symbol::Kind::GenericParam:`
`231`	`231`	`case Symbol::Kind::Shape:`
	`232`	`+ case Symbol::Kind::PackElement:`
`232`	`233`	`break;`
`233`	`234`	`}`
`234`	`235`