[Requirement Machine] Implement same-element requirements.

include/swift/Basic/Features.def

@@ -123,6 +123,7 @@ EXPERIMENTAL_FEATURE(NamedOpaqueTypes, false)
 EXPERIMENTAL_FEATURE(FlowSensitiveConcurrencyCaptures, false)
 EXPERIMENTAL_FEATURE(CodeItemMacros, false)
 EXPERIMENTAL_FEATURE(TupleConformances, false)
+EXPERIMENTAL_FEATURE(SameElementRequirements, false)
 
 SUPPRESSIBLE_LANGUAGE_FEATURE(ExtensionMacroAttr, 0, "@attached(extension)", true)
 

lib/AST/ASTPrinter.cpp

@@ -1854,6 +1854,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()) {
@@ -3256,6 +3259,10 @@ static bool usesFeatureTupleConformances(Decl *decl) {
   return false;
 }
 
+static bool usesFeatureSameElementRequirements(Decl *decl) {
+  return false;
+}
+
 static bool usesFeatureSymbolLinkageMarkers(Decl *decl) {
   auto &attrs = decl->getAttrs();
   return std::any_of(attrs.begin(), attrs.end(), [](auto *attr) {

lib/AST/GenericEnvironment.cpp

@@ -180,16 +180,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));
   }
 };
 

lib/AST/GenericSignature.cpp

