[RequirementMachine] Add same-length requirement inference for pack

expansion types.

Author: hborla

include/swift/AST/DiagnosticsSema.def

@@ -2545,6 +2545,10 @@ ERROR(requires_not_suitable_archetype,none,
       "generic parameter or associated type",
       (Type))
 
+ERROR(invalid_shape_requirement,none,
+      "invalid same-shape requirement between %0 and %1",
+      (Type, Type))
+
 ERROR(requires_generic_params_made_equal,none,
       "same-type requirement makes generic parameters %0 and %1 equivalent",
       (Type, Type))

lib/AST/RequirementMachine/ConcreteContraction.cpp

@@ -554,7 +554,8 @@ bool ConcreteContraction::performConcreteContraction(
     auto kind = req.req.getKind();
     switch (kind) {
     case RequirementKind::SameCount:
-      llvm_unreachable("Same-count requirement not supported here");
+      assert(req.req.getSecondType()->isTypeParameter());
+      continue;
 
     case RequirementKind::SameType: {
       auto constraintType = req.req.getSecondType();

lib/AST/RequirementMachine/Diagnostics.cpp

@@ -105,6 +105,20 @@ bool swift::rewriting::diagnoseRequirementErrors(
       break;
     }
 
+    case RequirementError::Kind::InvalidShapeRequirement: {
+      if (error.requirement.hasError())
+        break;
+
+      auto lhs = error.requirement.getFirstType();
+      auto rhs = error.requirement.getSecondType();
+
+      // FIXME: Add tailored messages for specific issues.
+      ctx.Diags.diagnose(loc, diag::invalid_shape_requirement,
+                         lhs, rhs);
+      diagnosedError = true;
+      break;
+    }
+
     case RequirementError::Kind::ConflictingRequirement: {
       auto requirement = error.requirement;
       auto conflict = error.conflictingRequirement;

lib/AST/RequirementMachine/Diagnostics.h

@@ -34,6 +34,8 @@ struct RequirementError {
     /// A type requirement on a trivially invalid subject type,
     /// e.g. Bool: Collection.
     InvalidRequirementSubject,
+    /// An invalid shape requirement, e.g. length(T...) == length(Int)
+    InvalidShapeRequirement,
     /// A pair of conflicting requirements, T == Int, T == String
     ConflictingRequirement,
     /// A recursive requirement, e.g. T == G<T.A>.
@@ -73,6 +75,11 @@ struct RequirementError {
     return {Kind::InvalidRequirementSubject, req, loc};
   }
 
+  static RequirementError forInvalidShapeRequirement(Requirement req,
+                                                     SourceLoc loc) {
+    return {Kind::InvalidShapeRequirement, req, loc};
+  }
+
   static RequirementError forConflictingRequirement(Requirement req,
                                                     SourceLoc loc) {
     return {Kind::ConflictingRequirement, req, loc};

lib/AST/RequirementMachine/RequirementLowering.cpp

@@ -350,6 +350,23 @@ static void desugarConformanceRequirement(Type subjectType, Type constraintType,
   }
 }
 
+/// Desugar same-shape requirements by equating the shapes of the
+/// root pack types, and diagnose shape requirements on non-pack
+/// types.
+static void desugarSameShapeRequirement(Type lhs, Type rhs, SourceLoc loc,
+                                        SmallVectorImpl<Requirement> &result,
+                                        SmallVectorImpl<RequirementError> &errors) {
+  // For now, only allow shape requirements directly between pack types.
+  if (!lhs->isTypeSequenceParameter() || !rhs->isTypeSequenceParameter()) {
+    errors.push_back(RequirementError::forInvalidShapeRequirement(
+        {RequirementKind::SameCount, lhs, rhs}, loc));
+  }
+
+  result.emplace_back(RequirementKind::SameCount,
+                      lhs->getRootGenericParam(),
+                      rhs->getRootGenericParam());
+}
+
 /// Convert a requirement where the subject type might not be a type parameter,
 /// or the constraint type in the conformance requirement might be a protocol
 /// composition, into zero or more "proper" requirements which can then be
@@ -362,7 +379,9 @@ swift::rewriting::desugarRequirement(Requirement req, SourceLoc loc,
 
   switch (req.getKind()) {
   case RequirementKind::SameCount:
-    llvm_unreachable("Same-count requirement not supported here");
+    desugarSameShapeRequirement(firstType, req.getSecondType(),
+                                loc, result, errors);
+    break;
 
   case RequirementKind::Conformance:
     desugarConformanceRequirement(firstType, req.getSecondType(),
@@ -474,6 +493,23 @@ struct InferRequirementsWalker : public TypeWalker {
       return Action::Continue;
     }
 
+    // Infer same-length requirements between pack references that
+    // are expanded in parallel.
+    if (auto packExpansion = ty->getAs<PackExpansionType>()) {
+      // Get all pack parameters referenced from the pattern.
+      SmallVector<Type, 2> packReferences;
+      packExpansion->getPatternType()->getTypeSequenceParameters(packReferences);
+
+      if (packReferences.size() > 1) {
+        auto first = packReferences.begin();
+        auto second = first + 1;
+        while (second != packReferences.end()) {
+          Requirement req(RequirementKind::SameCount, *first++, *second++);
+          desugarRequirement(req, SourceLoc(), reqs, errors);
+        }
+      }
+    }
+
     // Infer requirements from `@differentiable` function types.
     // For all non-`@noDerivative` parameter and result types:
     // - `@differentiable`, `@differentiable(_forward)`, or

lib/AST/RequirementMachine/RuleBuilder.cpp

@@ -299,7 +299,8 @@ void RuleBuilder::addRequirement(const Requirement &req,
 
   switch (req.getKind()) {
   case RequirementKind::SameCount:
-      llvm_unreachable("Same-count requirement not supported here");
+    // TODO
+    return;
 
   case RequirementKind::Conformance: {
     // A conformance requirement T : P becomes a rewrite rule