[RequirementMachine] Map same-shape requriements to rewrite rules.

A same-shape requirement 'length(T...) == length(U...)' becomes a rewrite
rule 'T.[shape] => U.[shape]'. Reduced shape rules will drop the [shape]
term from each side of the rule, and create a same-shape requirement
between the two type parameter packs.

Author: hborla

lib/AST/RequirementMachine/RequirementBuilder.cpp

@@ -250,6 +250,23 @@ void RequirementBuilder::addRequirementRules(ArrayRef<unsigned> rules) {
       llvm_unreachable("Invalid symbol kind");
     }
 
+    if (rule.getLHS().back().getKind() == Symbol::Kind::Shape) {
+      assert(rule.getRHS().back().getKind() == Symbol::Kind::Shape);
+
+      // Strip off the shape symbols from either side of the rule.
+      MutableTerm lhsTerm(rule.getLHS().begin(),
+                          rule.getLHS().end() - 1);
+      MutableTerm rhsTerm(rule.getRHS().begin(),
+                          rule.getRHS().end() - 1);
+
+      // Add a SameCount requirement between the two parameter packs.
+      auto constraintType = Map.getTypeForTerm(lhsTerm, GenericParams);
+      auto subjectType = Map.getTypeForTerm(rhsTerm, GenericParams);
+      Reqs.emplace_back(RequirementKind::SameCount,
+                        subjectType, constraintType);
+      return;
+    }
+
     assert(rule.getLHS().back().getKind() != Symbol::Kind::Protocol);
     auto constraintType = Map.getTypeForTerm(rule.getLHS(), GenericParams);
     auto subjectType = Map.getTypeForTerm(rule.getRHS(), GenericParams);

lib/AST/RequirementMachine/RuleBuilder.cpp

@@ -298,9 +298,25 @@ void RuleBuilder::addRequirement(const Requirement &req,
   MutableTerm constraintTerm;
 
   switch (req.getKind()) {
-  case RequirementKind::SameCount:
-    // TODO
-    return;
+  case RequirementKind::SameCount: {
+    // A same-shape requirement length(T...) == length(U...)
+    // becomes a rewrite rule:
+    //
+    //    T.[shape] => U.[shape]
+    auto otherType = CanType(req.getSecondType());
+    assert(otherType->isTypeSequenceParameter());
+
+    constraintTerm = (substitutions
+                      ? Context.getRelativeTermForType(
+                            otherType, *substitutions)
+                      : Context.getMutableTermForType(
+                            otherType, proto));
+
+    // Add the [shape] symbol to both sides.
+    subjectTerm.add(Symbol::forShape(Context));
+    constraintTerm.add(Symbol::forShape(Context));
+    break;
+  }
 
   case RequirementKind::Conformance: {
     // A conformance requirement T : P becomes a rewrite rule

lib/AST/RequirementMachine/Symbol.cpp

@@ -72,6 +72,14 @@ struct Symbol::Storage final
     GenericParam = param;
   }
 
+  /// A dummy type for overload resolution of the
+  /// 'shape' constructor for Storage.
+  struct ForShape {};
+
+  explicit Storage(ForShape shape) {
+    Kind = Kind::Shape;
+  }
+
   Storage(const ProtocolDecl *proto, Identifier name) {
     Kind = Symbol::Kind::AssociatedType;
     Proto = proto;
@@ -298,6 +306,30 @@ Symbol Symbol::forGenericParam(GenericTypeParamType *param,
   return symbol;
 }
 
+Symbol Symbol::forShape(RewriteContext &ctx) {
+  llvm::FoldingSetNodeID id;
+  id.AddInteger(unsigned(Kind::Shape));
+
+  void *insertPos = nullptr;
+  if (auto *symbol = ctx.Symbols.FindNodeOrInsertPos(id, insertPos))
+    return symbol;
+
+  unsigned size = Storage::totalSizeToAlloc<unsigned, Term>(0, 0);
+  void *mem = ctx.Allocator.Allocate(size, alignof(Storage));
+  auto *symbol = new (mem) Storage(Storage::ForShape());
+
+#ifndef NDEBUG
+  llvm::FoldingSetNodeID newID;
+  symbol->Profile(newID);
+  assert(id == newID);
+#endif
+
+  ctx.Symbols.InsertNode(symbol, insertPos);
+  ctx.SymbolHistogram.add(unsigned(Kind::Shape));
+
+  return symbol;
+}
+
 /// Creates a layout symbol, representing a layout constraint.
 Symbol Symbol::forLayout(LayoutConstraint layout,
                          RewriteContext &ctx) {

lib/AST/RequirementMachine/Symbol.h

@@ -204,6 +204,8 @@ class Symbol final {
   static Symbol forGenericParam(GenericTypeParamType *param,
                                 RewriteContext &ctx);
 
+  static Symbol forShape(RewriteContext &ctx);
+
   static Symbol forLayout(LayoutConstraint layout,
                           RewriteContext &ctx);
 

test/Generics/pack-shape-requirements.swift

@@ -0,0 +1,15 @@
+// RUN: %target-swift-frontend -typecheck -enable-experimental-variadic-generics %s -debug-generic-signatures 2>&1 | %FileCheck %s
+
+protocol P {
+  associatedtype A
+}
+
+// CHECK-LABEL: inferSameShape(ts:us:)
+// CHECK-NEXT: Generic signature: <@_typeSequence T, @_typeSequence U where T.count == U.count>
+func inferSameShape<@_typeSequence T, @_typeSequence U>(ts t: T..., us u: U...) where ((T, U)...): Any {
+}
+
+// CHECK-LABEL: desugarSameShape(ts:us:)
+// CHECK-NEXT: Generic signature: <@_typeSequence T, @_typeSequence U where T : P, T.count == U.count, U : P>
+func desugarSameShape<@_typeSequence T, @_typeSequence U>(ts t: T..., us u: U...) where T: P, U: P, ((T.A, U.A)...): Any {
+}