Sema: Add ConstraintKind::ShapeOf

Author: slavapestov

include/swift/Sema/Constraint.h

@@ -223,7 +223,9 @@ enum class ConstraintKind : char {
   ///
   /// Binds the RHS type to a tuple of the params of a function typed LHS. Note
   /// this discards function parameter flags.
-  BindTupleOfFunctionParams
+  BindTupleOfFunctionParams,
+  /// The first type is a type pack, and the second type is its reduced shape.
+  ShapeOf,
 };
 
 /// Classification of the different kinds of constraints.
@@ -705,6 +707,7 @@ class Constraint final : public llvm::ilist_node<Constraint>,
     case ConstraintKind::KeyPathApplication:
     case ConstraintKind::Defaultable:
     case ConstraintKind::BindTupleOfFunctionParams:
+    case ConstraintKind::ShapeOf:
       return ConstraintClassification::TypeProperty;
 
     case ConstraintKind::Disjunction:

include/swift/Sema/ConstraintSystem.h

@@ -5613,6 +5613,12 @@ class ConstraintSystem {
       ASTNode element, ContextualTypeInfo context, bool isDiscarded,
       TypeMatchOptions flags, ConstraintLocatorBuilder locator);
 
+  /// Simplify a shape constraint by binding the reduced shape of the
+  /// left hand side to the right hand side.
+  SolutionKind simplifyShapeOfConstraint(
+      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

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

lib/Sema/CSSimplify.cpp

@@ -2425,6 +2425,7 @@ ConstraintSystem::matchTupleTypes(TupleType *tuple1, TupleType *tuple2,
   case ConstraintKind::SyntacticElement:
   case ConstraintKind::BindTupleOfFunctionParams:
   case ConstraintKind::PackElementOf:
+  case ConstraintKind::ShapeOf:
     llvm_unreachable("Bad constraint kind in matchTupleTypes()");
   }
 
@@ -2597,6 +2598,7 @@ static bool matchFunctionRepresentations(FunctionType::ExtInfo einfo1,
   case ConstraintKind::SyntacticElement:
   case ConstraintKind::BindTupleOfFunctionParams:
   case ConstraintKind::PackElementOf:
+  case ConstraintKind::ShapeOf:
     return true;
   }
 
@@ -3013,6 +3015,7 @@ ConstraintSystem::matchFunctionTypes(FunctionType *func1, FunctionType *func2,
   case ConstraintKind::SyntacticElement:
   case ConstraintKind::BindTupleOfFunctionParams:
   case ConstraintKind::PackElementOf:
+  case ConstraintKind::ShapeOf:
     llvm_unreachable("Not a relational constraint");
   }
 
@@ -6366,6 +6369,7 @@ ConstraintSystem::matchTypes(Type type1, Type type2, ConstraintKind kind,
     case ConstraintKind::SyntacticElement:
     case ConstraintKind::BindTupleOfFunctionParams:
     case ConstraintKind::PackElementOf:
+    case ConstraintKind::ShapeOf:
       llvm_unreachable("Not a relational constraint");
     }
   }
@@ -7442,17 +7446,14 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifySubclassOfConstraint(
     return SolutionKind::Solved;
   }
 
