[RequirementMachine] Use pack expansion types as a source for requirement inference.

Previously, only the pattern type was added as a requirement inference source, but
the pack expansion type is necessary for inferring same-shape requirements for
packs that are expanded in parallel.

Author: hborla

lib/AST/RequirementMachine/ConcreteContraction.cpp

@@ -371,8 +371,12 @@ ConcreteContraction::substRequirement(const Requirement &req) const {
   auto firstType = req.getFirstType();
 
   switch (req.getKind()) {
-  case RequirementKind::SameShape:
-    llvm_unreachable("Same-shape requirement not supported here");
+  case RequirementKind::SameShape: {
+    auto substFirstType = substType(firstType);
+    auto substSecondType = substType(req.getSecondType());
+
+    return Requirement(req.getKind(), substFirstType, substSecondType);
+  }
 
   case RequirementKind::Superclass:
   case RequirementKind::SameType: {

lib/Sema/TypeCheckGeneric.cpp

@@ -648,8 +648,6 @@ GenericSignatureRequest::evaluate(Evaluator &evaluator,
           typeRepr = specifier->getBase();
 
         if (auto *packExpansion = dyn_cast<PackExpansionTypeRepr>(typeRepr)) {
-          typeRepr = packExpansion->getPatternType();
-
           paramOptions.setContext(TypeResolverContext::VariadicFunctionInput);
         } else {
           paramOptions.setContext(TypeResolverContext::FunctionInput);

lib/Sema/TypeCheckType.cpp

@@ -4228,8 +4228,13 @@ NeverNullType TypeResolver::resolvePackExpansionType(PackExpansionTypeRepr *repr
     return ErrorType::get(ctx);
   }
 
-  // The pattern type must contain at least one variadic generic parameter.
   if (!pair.second) {
+    // Non-pack variadic parameters parse as PackExpansionTypeReprs. These
+    // can only appear in variadic function parameter types.
+    if (options.getContext() == TypeResolverContext::VariadicFunctionInput)
+      return pair.first;
+
+    // Otherwise, the pattern type must contain at least one pack reference.
     diagnose(repr->getLoc(), diag::expansion_not_variadic, pair.first)
       .highlight(repr->getSourceRange());
     return ErrorType::get(ctx);

lib/Sema/TypeCheckType.h

@@ -314,6 +314,7 @@ class TypeResolutionOptions {
   bool isPackExpansionSupported() const {
     switch (context) {
     case Context::FunctionInput:
+    case Context::VariadicFunctionInput:
     case Context::TupleElement:
     case Context::GenericArgument:
       return true;
@@ -335,7 +336,6 @@ class TypeResolutionOptions {
     case Context::PatternBindingDecl:
     case Context::EditorPlaceholderExpr:
     case Context::ClosureExpr:
-    case Context::VariadicFunctionInput:
     case Context::InoutFunctionInput:
     case Context::FunctionResult:
     case Context::SubscriptDecl:

test/Generics/pack-shape-requirements.swift

@@ -59,3 +59,9 @@ struct Ts<T...> {
     }
   }
 }
+
+// CHECK-LABEL: expandedParameters(_:transform:)
+// CHECK-NEXT: Generic signature: <T..., Result... where T.shape == Result.shape>
+func expandedParameters<T..., Result...>(_ t: T..., transform: ((T) -> Result)...) -> (Result...) {
+  fatalError()
+}