@@ -105,8 +105,16 @@ void GenericSignatureImpl::forEachParam(
 
   for (auto req : getRequirements()) {
     GenericTypeParamType *gp;
+    bool isCanonical = false;
     switch (req.getKind()) {
     case RequirementKind::SameType: {
+      if (req.getSecondType()->isParameterPack() != 
+          req.getFirstType()->isParameterPack()) {
+        // This is a same-element requirement, which does not make
+        // type parameters non-canonical.
+        isCanonical = true;
+      }
+
       if (auto secondGP = req.getSecondType()->getAs<GenericTypeParamType>()) {
         // If two generic parameters are same-typed, then the right-hand one
         // is non-canonical.
@@ -136,7 +144,7 @@ void GenericSignatureImpl::forEachParam(
     }
 
     unsigned index = GenericParamKey(gp).findIndexIn(genericParams);
-    genericParamsAreCanonical[index] = false;
+    genericParamsAreCanonical[index] = isCanonical;
   }
 
   // Call the callback with each parameter and the result of the above analysis.
@@ -938,7 +946,12 @@ void GenericSignature::verify(ArrayRef<Requirement> reqts) const {
           llvm::errs() << "\n";
           abort();
         }
-        if (compareDependentTypes(firstType, secondType) >= 0) {
+        // FIXME: Same-element rewrite rules are of the form [element].T => U,
+        // which turns into a same-type requirement repeat U == each T. The
+        // pack element must always be on the left side of the rule, so type
+        // parameters can appear out of order in the final same-type requirement.
+        if (firstType->isParameterPack() == secondType->isParameterPack() &&
+            compareDependentTypes(firstType, secondType) >= 0) {
           llvm::errs() << "Out-of-order type parameters: ";
           reqt.dump(llvm::errs());
           llvm::errs() << "\n";

lib/AST/RequirementMachine/GenericSignatureQueries.cpp

@@ -278,6 +278,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();
     }
@@ -788,6 +789,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;
 
@@ -827,6 +829,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;

lib/AST/RequirementMachine/HomotopyReduction.cpp

@@ -743,7 +743,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);

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:

lib/AST/RequirementMachine/MinimalConformances.cpp

@@ -286,6 +286,7 @@ static const ProtocolDecl *getParentConformanceForTerm(Term lhs) {
   case Symbol::Kind::ConcreteType:
   case Symbol::Kind::ConcreteConformance:
   case Symbol::Kind::Shape:
+  case Symbol::Kind::PackElement:
     break;
   }
 
@@ -552,6 +553,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";

lib/AST/RequirementMachine/PropertyUnification.cpp

@@ -646,6 +646,7 @@ void PropertyMap::addProperty(
   case Symbol::Kind::GenericParam:
   case Symbol::Kind::AssociatedType:
   case Symbol::Kind::Shape:
+  case Symbol::Kind::PackElement:
     break;
   }
 

lib/AST/RequirementMachine/RequirementBuilder.cpp

@@ -248,6 +248,7 @@ void RequirementBuilder::addRequirementRules(ArrayRef<unsigned> rules) {
       case Symbol::Kind::AssociatedType:
       case Symbol::Kind::GenericParam:
       case Symbol::Kind::Shape:
+      case Symbol::Kind::PackElement:
         break;
       }
 

lib/AST/RequirementMachine/RequirementLowering.cpp

@@ -218,13 +218,16 @@ static void desugarSameTypeRequirement(Requirement req, SourceLoc loc,
         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));
-        recordedErrors = true;
-        return true;
+      auto &ctx = firstType->getASTContext();
+      if (!ctx.LangOpts.hasFeature(Feature::SameElementRequirements)) {
+        // 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));
+          recordedErrors = true;
+          return true;
+        }
       }
 
       if (firstType->isTypeParameter() && secondType->isTypeParameter()) {

lib/AST/RequirementMachine/RequirementMachine.cpp

@@ -560,7 +560,7 @@ void RequirementMachine::dump(llvm::raw_ostream &out) const {
     for (auto paramTy : Params) {
       out << " " << Type(paramTy);
       if (paramTy->isParameterPack())
-        out << "…";
+        out << " " << paramTy;
     }
     out << " >";
   }

lib/AST/RequirementMachine/RequirementMachineRequests.cpp

@@ -963,13 +963,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);

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;
 

lib/AST/RequirementMachine/RewriteSystem.cpp

@@ -569,7 +569,9 @@ void RewriteSystem::verifyRewriteRules(ValidityPolicy policy) const {
       }
 
       if (index != 0) {
-        ASSERT_RULE(symbol.getKind() != Symbol::Kind::GenericParam);
+        ASSERT_RULE(symbol.getKind() != Symbol::Kind::GenericParam ||
+                    (index == 1 &&
+                     lhs[index - 1].getKind() == Symbol::Kind::PackElement));
       }
 
       if (!rule.isLHSSimplified() &&
@@ -624,7 +626,9 @@ void RewriteSystem::verifyRewriteRules(ValidityPolicy policy) const {
       }
 
       if (index != 0) {
-        ASSERT_RULE(symbol.getKind() != Symbol::Kind::GenericParam);
+        ASSERT_RULE(symbol.getKind() != Symbol::Kind::GenericParam ||
+                    (index == 1 &&
+                     lhs[index - 1].getKind() == Symbol::Kind::PackElement));
       }
 
       if (!rule.isRHSSimplified() &&

lib/AST/RequirementMachine/Rule.cpp

@@ -73,6 +73,7 @@ const ProtocolDecl *Rule::isAnyConformanceRule() const {
     case Symbol::Kind::AssociatedType:
     case Symbol::Kind::GenericParam:
     case Symbol::Kind::Shape:
+    case Symbol::Kind::PackElement:
       break;
     }
 
@@ -146,6 +147,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:
 ///

lib/AST/RequirementMachine/Rule.h

@@ -137,6 +137,8 @@ class Rule final {
 
   bool isCircularConformanceRule() const;
 
+  bool isSameElementRule() const;
+
   /// See above for an explanation of these predicates.
   bool isPermanent() const {
     return Permanent;

lib/AST/RequirementMachine/RuleBuilder.cpp

@@ -367,6 +367,7 @@ void RuleBuilder::addRequirement(const Requirement &req,
 
   case RequirementKind::SameType: {
     auto otherType = CanType(req.getSecondType());
+    auto elementSymbol = Symbol::forPackElement(Context);
 
     if (!otherType->isTypeParameter()) {
       // A concrete same-type requirement T == C<X, Y> becomes a
@@ -380,6 +381,16 @@ void RuleBuilder::addRequirement(const Requirement &req,
                    : Context.getSubstitutionSchemaFromType(
                         otherType, proto, result));
 
+      // If 'T' is a parameter pack, this is a same-element
+      // requirement that becomes the following rewrite rule:
+      //
+      //   [element].T.[concrete: C<X, Y>] => [element].T
+      if (subjectType->isParameterPack()) {
+        llvm::SmallVector<Symbol, 3> subjectSymbols{elementSymbol};
+        subjectSymbols.append(subjectTerm.begin(), subjectTerm.end());
+        subjectTerm = MutableTerm(std::move(subjectSymbols));
+      }
+
       constraintTerm = subjectTerm;
       constraintTerm.add(Symbol::forConcreteType(otherType, result, Context));
       break;
@@ -390,6 +401,20 @@ void RuleBuilder::addRequirement(const Requirement &req,
                             otherType, *substitutions)
                       : Context.getMutableTermForType(
                             otherType, proto));
+
+    if (subjectType->isParameterPack() != otherType->isParameterPack()) {
+      // This is a same-element requirement.
+      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;
   }
   }
@@ -565,4 +590,4 @@ void RuleBuilder::collectPackShapeRules(ArrayRef<GenericTypeParamType *> generic
       }
     }
   }
-}
+}