AST: Move getReducedShape() from CSSimplify.cpp to a method on TypeBase

Author: slavapestov

include/swift/AST/Types.h

@@ -1290,6 +1290,18 @@ class alignas(1 << TypeAlignInBits) TypeBase
   /// Error.
   bool isExistentialWithError();
 
+  /// Returns the reduced shape of the type, which represents an equivalence
+  /// class for the same-shape generic requirement:
+  ///
+  /// - The shape of a scalar type is always the empty tuple type ().
+  /// - The shape of a pack archetype is computed from the generic signature
+  ///   using same-shape requirements.
+  /// - The shape of a pack type is computed recursively from its elements.
+  ///
+  /// Two types satisfy a same-shape requirement if their reduced shapes are
+  /// equal as canonical types.
+  CanType getReducedShape();
+
   SWIFT_DEBUG_DUMP;
   void dump(raw_ostream &os, unsigned indent = 0) const;
 
@@ -6028,7 +6040,7 @@ class PackArchetypeType final
       LayoutConstraint Layout);
 
   // Returns the reduced shape type for this pack archetype.
-  Type getShape() const;
+  CanType getReducedShape() const;
 
   static bool classof(const TypeBase *T) {
     return T->getKind() == TypeKind::PackArchetype;
@@ -6481,6 +6493,8 @@ class PackType final : public TypeBase, public llvm::FoldingSetNode,
 
   PackType *flattenPackTypes();
 
+  CanTypeWrapper<PackType> getReducedShape();
+
 public:
   void Profile(llvm::FoldingSetNodeID &ID) const {
     Profile(ID, getElementTypes());
@@ -6553,6 +6567,8 @@ class PackExpansionType : public TypeBase, public llvm::FoldingSetNode {
 
   PackExpansionType *expand();
 
+  CanType getReducedShape();
+
 public:
   void Profile(llvm::FoldingSetNodeID &ID) {
     Profile(ID, getPatternType(), getCountType());

lib/AST/ParameterPack.cpp

@@ -15,10 +15,11 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "swift/AST/Types.h"
+#include "swift/AST/ASTContext.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/Type.h"
+#include "swift/AST/Types.h"
 #include "llvm/ADT/SmallVector.h"
 
 using namespace swift;
@@ -123,6 +124,19 @@ PackExpansionType *PackExpansionType::expand() {
   return PackExpansionType::get(packType, countType);
 }
 
+CanType PackExpansionType::getReducedShape() {
+  if (auto *archetypeType = countType->getAs<PackArchetypeType>()) {
+    auto shape = archetypeType->getReducedShape();
+    return CanType(PackExpansionType::get(shape, shape));
+  } else if (auto *packType = countType->getAs<PackType>()) {
+    auto shape = packType->getReducedShape();
+    return CanType(PackExpansionType::get(shape, shape));
+  }
+
+  assert(countType->is<PlaceholderType>());
+  return getASTContext().TheEmptyTupleType;
+}
+
 bool TupleType::containsPackExpansionType() const {
   for (auto elt : getElements()) {
     if (elt.getType()->is<PackExpansionType>())
@@ -266,6 +280,44 @@ PackType *PackType::flattenPackTypes() {
   return PackType::get(getASTContext(), elts);
 }
 
+CanPackType PackType::getReducedShape() {
+  SmallVector<Type, 4> elts;
+
+  auto &ctx = getASTContext();
+
+  for (auto elt : getElementTypes()) {
+    // T... => shape(T)...
+    if (auto *packExpansionType = elt->getAs<PackExpansionType>()) {
+      elts.push_back(packExpansionType->getReducedShape());
+      continue;
+    }
+
+    // Use () as a placeholder for scalar shape.
+    assert(!elt->is<PackArchetypeType>() &&
+           "Pack archetype outside of a pack expansion");
+    elts.push_back(ctx.TheEmptyTupleType);
+  }
+
+  return CanPackType(PackType::get(ctx, elts));
+}
+
+CanType TypeBase::getReducedShape() {
+  if (auto *packArchetype = getAs<PackArchetypeType>())
+    return packArchetype->getReducedShape();
+
+  if (auto *packType = getAs<PackType>())
+    return packType->getReducedShape();
+
+  if (auto *expansionType = getAs<PackExpansionType>())
+    return expansionType->getReducedShape();
+
+  assert(!isTypeVariableOrMember());
+  assert(!hasTypeParameter());
+
+  // Use () as a placeholder for scalar shape.
+  return getASTContext().TheEmptyTupleType;
+}
+
 unsigned ParameterList::getOrigParamIndex(SubstitutionMap subMap,
                                           unsigned substIndex) const {
   unsigned remappedIndex = substIndex;

lib/AST/Type.cpp

@@ -4013,9 +4013,9 @@ PackArchetypeType::get(const ASTContext &Ctx,
       {ShapeType}));
 }
 
-Type PackArchetypeType::getShape() const {
+CanType PackArchetypeType::getReducedShape() const {
   auto shapeType = getTrailingObjects<PackShape>()->shapeType;
-  return getGenericEnvironment()->mapTypeIntoContext(shapeType);
+  return getGenericEnvironment()->mapTypeIntoContext(shapeType)->getCanonicalType();
 }
 
 ElementArchetypeType::ElementArchetypeType(

lib/Sema/CSGen.cpp

@@ -2927,7 +2927,7 @@ namespace {
           return;
 
         if (auto archetype = type->getAs<PackArchetypeType>()) {
-          shapeType = archetype->getShape();
+          shapeType = archetype->getReducedShape();
         }
       });
 

lib/Sema/CSSimplify.cpp

@@ -12603,40 +12603,6 @@ ConstraintSystem::simplifyDynamicCallableApplicableFnConstraint(
   return SolutionKind::Solved;
 }
 
-/// FIXME: Move this elsewhere if it is broadly useful
-static Type getReducedShape(Type type, ASTContext &ctx) {
-  // 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>()) {
-    SmallVector<Type, 2> elts;
-
-    for (auto elt : packType->getElementTypes()) {
-      // T... => shape(T)...
-      if (auto *packExpansionType = elt->getAs<PackExpansionType>()) {
-        if (packExpansionType->getCountType()->is<PlaceholderType>()) {
-          elts.push_back(ctx.TheEmptyTupleType);
-          continue;
-        }
-        auto shapeType = getReducedShape(packExpansionType->getCountType(), ctx);
-        elts.push_back(PackExpansionType::get(shapeType, shapeType));
-      }
-
-      // Use () as a placeholder for scalar shape.
-      elts.push_back(ctx.TheEmptyTupleType);
-    }
-
-    return PackType::get(ctx, elts);
-  }
-
-  assert(!type->isTypeVariableOrMember());
-
-  // Use () as a placeholder for scalar shape.
-  return ctx.TheEmptyTupleType;
-}
-
 ConstraintSystem::SolutionKind ConstraintSystem::simplifyShapeOfConstraint(
     Type type1, Type type2, TypeMatchOptions flags,
     ConstraintLocatorBuilder locator) {
@@ -12667,12 +12633,9 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyShapeOfConstraint(
       return formUnsolved();
   }
 
-  if (Type shape = getReducedShape(type1, getASTContext())) {
-    addConstraint(ConstraintKind::Bind, shape, type2, locator);
-    return SolutionKind::Solved;
-  }
-
-  return SolutionKind::Error;
+  auto shape = type1->getReducedShape();
+  addConstraint(ConstraintKind::Bind, shape, type2, locator);
+  return SolutionKind::Solved;
 }
 
 static llvm::PointerIntPair<Type, 3, unsigned>