-  auto formUnsolved = [&](bool activate = false) {
+  auto formUnsolved = [&]() {
     // If we're supposed to generate constraints, do so.
     if (flags.contains(TMF_GenerateConstraints)) {
-      auto *conformance = Constraint::create(
+      auto *subclassOf = Constraint::create(
           *this, ConstraintKind::SubclassOf, type, classType,
           getConstraintLocator(locator));
 
-      addUnsolvedConstraint(conformance);
-      if (activate)
-        activateConstraint(conformance);
-
+      addUnsolvedConstraint(subclassOf);
       return SolutionKind::Solved;
     }
 
@@ -12510,6 +12511,74 @@ ConstraintSystem::simplifyDynamicCallableApplicableFnConstraint(
   return SolutionKind::Solved;
 }
 
+static Type getReducedShape(Type type) {
+  // Pack archetypes know their reduced shape
+  if (auto *packArchetype = type->getAs<PackArchetypeType>())
+    return packArchetype->getShape();
+
+  // Reduced shape of pack is computed recursively
+  if (auto *packType = type->getAs<PackType>()) {
+    auto &ctx = type->getASTContext();
+    SmallVector<Type, 2> elts;
+
+    for (auto elt : packType->getElementTypes()) {
+      // T... => shape(T)...
+      if (auto *packExpansionType = elt->getAs<PackExpansionType>()) {
+        auto shapeType = getReducedShape(packExpansionType->getCountType());
+        assert(shapeType && "Should not end up here if pack type's shape "
+                            "is still potentially unknown");
+        elts.push_back(PackExpansionType::get(shapeType, shapeType));
+      }
+
+      // Use () as a placeholder for scalar shape
+      elts.push_back(ctx.TheEmptyTupleType);
+    }
+
+    return PackType::get(ctx, elts);
+  }
+
+  // Getting the shape of any other type is an error.
+  return Type();
+}
+
+ConstraintSystem::SolutionKind ConstraintSystem::simplifyShapeOfConstraint(
+    Type type1, Type type2, TypeMatchOptions flags,
+    ConstraintLocatorBuilder locator) {
+  // Recursively replace all type variables with fixed bindings if
+  // possible.
+  type1 = simplifyType(type1, flags);
+
+  auto formUnsolved = [&]() {
+    // If we're supposed to generate constraints, do so.
+    if (flags.contains(TMF_GenerateConstraints)) {
+      auto *shapeOf = Constraint::create(
+          *this, ConstraintKind::ShapeOf, type1, type2,
+          getConstraintLocator(locator));
+
+      addUnsolvedConstraint(shapeOf);
+      return SolutionKind::Solved;
+    }
+
+    return SolutionKind::Unsolved;
+  };
+
+  // We can't compute a reduced shape if the input type still
+  // contains type variables that might bind to pack archetypes.
+  SmallPtrSet<TypeVariableType *, 2> typeVars;
+  type1->getTypeVariables(typeVars);
+  for (auto *typeVar : typeVars) {
+    if (typeVar->getImpl().canBindToPack())
+      return formUnsolved();
+  }
+
+  if (Type shape = getReducedShape(type1)) {
+    addConstraint(ConstraintKind::Bind, shape, type2, locator);
+    return SolutionKind::Solved;
+  }
+
+  return SolutionKind::Error;
+}
+
 static llvm::PointerIntPair<Type, 3, unsigned>
 getBaseTypeForPointer(TypeBase *type) {
   unsigned unwrapCount = 0;
@@ -13847,6 +13916,9 @@ ConstraintSystem::addConstraintImpl(ConstraintKind kind, Type first,
   case ConstraintKind::PackElementOf:
     return simplifyPackElementOfConstraint(first, second, subflags, locator);
 
+  case ConstraintKind::ShapeOf:
+    return simplifyShapeOfConstraint(first, second, subflags, locator);
+
   case ConstraintKind::ValueMember:
   case ConstraintKind::UnresolvedValueMember:
   case ConstraintKind::ValueWitness:
@@ -14010,8 +14082,18 @@ void ConstraintSystem::addConstraint(Requirement req,
   bool conformsToAnyObject = false;
   Optional<ConstraintKind> kind;
   switch (req.getKind()) {
-  case RequirementKind::SameShape:
-    llvm_unreachable("Same-shape requirement not supported here");
+  case RequirementKind::SameShape: {
+    auto type1 = req.getFirstType();
+    auto type2 = req.getSecondType();
+
+    // FIXME: Locator for diagnostics
+    auto typeVar = createTypeVariable(getConstraintLocator(locator),
+                                      TVO_CanBindToPack);
+
+    addConstraint(ConstraintKind::ShapeOf, type1, typeVar, locator);
+    addConstraint(ConstraintKind::ShapeOf, type2, typeVar, locator);
+    return;
+  }
 
   case RequirementKind::Conformance:
     kind = ConstraintKind::ConformsTo;
@@ -14421,6 +14503,11 @@ ConstraintSystem::simplifyConstraint(const Constraint &constraint) {
     return simplifyPackElementOfConstraint(
         constraint.getFirstType(), constraint.getSecondType(), /*flags*/None,
         constraint.getLocator());
+
+  case ConstraintKind::ShapeOf:
+    return simplifyShapeOfConstraint(
+        constraint.getFirstType(), constraint.getSecondType(), /*flags*/ None,
+        constraint.getLocator());
   }
 
   llvm_unreachable("Unhandled ConstraintKind in switch.");

lib/Sema/Constraint.cpp

@@ -80,6 +80,7 @@ Constraint::Constraint(ConstraintKind Kind, Type First, Type Second,
   case ConstraintKind::PropertyWrapper:
   case ConstraintKind::BindTupleOfFunctionParams:
   case ConstraintKind::PackElementOf:
+  case ConstraintKind::ShapeOf:
     assert(!First.isNull());
     assert(!Second.isNull());
     break;
@@ -167,6 +168,7 @@ Constraint::Constraint(ConstraintKind Kind, Type First, Type Second, Type Third,
   case ConstraintKind::SyntacticElement:
   case ConstraintKind::BindTupleOfFunctionParams:
   case ConstraintKind::PackElementOf:
+  case ConstraintKind::ShapeOf:
     llvm_unreachable("Wrong constructor");
 
   case ConstraintKind::KeyPath:
@@ -313,6 +315,7 @@ Constraint *Constraint::clone(ConstraintSystem &cs) const {
   case ConstraintKind::PropertyWrapper:
   case ConstraintKind::BindTupleOfFunctionParams:
   case ConstraintKind::PackElementOf:
+  case ConstraintKind::ShapeOf:
     return create(cs, getKind(), getFirstType(), getSecondType(), getLocator());
 
   case ConstraintKind::ApplicableFunction:
@@ -553,6 +556,10 @@ void Constraint::print(llvm::raw_ostream &Out, SourceManager *sm, unsigned inden
     Out << " element of pack expansion pattern ";
     break;
 
+  case ConstraintKind::ShapeOf:
+    Out << " shape of ";
+    break;
+
   case ConstraintKind::Disjunction:
     llvm_unreachable("disjunction handled above");
   case ConstraintKind::Conjunction:
@@ -718,6 +725,7 @@ gatherReferencedTypeVars(Constraint *constraint,
   case ConstraintKind::PropertyWrapper:
   case ConstraintKind::BindTupleOfFunctionParams:
   case ConstraintKind::PackElementOf:
+  case ConstraintKind::ShapeOf:
     constraint->getFirstType()->getTypeVariables(typeVars);
     constraint->getSecondType()->getTypeVariables(typeVars);
     break;

lib/Sema/ConstraintSystem.cpp

@@ -1773,8 +1773,6 @@ void ConstraintSystem::openGenericRequirement(
 
   auto kind = req.getKind();
   switch (kind) {
-  case RequirementKind::SameShape:
-    llvm_unreachable("Same-shape requirement not supported here");
   case RequirementKind::Conformance: {
     auto protoDecl = req.getProtocolDecl();
     // Determine whether this is the protocol 'Self' constraint we should
@@ -1788,6 +1786,7 @@ void ConstraintSystem::openGenericRequirement(
   }
   case RequirementKind::Superclass:
   case RequirementKind::SameType:
+  case RequirementKind::SameShape:
     openedReq = Requirement(kind, openedFirst, substFn(req.getSecondType()));
     break;
   case RequirementKind::Layout:

test/Constraints/variadic_generic_constraints.swift

@@ -51,4 +51,16 @@ func takesParallelSequences<T..., U...>(t: T..., u: U...) where T: Sequence, U:
 takesParallelSequences()  // ok
 takesParallelSequences(t: Array<Int>(), u: Set<Int>())  // ok
 takesParallelSequences(t: Array<String>(), Set<Int>(), u: Set<String>(), Array<Int>())  // ok
-takesParallelSequences(t: Array<String>(), Set<Int>(), u: Array<Int>(), Set<String>())  // expected-error {{global function 'takesParallelSequences(t:u:)' requires the types 'String' and 'Int' be equivalent}}
+takesParallelSequences(t: Array<String>(), Set<Int>(), u: Array<Int>(), Set<String>())  // expected-error {{global function 'takesParallelSequences(t:u:)' requires the types 'String' and 'Int' be equivalent}}
+
+// Same-shape requirements
+
+func zip<T..., U...>(t: T..., u: U...) -> ((T, U)...) {}
+
+let _ = zip()  // ok
+let _ = zip(t: 1, u: "hi")  // ok
+let _ = zip(t: 1, 2, u: "hi", "hello")  // ok
+let _ = zip(t: 1, 2, 3, u: "hi", "hello", "greetings")  // ok
+
+// FIXME: Bad diagnostic
+let _ = zip(t: 1, u: "hi", "hello", "greetings")  // expected-error {{type of expression is ambiguous without more context}}