RequirementMachine: Correct desugaring of same-type requirements with pack expansion on both sides

Author: slavapestov

include/swift/AST/TypeMatcher.h

@@ -51,6 +51,13 @@ namespace swift {
 /// or false to indicate that matching should exit early.
 template<typename ImplClass>
 class TypeMatcher {
+public:
+  enum class Position : uint8_t {
+    Type,
+    Shape
+  };
+
+private:
   class MatchVisitor : public CanTypeVisitor<MatchVisitor, bool, Type, Type> {
     TypeMatcher &Matcher;
 
@@ -192,11 +199,24 @@ class TypeMatcher {
 
     bool visitPackExpansionType(CanPackExpansionType firstPE, Type secondType,
                                 Type sugaredFirstType) {
-      if (auto secondInOut = secondType->getAs<PackExpansionType>()) {
-        return this->visit(firstPE.getPatternType(),
-                           secondInOut->getPatternType(),
-                           sugaredFirstType->castTo<PackExpansionType>()
-                             ->getPatternType());
+      if (auto secondExpansion = secondType->getAs<PackExpansionType>()) {
+        if (!this->visit(firstPE.getPatternType(),
+                         secondExpansion->getPatternType(),
+                         sugaredFirstType->castTo<PackExpansionType>()
+                           ->getPatternType())) {
+          return false;
+        }
+
+        Matcher.asDerived().pushPosition(Position::Shape);
+        if (!this->visit(firstPE.getCountType(),
+                         secondExpansion->getCountType(),
+                         sugaredFirstType->castTo<PackExpansionType>()
+                           ->getCountType())) {
+          return false;
+        }
+        Matcher.asDerived().popPosition(Position::Shape);
+
+        return true;
       }
 
       return mismatch(firstPE.getPointer(), secondType, sugaredFirstType);
@@ -524,6 +544,9 @@ class TypeMatcher {
 
   bool alwaysMismatchTypeParameters() const { return false; }
 
+  void pushPosition(Position pos) {}
+  void popPosition(Position pos) {}
+
   ImplClass &asDerived() { return static_cast<ImplClass &>(*this); }
 
   const ImplClass &asDerived() const {

lib/AST/RequirementMachine/PropertyUnification.cpp

@@ -24,7 +24,6 @@
 
 #include "swift/AST/Decl.h"
 #include "swift/AST/LayoutConstraint.h"
-#include "swift/AST/TypeMatcher.h"
 #include "swift/AST/Types.h"
 #include <algorithm>
 #include <vector>

lib/AST/RequirementMachine/RequirementLowering.cpp

@@ -180,6 +180,7 @@ static void desugarSameTypeRequirement(Requirement req, SourceLoc loc,
     SourceLoc loc;
     SmallVectorImpl<Requirement> &result;
     SmallVectorImpl<RequirementError> &errors;
+    SmallVector<Position, 2> stack;
 
   public:
     bool recordedErrors = false;
@@ -191,32 +192,53 @@ static void desugarSameTypeRequirement(Requirement req, SourceLoc loc,
 
     bool alwaysMismatchTypeParameters() const { return true; }
 
+    void pushPosition(Position pos) {
+      stack.push_back(pos);
+    }
+
+    void popPosition(Position pos) {
+      assert(stack.back() == pos);
+      stack.pop_back();
+    }
+
+    Position getPosition() const {
+      if (stack.empty()) return Position::Type;
+      return stack.back();
+    }
+
     bool mismatch(TypeBase *firstType, TypeBase *secondType,
                   Type sugaredFirstType) {
+      RequirementKind kind;
+      switch (getPosition()) {
+      case Position::Type:
+        kind = RequirementKind::SameType;
+        break;
+      case Position::Shape:
+        kind = RequirementKind::SameShape;
+        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(
-            {RequirementKind::SameType, sugaredFirstType, secondType}, loc));
+            {kind, sugaredFirstType, secondType}, loc));
         recordedErrors = true;
         return true;
       }
 
       if (firstType->isTypeParameter() && secondType->isTypeParameter()) {
-        result.emplace_back(RequirementKind::SameType,
-                            sugaredFirstType, secondType);
+        result.emplace_back(kind, sugaredFirstType, secondType);
         return true;
       }
 
       if (firstType->isTypeParameter()) {
-        result.emplace_back(RequirementKind::SameType,
-                            sugaredFirstType, secondType);
+        result.emplace_back(kind, sugaredFirstType, secondType);
         return true;
       }
 
       if (secondType->isTypeParameter()) {
-        result.emplace_back(RequirementKind::SameType,
-                            secondType, sugaredFirstType);
+        result.emplace_back(kind, secondType, sugaredFirstType);
         return true;
       }
 

test/Generics/pack-shape-requirements.swift

@@ -83,3 +83,27 @@ protocol Q: P where A: Q {}
 // CHECK-LABEL: sameType2
 // CHECK-NEXT: Generic signature: <each T, each U where repeat each T : Q, repeat each U : Q, repeat (each T).[P]A.[P]A == (each U).[P]A.[P]A>
 func sameType2<each T, each U>(_: repeat (each T, each U)) where repeat each T: Q, repeat each U: Q, repeat (each T).A.A == (each U).A.A {}
+
+
+//////
+///
+/// A same-type requirement between two pack expansion types
+/// should desugar to a same-shape requirement between their
+/// count types and a same-type requirement between their
+/// element types.
+///
+//////
+
+typealias First<T, U> = T
+typealias Shape<each T> = (repeat First<(), each T>)
+
+// CHECK-LABEL: sameTypeDesugar1
+// CHECK-NEXT: Generic signature: <each T, each U where (repeat (each T, each U)) : Any>
+func sameTypeDesugar1<each T, each U>(t: repeat each T, u: repeat each U)
+  where Shape<repeat each T> == Shape<repeat each U> {}
+
+// CHECK-LABEL: sameTypeDesugar2
+// CHECK-NEXT: Generic signature: <each T, each U where repeat each T : P, (repeat (each T, each U)) : Any, repeat each U : P>
+func sameTypeDesugar2<each T: P, each U: P>(t: repeat each T, u: repeat each U)
+  where Shape<repeat (each T).A> == Shape<repeat (each U).A> {}
+