[RequirementMachine] Generalize connected components to process same

shape requirements.

Author: hborla

lib/AST/RequirementMachine/RequirementBuilder.cpp

@@ -43,6 +43,7 @@ struct ConnectedComponent {
   Type ConcreteType;
 
   void buildRequirements(Type subjectType,
+                         RequirementKind kind,
                          std::vector<Requirement> &reqs,
                          std::vector<ProtocolTypeAlias> &aliases);
 };
@@ -68,6 +69,7 @@ struct ConnectedComponent {
 ///
 ///   A == X, B == X, C == X, D == X
 void ConnectedComponent::buildRequirements(Type subjectType,
+                                           RequirementKind kind,
                                            std::vector<Requirement> &reqs,
                                            std::vector<ProtocolTypeAlias> &aliases) {
   std::sort(Members.begin(), Members.end(),
@@ -81,12 +83,14 @@ void ConnectedComponent::buildRequirements(Type subjectType,
     }
 
     for (auto constraintType : Members) {
-      reqs.emplace_back(RequirementKind::SameType,
-                        subjectType, constraintType);
+      reqs.emplace_back(kind, subjectType, constraintType);
       subjectType = constraintType;
     }
 
   } else {
+    // Shape requirements cannot be concrete.
+    assert(kind == RequirementKind::SameType);
+
     // If there are multiple protocol typealiases in the connected component,
     // lower them all to a series of identical concrete-type aliases.
     for (auto name : Aliases) {
@@ -135,7 +139,7 @@ class RequirementBuilder {
 
   // Temporary state populated by addRequirementRules() and
   // addTypeAliasRules().
-  llvm::SmallDenseMap<TypeBase *, ConnectedComponent> Components;
+  llvm::SmallDenseMap<Term, ConnectedComponent> Components;
 
 public:
   // Results.
@@ -229,7 +233,7 @@ void RequirementBuilder::addRequirementRules(ArrayRef<unsigned> rules) {
         if (ReconstituteSugar)
           concreteType = concreteType->reconstituteSugar(/*recursive=*/true);
 
-        auto &component = Components[subjectType.getPointer()];
+        auto &component = Components[rule.getRHS()];
         assert(!component.ConcreteType);
         component.ConcreteType = concreteType;
         return;
@@ -250,28 +254,18 @@ void RequirementBuilder::addRequirementRules(ArrayRef<unsigned> rules) {
       llvm_unreachable("Invalid symbol kind");
     }
 
-    if (rule.getLHS().back().getKind() == Symbol::Kind::Shape) {
-      assert(rule.getRHS().back().getKind() == Symbol::Kind::Shape);
+    assert(rule.getLHS().back().getKind() != Symbol::Kind::Protocol);
 
-      // Strip off the shape symbols from either side of the rule.
-      MutableTerm lhsTerm(rule.getLHS().begin(),
-                          rule.getLHS().end() - 1);
-      MutableTerm rhsTerm(rule.getRHS().begin(),
-                          rule.getRHS().end() - 1);
-
-      // Add a SameCount requirement between the two parameter packs.
-      auto constraintType = Map.getTypeForTerm(lhsTerm, GenericParams);
-      auto subjectType = Map.getTypeForTerm(rhsTerm, GenericParams);
-      Reqs.emplace_back(RequirementKind::SameCount,
-                        subjectType, constraintType);
-      return;
+    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);
     }
 
-    assert(rule.getLHS().back().getKind() != Symbol::Kind::Protocol);
-    auto constraintType = Map.getTypeForTerm(rule.getLHS(), GenericParams);
-    auto subjectType = Map.getTypeForTerm(rule.getRHS(), GenericParams);
-
-    Components[subjectType.getPointer()].Members.push_back(constraintType);
+    auto constraintType = Map.getTypeForTerm(constraintTerm, GenericParams);
+    Components[rule.getRHS()].Members.push_back(constraintType);
   };
 
   if (Debug) {
@@ -319,11 +313,11 @@ void RequirementBuilder::addTypeAliasRules(ArrayRef<unsigned> rules) {
       if (ReconstituteSugar)
         concreteType = concreteType->reconstituteSugar(/*recursive=*/true);
 
-      auto &component = Components[subjectType.getPointer()];
+      auto &component = Components[rule.getRHS()];
       assert(!component.ConcreteType);
-      Components[subjectType.getPointer()].ConcreteType = concreteType;
+      Components[rule.getRHS()].ConcreteType = concreteType;
     } else {
-      Components[subjectType.getPointer()].Aliases.push_back(name);
+      Components[rule.getRHS()].Aliases.push_back(name);
     }
   }
 }
@@ -332,7 +326,18 @@ void RequirementBuilder::processConnectedComponents() {
   // Now, convert each connected component into a series of same-type
   // requirements.
   for (auto &pair : Components) {
-    pair.second.buildRequirements(pair.first, Reqs, Aliases);
+    MutableTerm subjectTerm(pair.first);
+    RequirementKind kind;
+    if (subjectTerm.back().getKind() == Symbol::Kind::Shape) {
+      kind = RequirementKind::SameCount;
+      // Strip off the shape symbol from the subject term.
+      subjectTerm = MutableTerm{subjectTerm.begin(), subjectTerm.end() - 1};
+    } else {
+      kind = RequirementKind::SameType;
+    }
+
+    auto subjectType = Map.getTypeForTerm(subjectTerm, GenericParams);
+    pair.second.buildRequirements(subjectType, kind, Reqs, Aliases);
   }
 }
 

test/Generics/pack-shape-requirements.swift

@@ -13,3 +13,33 @@ func inferSameShape<@_typeSequence T, @_typeSequence U>(ts t: T..., us u: U...)
 // CHECK-NEXT: Generic signature: <@_typeSequence T, @_typeSequence U where T : P, T.count == U.count, U : P>
 func desugarSameShape<@_typeSequence T, @_typeSequence U>(ts t: T..., us u: U...) where T: P, U: P, ((T.A, U.A)...): Any {
 }
+
+// CHECK-LABEL: multipleSameShape1(ts:us:vs:)
+// CHECK-NEXT: Generic signature: <@_typeSequence T, @_typeSequence U, @_typeSequence V where T.count == U.count, U.count == V.count>
+func multipleSameShape1<@_typeSequence T, @_typeSequence U, @_typeSequence V>(ts t: T..., us u: U..., vs v: V...) where ((T, U, V)...): Any {
+}
+
+// CHECK-LABEL: multipleSameShape2(ts:us:vs:)
+// CHECK-NEXT: Generic signature: <@_typeSequence T, @_typeSequence U, @_typeSequence V where T.count == U.count, U.count == V.count>
+func multipleSameShape2<@_typeSequence T, @_typeSequence U, @_typeSequence V>(ts t: T..., us u: U..., vs v: V...) where ((V, T, U)...): Any {
+}
+
+// CHECK-LABEL: multipleSameShape3(ts:us:vs:)
+// CHECK-NEXT: Generic signature: <@_typeSequence T, @_typeSequence U, @_typeSequence V where T.count == U.count, U.count == V.count>
+func multipleSameShape3<@_typeSequence T, @_typeSequence U, @_typeSequence V>(ts t: T..., us u: U..., vs v: V...) where ((U, V, T)...): Any {
+}
+
+// CHECK-LABEL: multipleSameShape4(ts:us:vs:)
+// CHECK-NEXT: Generic signature: <@_typeSequence T, @_typeSequence U, @_typeSequence V where T.count == U.count, U.count == V.count>
+func multipleSameShape4<@_typeSequence T, @_typeSequence U, @_typeSequence V>(ts t: T..., us u: U..., vs v: V...) where ((U, T, V)...): Any {
+}
+
+// CHECK-LABEL: multipleSameShape5(ts:us:vs:)
+// CHECK-NEXT: Generic signature: <@_typeSequence T, @_typeSequence U, @_typeSequence V where T.count == U.count, U.count == V.count>
+func multipleSameShape5<@_typeSequence T, @_typeSequence U, @_typeSequence V>(ts t: T..., us u: U..., vs v: V...) where ((T, V, U)...): Any {
+}
+
+// CHECK-LABEL: multipleSameShape6(ts:us:vs:)
+// CHECK-NEXT: Generic signature: <@_typeSequence T, @_typeSequence U, @_typeSequence V where T.count == U.count, U.count == V.count>
+func multipleSameShape6<@_typeSequence T, @_typeSequence U, @_typeSequence V>(ts t: T..., us u: U..., vs v: V...) where ((V, U, T)...): Any {
+}