Merge pull request swiftlang#67502 from slavapestov/rqm-shape-req-fixes

RequirementMachine: Fix some edge cases with shape requirements

Author: slavapestov

include/swift/AST/Type.h

@@ -214,7 +214,7 @@ enum class ForeignRepresentableKind : uint8_t {
 /// therefore, the result type is in covariant position relative to the function
 /// type.
 struct TypePosition final {
-  enum : uint8_t { Covariant, Contravariant, Invariant };
+  enum : uint8_t { Covariant, Contravariant, Invariant, Shape };
 
 private:
   decltype(Covariant) kind;
@@ -224,6 +224,7 @@ struct TypePosition final {
 
   TypePosition flipped() const {
     switch (kind) {
+    case Shape:
     case Invariant:
       return *this;
     case Covariant:

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/InterfaceType.cpp

@@ -403,7 +403,7 @@ Type PropertyMap::getTypeForTerm(const MutableTerm &term,
 /// Concrete type terms are written in terms of generic parameter types that
 /// have a depth of 0, and an index into an array of substitution terms.
 ///
-/// See RewriteSystemBuilder::getConcreteSubstitutionSchema().
+/// See RewriteSystemBuilder::getSubstitutionSchemaFromType().
 unsigned RewriteContext::getGenericParamIndex(Type type) {
   auto *paramTy = type->castTo<GenericTypeParamType>();
   assert(paramTy->getDepth() == 0);
@@ -429,6 +429,7 @@ RewriteContext::getRelativeTermForType(CanType typeWitness,
   // Get the substitution S corresponding to τ_0_n.
   unsigned index = getGenericParamIndex(typeWitness->getRootGenericParam());
   result = MutableTerm(substitutions[index]);
+  assert(result.back().getKind() != Symbol::Kind::Shape);
 
   // If the substitution is a term consisting of a single protocol symbol
   // [P], save P for later.
@@ -471,7 +472,7 @@ RewriteContext::getRelativeTermForType(CanType typeWitness,
 }
 
 /// Reverses the transformation performed by
-/// RewriteSystemBuilder::getConcreteSubstitutionSchema().
+/// RewriteSystemBuilder::getSubstitutionSchemaFromType().
 Type PropertyMap::getTypeFromSubstitutionSchema(
     Type schema, ArrayRef<Term> substitutions,
     ArrayRef<GenericTypeParamType *> genericParams,
@@ -481,11 +482,38 @@ Type PropertyMap::getTypeFromSubstitutionSchema(
   if (!schema->hasTypeParameter())
     return schema;
 
-  return schema.transformRec([&](Type t) -> llvm::Optional<Type> {
+  return schema.transformWithPosition(
+      TypePosition::Invariant,
+      [&](Type t, TypePosition pos) -> llvm::Optional<Type> {
     if (t->is<GenericTypeParamType>()) {
       auto index = RewriteContext::getGenericParamIndex(t);
       auto substitution = substitutions[index];
 
+      bool isShapePosition = (pos == TypePosition::Shape);
+      bool isShapeTerm = (substitution.back() == Symbol::forShape(Context));
+      if (isShapePosition != isShapeTerm) {
+        llvm::errs() << "Shape vs. type mixup\n\n";
+        schema.dump(llvm::errs());
+        llvm::errs() << "Substitutions:\n";
+        for (auto otherSubst : substitutions) {
+          llvm::errs() << "- ";
+          otherSubst.dump(llvm::errs());
+          llvm::errs() << "\n";
+        }
+        llvm::errs() << "\n";
+        dump(llvm::errs());
+
+        abort();
+      }
+
+      // Undo the thing where the count type of a PackExpansionType
+      // becomes a shape term.
+      if (isShapeTerm) {
+        MutableTerm mutTerm(substitution.begin(),
+                            substitution.end() - 1);
+        substitution = Term::get(mutTerm, Context);
+      }
+
       // Prepend the prefix of the lookup key to the substitution.
       if (prefix.empty()) {
         // Skip creation of a new MutableTerm in the case where the
@@ -535,17 +563,31 @@ RewriteContext::getRelativeSubstitutionSchemaFromType(
   if (!concreteType->hasTypeParameter())
     return concreteType;
 
-  return CanType(concreteType.transformRec([&](Type t) -> llvm::Optional<Type> {
+  return CanType(concreteType.transformWithPosition(
+      TypePosition::Invariant,
+      [&](Type t, TypePosition pos) -> llvm::Optional<Type> {
+
     if (!t->isTypeParameter())
       return llvm::None;
 
     auto term = getRelativeTermForType(CanType(t), substitutions);
 
-    unsigned newIndex = result.size();
+    // PackExpansionType(pattern=T, count=U) becomes
+    // PackExpansionType(pattern=τ_0_0, count=τ_0_1) with
+    //
+    // τ_0_0 := T
+    // τ_0_1 := U.[shape]
+    if (pos == TypePosition::Shape) {
+      assert(false);
+      term.add(Symbol::forShape(*this));
+    }
+
+    unsigned index = result.size();
+
     result.push_back(Term::get(term, *this));
 
     return CanGenericTypeParamType::get(/*isParameterPack=*/ false,
-                                        /*depth=*/ 0, newIndex,
+                                        /*depth=*/ 0, index,
                                         Context);
   }));
 }
@@ -566,12 +608,26 @@ RewriteContext::getSubstitutionSchemaFromType(CanType concreteType,
   if (!concreteType->hasTypeParameter())
     return concreteType;
 
-  return CanType(concreteType.transformRec([&](Type t) -> llvm::Optional<Type> {
+  return CanType(concreteType.transformWithPosition(
+      TypePosition::Invariant,
+      [&](Type t, TypePosition pos)
+          -> llvm::Optional<Type> {
+
     if (!t->isTypeParameter())
       return llvm::None;
 
+    // PackExpansionType(pattern=T, count=U) becomes
+    // PackExpansionType(pattern=τ_0_0, count=τ_0_1) with
+    //
+    // τ_0_0 := T
+    // τ_0_1 := U.[shape]
+    MutableTerm term = getMutableTermForType(CanType(t), proto);
+    if (pos == TypePosition::Shape)
+      term.add(Symbol::forShape(*this));
+
     unsigned index = result.size();
-    result.push_back(getTermForType(CanType(t), proto));
+
+    result.push_back(Term::get(term, *this));
 
     return CanGenericTypeParamType::get(/*isParameterPack=*/ false,
                                         /*depth=*/0, index,

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;
       }
 

lib/AST/RequirementMachine/Symbol.h

@@ -57,7 +57,7 @@ class Term;
 /// This transformation allows DependentMemberTypes to be manipulated as
 /// terms, with the actual concrete type structure remaining opaque to
 /// the requirement machine. This transformation is implemented in
-/// RewriteContext::getConcreteSubstitutionSchema().
+/// RewriteContext::getSubstitutionSchemaFromType().
 ///
 /// For example, the superclass requirement
 /// "T : MyClass<U.X, (Int) -> V.A.B>" is denoted with a symbol

lib/AST/RequirementMachine/TypeDifference.cpp

@@ -159,6 +159,9 @@ namespace {
       bool rhsAbstract = rhsType->isTypeParameter();
 
       if (lhsAbstract && rhsAbstract) {
+        // FIXME: same-element requirements
+        assert(lhsType->isParameterPack() == rhsType->isParameterPack());
+
         unsigned lhsIndex = RewriteContext::getGenericParamIndex(lhsType);
         unsigned rhsIndex = RewriteContext::getGenericParamIndex(rhsType);
 

lib/AST/Type.cpp

@@ -4677,7 +4677,7 @@ case TypeKind::Id:
       return Type();
 
     Type transformedCount =
-        expand->getCountType().transformWithPosition(pos, fn);
+        expand->getCountType().transformWithPosition(TypePosition::Shape, fn);
     if (!transformedCount)
       return Type();
 

lib/Sema/ConstraintSystem.cpp

@@ -2199,6 +2199,7 @@ static Type typeEraseExistentialSelfReferences(Type refTy, Type baseTy,
 
           case TypePosition::Contravariant:
           case TypePosition::Invariant:
+          case TypePosition::Shape:
             return Type(t);
           }
 

test/Generics/pack-shape-requirements.swift

@@ -83,3 +83,46 @@ 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> {}
+
+/// More complex example involving concrete type matching in
+/// property map construction
+
+protocol PP {
+  associatedtype A
+}
+
+struct G<each T> {}
+
+// CHECK-LABEL: sameTypeMatch1
+// CHECK-NEXT: <T, each U, each V where T : PP, repeat each U : PP, repeat each V : PP, T.[PP]A == G<repeat (each U).[PP]A>, repeat (each U).[PP]A == (each V).[PP]A>
+func sameTypeMatch1<T: PP, each U: PP, each V: PP>(t: T, u: repeat each U, v: repeat each V)
+  where T.A == G<repeat (each U).A>, T.A == G<repeat (each V).A>,
+        (repeat (each U, each V)) : Any {}
+
+// CHECK-LABEL: sameTypeMatch2
+// CHECK-NEXT: <T, each U, each V where T : PP, repeat each U : PP, (repeat (each U, each V)) : Any, repeat each V : PP, T.[PP]A == (/* shape: each U */ repeat ())>
+func sameTypeMatch2<T: PP, each U: PP, each V: PP>(t: T, u: repeat each U, v: repeat each V)
+  where T.A == Shape<repeat each U>, T.A == Shape<repeat each V> {}