[ConstraintSystem] Add same-shape constraint

The constraint takes two pack types and makes sure that their
reduced shapes are equal. This helps with diagnostics because
constraint has access to the original pack expansion pattern
types.

Author: xedin

include/swift/Sema/Constraint.h

@@ -233,6 +233,8 @@ enum class ConstraintKind : char {
   /// an overload. The second type is a PackType containing the explicit
   /// generic arguments.
   ExplicitGenericArguments,
+  /// Both (first and second) pack types should have the same reduced shape.
+  SameShape,
 };
 
 /// Classification of the different kinds of constraints.
@@ -701,6 +703,7 @@ class Constraint final : public llvm::ilist_node<Constraint>,
     case ConstraintKind::DefaultClosureType:
     case ConstraintKind::UnresolvedMemberChainBase:
     case ConstraintKind::PackElementOf:
+    case ConstraintKind::SameShape:
       return ConstraintClassification::Relational;
 
     case ConstraintKind::ValueMember:

include/swift/Sema/ConstraintSystem.h

@@ -4871,6 +4871,12 @@ class ConstraintSystem {
       Type type1, Type type2, TypeMatchOptions flags,
       ConstraintLocatorBuilder locator);
 
+  /// Simplify a same-shape constraint by comparing the reduced shape of the
+  /// left hand side to the right hand side.
+  SolutionKind simplifySameShapeConstraint(Type type1, Type type2,
+                                           TypeMatchOptions flags,
+                                           ConstraintLocatorBuilder locator);
+
 public: // FIXME: Public for use by static functions.
   /// Simplify a conversion constraint with a fix applied to it.
   SolutionKind simplifyFixConstraint(ConstraintFix *fix, Type type1, Type type2,

lib/Sema/CSBindings.cpp

@@ -1480,6 +1480,7 @@ void PotentialBindings::infer(Constraint *constraint) {
   case ConstraintKind::ShapeOf:
   case ConstraintKind::ExplicitGenericArguments:
   case ConstraintKind::PackElementOf:
+  case ConstraintKind::SameShape:
     // Constraints from which we can't do anything.
     break;
 

lib/Sema/CSSimplify.cpp

@@ -2276,6 +2276,7 @@ ConstraintSystem::matchTupleTypes(TupleType *tuple1, TupleType *tuple2,
   case ConstraintKind::PackElementOf:
   case ConstraintKind::ShapeOf:
   case ConstraintKind::ExplicitGenericArguments:
+  case ConstraintKind::SameShape:
     llvm_unreachable("Bad constraint kind in matchTupleTypes()");
   }
 
@@ -2489,13 +2490,9 @@ ConstraintSystem::matchPackExpansionTypes(PackExpansionType *expansion1,
                                           ConstraintKind kind, TypeMatchOptions flags,
                                           ConstraintLocatorBuilder locator) {
   // The count types of two pack expansion types must have the same shape.
-  auto *shapeLoc = getConstraintLocator(
-      locator.withPathElement(ConstraintLocator::PackShape));
-  auto *shapeTypeVar = createTypeVariable(shapeLoc, TVO_CanBindToPack);
-  addConstraint(ConstraintKind::ShapeOf,
-                expansion1->getCountType(), shapeTypeVar, shapeLoc);
-  addConstraint(ConstraintKind::ShapeOf,
-                expansion2->getCountType(), shapeTypeVar, shapeLoc);
+  addConstraint(ConstraintKind::SameShape, expansion1->getCountType(),
+                expansion2->getCountType(),
+                locator.withPathElement(ConstraintLocator::PackShape));
 
   auto pattern1 = expansion1->getPatternType();
   auto pattern2 = expansion2->getPatternType();
@@ -2655,6 +2652,7 @@ static bool matchFunctionRepresentations(FunctionType::ExtInfo einfo1,
   case ConstraintKind::PackElementOf:
   case ConstraintKind::ShapeOf:
   case ConstraintKind::ExplicitGenericArguments:
+  case ConstraintKind::SameShape:
     return true;
   }
 
@@ -3162,6 +3160,7 @@ ConstraintSystem::matchFunctionTypes(FunctionType *func1, FunctionType *func2,
   case ConstraintKind::PackElementOf:
   case ConstraintKind::ShapeOf:
   case ConstraintKind::ExplicitGenericArguments:
+  case ConstraintKind::SameShape:
     llvm_unreachable("Not a relational constraint");
   }
 
@@ -6815,6 +6814,7 @@ ConstraintSystem::matchTypes(Type type1, Type type2, ConstraintKind kind,
     case ConstraintKind::PackElementOf:
     case ConstraintKind::ShapeOf:
     case ConstraintKind::ExplicitGenericArguments:
+    case ConstraintKind::SameShape:
       llvm_unreachable("Not a relational constraint");
     }
   }
@@ -13190,6 +13190,18 @@ ConstraintSystem::simplifyDynamicCallableApplicableFnConstraint(
   return SolutionKind::Solved;
 }
 
+static bool hasUnresolvedPackVars(Type type) {
+  // We can't compute a reduced shape if the input type still
+  // contains type variables that might bind to pack archetypes
+  // or pack expansions.
+  SmallPtrSet<TypeVariableType *, 2> typeVars;
+  type->getTypeVariables(typeVars);
+  return llvm::any_of(typeVars, [](const TypeVariableType *typeVar) {
+    return typeVar->getImpl().canBindToPack() ||
+           typeVar->getImpl().isPackExpansion();
+  });
+}
+
 ConstraintSystem::SolutionKind ConstraintSystem::simplifyShapeOfConstraint(
     Type type1, Type type2, TypeMatchOptions flags,
     ConstraintLocatorBuilder locator) {
@@ -13231,6 +13243,51 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyShapeOfConstraint(
   return SolutionKind::Solved;
 }
 
+ConstraintSystem::SolutionKind ConstraintSystem::simplifySameShapeConstraint(
+    Type type1, Type type2, TypeMatchOptions flags,
+    ConstraintLocatorBuilder locator) {
+  type1 = simplifyType(type1);
+  type2 = simplifyType(type2);
+
+  auto formUnsolved = [&]() {
+    // If we're supposed to generate constraints, do so.
+    if (flags.contains(TMF_GenerateConstraints)) {
+      auto *sameShape =
+          Constraint::create(*this, ConstraintKind::SameShape, type1, type2,
+                             getConstraintLocator(locator));
+
+      addUnsolvedConstraint(sameShape);
+      return SolutionKind::Solved;
+    }
+
+    return SolutionKind::Unsolved;
+  };
+
+  if (hasUnresolvedPackVars(type1) || hasUnresolvedPackVars(type2))
+    return formUnsolved();
+
+  auto shape1 = type1->getReducedShape();
+  auto shape2 = type2->getReducedShape();
+
+  if (shape1->isEqual(shape2))
+    return SolutionKind::Solved;
+
+  if (shouldAttemptFixes()) {
+    if (type1->hasPlaceholder() || type2->hasPlaceholder())
+      return SolutionKind::Solved;
+
+    unsigned impact = 1;
+    if (locator.endsWith<LocatorPathElt::AnyRequirement>())
+      impact = assessRequirementFailureImpact(*this, shape1, locator);
+
+    auto *fix = SkipSameShapeRequirement::create(*this, type1, type2,
+                                                 getConstraintLocator(locator));
+    return recordFix(fix, impact) ? SolutionKind::Error : SolutionKind::Solved;
+  }
+
+  return SolutionKind::Error;
+}
+
 ConstraintSystem::SolutionKind
 ConstraintSystem::simplifyExplicitGenericArgumentsConstraint(
     Type type1, Type type2, TypeMatchOptions flags,
@@ -14718,6 +14775,9 @@ ConstraintSystem::addConstraintImpl(ConstraintKind kind, Type first,
   case ConstraintKind::ShapeOf:
     return simplifyShapeOfConstraint(first, second, subflags, locator);
 
+  case ConstraintKind::SameShape:
+    return simplifySameShapeConstraint(first, second, subflags, locator);
+
   case ConstraintKind::ExplicitGenericArguments:
     return simplifyExplicitGenericArgumentsConstraint(
         first, second, subflags, locator);
@@ -14889,13 +14949,7 @@ void ConstraintSystem::addConstraint(Requirement req,
     auto type1 = req.getFirstType();
     auto type2 = req.getSecondType();
 
-    auto *shapeLoc = getConstraintLocator(
-        locator.withPathElement(ConstraintLocator::PackShape));
-    auto typeVar = createTypeVariable(shapeLoc,
-                                      TVO_CanBindToPack);
-
-    addConstraint(ConstraintKind::ShapeOf, type1, typeVar, locator);
-    addConstraint(ConstraintKind::ShapeOf, type2, typeVar, locator);
+    addConstraint(ConstraintKind::SameShape, type1, type2, locator);
     return;
   }
 
@@ -15319,6 +15373,11 @@ ConstraintSystem::simplifyConstraint(const Constraint &constraint) {
         constraint.getFirstType(), constraint.getSecondType(), /*flags*/ None,
         constraint.getLocator());
 
+  case ConstraintKind::SameShape:
+    return simplifySameShapeConstraint(constraint.getFirstType(),
+                                       constraint.getSecondType(),
+                                       /*flags*/ None, constraint.getLocator());
+
   case ConstraintKind::ExplicitGenericArguments:
     return simplifyExplicitGenericArgumentsConstraint(
         constraint.getFirstType(), constraint.getSecondType(),

lib/Sema/Constraint.cpp

@@ -82,6 +82,7 @@ Constraint::Constraint(ConstraintKind Kind, Type First, Type Second,
   case ConstraintKind::PackElementOf:
   case ConstraintKind::ShapeOf:
   case ConstraintKind::ExplicitGenericArguments:
+  case ConstraintKind::SameShape:
     assert(!First.isNull());
     assert(!Second.isNull());
     break;
@@ -171,6 +172,7 @@ Constraint::Constraint(ConstraintKind Kind, Type First, Type Second, Type Third,
   case ConstraintKind::PackElementOf:
   case ConstraintKind::ShapeOf:
   case ConstraintKind::ExplicitGenericArguments:
+  case ConstraintKind::SameShape:
     llvm_unreachable("Wrong constructor");
 
   case ConstraintKind::KeyPath:
@@ -319,6 +321,7 @@ Constraint *Constraint::clone(ConstraintSystem &cs) const {
   case ConstraintKind::PackElementOf:
   case ConstraintKind::ShapeOf:
   case ConstraintKind::ExplicitGenericArguments:
+  case ConstraintKind::SameShape:
     return create(cs, getKind(), getFirstType(), getSecondType(), getLocator());
 
   case ConstraintKind::ApplicableFunction:
@@ -568,6 +571,10 @@ void Constraint::print(llvm::raw_ostream &Out, SourceManager *sm,
     Out << " shape of ";
     break;
 
+  case ConstraintKind::SameShape:
+    Out << " same-shape ";
+    break;
+
   case ConstraintKind::ExplicitGenericArguments:
     Out << " explicit generic argument binding ";
     break;
@@ -740,6 +747,7 @@ gatherReferencedTypeVars(Constraint *constraint,
   case ConstraintKind::PackElementOf:
   case ConstraintKind::ShapeOf:
   case ConstraintKind::ExplicitGenericArguments:
+  case ConstraintKind::SameShape:
     constraint->getFirstType()->getTypeVariables(typeVars);
     constraint->getSecondType()->getTypeVariables(typeVars);
     break;

test/Constraints/pack-expansion-expressions.swift

@@ -133,7 +133,7 @@ func tupleExpansion<each T, each U>(
   _ = zip(repeat each tuple1.element, with: repeat each tuple1.element)
 
   _ = zip(repeat each tuple1.element, with: repeat each tuple2.element)
-  // expected-error@-1 {{global function 'zip(_:with:)' requires the type packs 'each U' and 'each T' have the same shape}}
+  // expected-error@-1 {{global function 'zip(_:with:)' requires the type packs 'each T' and 'each U' have the same shape}}
 }
 
 protocol Generatable {

test/Constraints/variadic_generic_constraints.swift

@@ -73,5 +73,5 @@ func goodCallToZip<each T, each U>(t: repeat each T, u: repeat each U) where (re
 
 func badCallToZip<each T, each U>(t: repeat each T, u: repeat each U) {
   _ = zip(t: repeat each t, u: repeat each u)
-  // expected-error@-1 {{global function 'zip(t:u:)' requires the type packs 'each U' and 'each T' have the same shape}}
+  // expected-error@-1 {{global function 'zip(t:u:)' requires the type packs 'each T' and 'each U' have the same shape}}
 }

test/Constraints/variadic_generic_functions.swift

@@ -32,9 +32,15 @@ func call() {
 
   let x: String = multipleParameters(xs: "", ys: "")
   let (one, two) = multipleParameters(xs: "", 5.0, ys: "", 5.0)
-  multipleParameters(xs: "", 5.0, ys: 5.0, "") // expected-error {{type of expression is ambiguous without more context}}
+  multipleParameters(xs: "", 5.0, ys: 5.0, "") // expected-error {{conflicting arguments to generic parameter 'each T' ('Pack{Double, String}' vs. 'Pack{String, String}' vs. 'Pack{String, Double}' vs. 'Pack{Double, Double}')}}
 
   func multipleSequences<each T, each U>(xs: repeat each T, ys: repeat each U) -> (repeat each T) {
+    return (repeat each ys)
+    // expected-error@-1 {{pack expansion requires that 'each U' and 'each T' have the same shape}}
+    // expected-error@-2 {{cannot convert return expression of type '(repeat each U)' to return type '(repeat each T)'}}
+  }
+
+  func multipleSequencesWithSameShape<each T, each U>(xs: repeat each T, ys: repeat each U) -> (repeat each T) where (repeat (each T, each U)): Any {
     return (repeat each ys)
     // expected-error@-1 {{cannot convert return expression of type '(repeat each U)' to return type '(repeat each T)'}}
   }