Merge pull request #70227 from simanerush/same-element-requirements

simanerush · web-flow · commit 6400a294a88a · 2024-08-06T14:24:33.000-07:00
[Requirement Machine] Implement same-element requirements.
diff --git a/include/swift/Basic/Features.def b/include/swift/Basic/Features.def
@@ -231,6 +231,7 @@ EXPERIMENTAL_FEATURE(CodeItemMacros, false)
 EXPERIMENTAL_FEATURE(PreambleMacros, false)
 EXPERIMENTAL_FEATURE(TupleConformances, false)
 EXPERIMENTAL_FEATURE(FullTypedThrows, false)
+EXPERIMENTAL_FEATURE(SameElementRequirements, false)
 
 // Whether to enable @_used and @_section attributes
 EXPERIMENTAL_FEATURE(SymbolLinkageMarkers, true)
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/FeatureSet.cpp b/lib/AST/FeatureSet.cpp
@@ -134,6 +134,7 @@ UNINTERESTING_FEATURE(StaticExclusiveOnly)
 UNINTERESTING_FEATURE(ExtractConstantsFromMembers)
 UNINTERESTING_FEATURE(FixedArrays)
 UNINTERESTING_FEATURE(GroupActorErrors)
+UNINTERESTING_FEATURE(SameElementRequirements)
 
 static bool usesFeatureSendingArgsAndResults(Decl *decl) {
   auto isFunctionTypeWithSending = [](Type type) {
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
@@ -936,6 +936,7 @@ void GenericSignature::verify(ArrayRef<Requirement> reqts) const {
           llvm::errs() << "\n";
           dumpAndAbort();
         }
+
         if (compareDependentTypes(firstType, secondType) >= 0) {
           llvm::errs() << "Out-of-order type parameters: ";
           reqt.dump(llvm::errs());
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();
     }
@@ -804,6 +805,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 +845,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,24 +229,49 @@ void RequirementBuilder::addRequirementRules(ArrayRef<unsigned> rules) {
       case Symbol::Kind::AssociatedType:
       case Symbol::Kind::GenericParam:
       case Symbol::Kind::Shape:
+      case Symbol::Kind::PackElement:
         break;
       }
 
       llvm_unreachable("Invalid symbol kind");
     }
 
+    MutableTerm constraintTerm = MutableTerm(rule.getLHS());
+    MutableTerm subjectTerm = MutableTerm(rule.getRHS());
+
+    RewriteContext &ctx = this->System.getRewriteContext();
+
+    // Drop the [element] symbol from lhs to determine if we need to swap the
+    // sides.
+    if (constraintTerm[0].getKind() == Symbol::Kind::PackElement) {
+      constraintTerm =
+          MutableTerm(constraintTerm.begin() + 1, constraintTerm.end());
+
+      // Make sure that the shorter term is ordered first.
+      if (constraintTerm.compare(subjectTerm, ctx) == -1) {
+        MutableTerm tempTerm = subjectTerm;
+        subjectTerm = constraintTerm;
+        constraintTerm = tempTerm;
+      }
+    }
+
     ASSERT(rule.getLHS().back().getKind() != Symbol::Kind::Protocol);
 
-    MutableTerm constraintTerm(rule.getLHS());
     if (constraintTerm.back().getKind() == Symbol::Kind::Shape) {
       ASSERT(rule.getRHS().back().getKind() == Symbol::Kind::Shape);
       // Strip off the shape symbol from the constraint term.
       constraintTerm = MutableTerm(constraintTerm.begin(),
                                    constraintTerm.end() - 1);
     }
 
+    if (constraintTerm.front().getKind() == Symbol::Kind::PackElement) {
+      // Strip off the element symbol from the constraint term.
+      constraintTerm = MutableTerm(constraintTerm.begin() + 1,
+                                   constraintTerm.end());
+    }
+
     auto constraintType = Map.getTypeForTerm(constraintTerm, GenericParams);
-    Components[rule.getRHS()].Members.push_back(constraintType);
+    Components[Term::get(subjectTerm, ctx)].Members.push_back(constraintType);
   };
 
   if (Debug) {
@@ -313,6 +338,10 @@ void RequirementBuilder::processConnectedComponents() {
       subjectTerm = MutableTerm(subjectTerm.begin(), subjectTerm.end() - 1);
     } else {
       kind = RequirementKind::SameType;
+      if (subjectTerm.front().getKind() == Symbol::Kind::PackElement) {
+        // Strip off the element symbol from the subject term.
+        subjectTerm = MutableTerm(subjectTerm.begin() + 1, subjectTerm.end());
+      }
     }
 
     auto subjectType = Map.getTypeForTerm(subjectTerm, GenericParams);
diff --git a/lib/AST/RequirementMachine/RequirementLowering.cpp b/lib/AST/RequirementMachine/RequirementLowering.cpp
@@ -220,12 +220,15 @@ 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;
+      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));
+          return true;
+        }
       }
 
       if (firstType->isTypeParameter() && secondType->isTypeParameter()) {
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.cpp b/lib/AST/RequirementMachine/RewriteContext.cpp
@@ -200,6 +200,7 @@ static DebugOptions parseDebugFlags(StringRef debugFlags) {
 
 RewriteContext::RewriteContext(ASTContext &ctx)
     : TheShapeSymbol(nullptr),
+      ThePackElementSymbol(nullptr),
       Context(ctx),
       Stats(ctx.Stats),
       SymbolHistogram(Symbol::NumKinds),
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,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() &&
@@ -610,7 +612,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() &&
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
@@ -380,6 +380,14 @@ 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()) {
+        subjectTerm.prepend(Symbol::forPackElement(Context));
+      }
+
       constraintTerm = subjectTerm;
       constraintTerm.add(Symbol::forConcreteType(otherType, result, Context));
       break;
@@ -390,6 +398,16 @@ void RuleBuilder::addRequirement(const Requirement &req,
                             otherType, *substitutions)
                       : Context.getMutableTermForType(
                             otherType, proto));
+
+    if (subjectType->isParameterPack() != otherType->isParameterPack()) {
+      // This is a same-element requirement.
+      if (subjectType->isParameterPack()) {
+        subjectTerm.prepend(Symbol::forPackElement(Context));
+      } else {
+        constraintTerm.prepend(Symbol::forPackElement(Context));
+      }
+    }
+
     break;
   }
   }
diff --git a/lib/AST/RequirementMachine/Symbol.cpp b/lib/AST/RequirementMachine/Symbol.cpp
diff --git a/lib/AST/RequirementMachine/Symbol.h b/lib/AST/RequirementMachine/Symbol.h
diff --git a/lib/AST/RequirementMachine/Term.cpp b/lib/AST/RequirementMachine/Term.cpp
diff --git a/lib/AST/RequirementMachine/Term.h b/lib/AST/RequirementMachine/Term.h
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
`@@ -936,6 +936,7 @@ void GenericSignature::verify(ArrayRef<Requirement> reqts) const {`
`936`	`936`	`llvm::errs() << "\n";`
`937`	`937`	`dumpAndAbort();`
`938`	`938`	`}`
	`939`	`+`
`939`	`940`	`if (compareDependentTypes(firstType, secondType) >= 0) {`
`940`	`941`	`llvm::errs() << "Out-of-order type parameters: ";`
`941`	`942`	`reqt.dump(llvm::errs());`
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`