Merge pull request #76910 from slavapestov/open-type-perf-fix

slavapestov · web-flow · commit fc7068ee402a · 2024-10-08T16:17:58.000-04:00
Fix performance regression in ConstraintSystem::openType()
diff --git a/include/swift/AST/TypeTransform.h b/include/swift/AST/TypeTransform.h
@@ -25,6 +25,12 @@ namespace swift {
 
 template<typename Derived>
 class TypeTransform {
+public:
+  ASTContext &ctx;
+
+  TypeTransform(ASTContext &ctx) : ctx(ctx) {}
+
+private:
   Derived &asDerived() { return *static_cast<Derived *>(this); }
 
   /// \param pos The variance position of the result type.
@@ -98,7 +104,6 @@ class TypeTransform {
 
     // Recur into children of this type.
     TypeBase *base = t.getPointer();
-    auto &ctx = base->getASTContext();
 
     switch (base->getKind()) {
 #define BUILTIN_TYPE(Id, Parent) \
@@ -467,9 +472,27 @@ case TypeKind::Id:
         return newUnderlyingTy;
 
       auto oldSubMap = alias->getSubstitutionMap();
-      auto newSubMap = asDerived().transformSubMap(oldSubMap);
-      if (oldSubMap && !newSubMap)
-        return Type();
+      SubstitutionMap newSubMap;
+
+      // We leave the old behavior behind for ConstraintSystem::openType(), where
+      // preserving sugar introduces a performance penalty.
+      if (asDerived().shouldDesugarTypeAliases()) {
+        for (auto oldReplacementType : oldSubMap.getReplacementTypes()) {
+          Type newReplacementType = doIt(oldReplacementType, TypePosition::Invariant);
+          if (!newReplacementType)
+            return Type();
+
+          // If anything changed with the replacement type, we lose the sugar.
+          if (newReplacementType.getPointer() != oldReplacementType.getPointer())
+            return newUnderlyingTy;
+        }
+
+        newSubMap = oldSubMap;
+      } else {
+        newSubMap = asDerived().transformSubMap(oldSubMap);
+        if (oldSubMap && !newSubMap)
+          return Type();
+      }
 
       if (oldParentType.getPointer() == newParentType.getPointer() &&
           oldUnderlyingTy.getPointer() == newUnderlyingTy.getPointer() &&
@@ -684,14 +707,16 @@ case TypeKind::Id:
       if (!anyChanged)
         return t;
 
-      // Handle vanishing tuples -- If the transform would yield a singleton
-      // tuple, and we didn't start with one, flatten to produce the
-      // element type.
-      if (elements.size() == 1 &&
-          !elements[0].getType()->is<PackExpansionType>() &&
-          !(tuple->getNumElements() == 1 &&
-            !tuple->getElementType(0)->is<PackExpansionType>())) {
-        return elements[0].getType();
+      if (asDerived().shouldUnwrapVanishingTuples()) {
+        // Handle vanishing tuples -- If the transform would yield a singleton
+        // tuple, and we didn't start with one, flatten to produce the
+        // element type.
+        if (elements.size() == 1 &&
+            !elements[0].getType()->is<PackExpansionType>() &&
+            !(tuple->getNumElements() == 1 &&
+              !tuple->getElementType(0)->is<PackExpansionType>())) {
+          return elements[0].getType();
+        }
       }
 
       return TupleType::get(elements, ctx);
@@ -1002,6 +1027,10 @@ case TypeKind::Id:
     return SubstitutionMap::get(sig, newSubs, LookUpConformanceInModule());
   }
 
+  bool shouldUnwrapVanishingTuples() const { return true; }
+
+  bool shouldDesugarTypeAliases() const { return false; }
+
   CanType transformSILField(CanType fieldTy, TypePosition pos) {
     return doIt(fieldTy, pos)->getCanonicalType();
   }
diff --git a/include/swift/Sema/ConstraintSystem.h b/include/swift/Sema/ConstraintSystem.h
@@ -4318,7 +4318,6 @@ class ConstraintSystem {
   Type openType(Type type, OpenedTypeMap &replacements,
                 ConstraintLocatorBuilder locator);
 
-private:
   /// "Open" an opaque archetype type, similar to \c openType.
   Type openOpaqueType(OpaqueTypeArchetypeType *type,
                       ConstraintLocatorBuilder locator);
@@ -4340,7 +4339,6 @@ class ConstraintSystem {
     ASSERT(erased);
   }
 
-public:
   /// Recurse over the given type and open any opaque archetype types.
   Type openOpaqueType(Type type, ContextualTypePurpose context,
                       ConstraintLocatorBuilder locator);
diff --git a/lib/AST/GenericEnvironment.cpp b/lib/AST/GenericEnvironment.cpp
@@ -524,7 +524,8 @@ namespace {
 struct MapTypeIntoContext: TypeTransform<MapTypeIntoContext> {
   GenericEnvironment *env;
 
-  explicit MapTypeIntoContext(GenericEnvironment *env) : env(env) {}
+  explicit MapTypeIntoContext(GenericEnvironment *env, ASTContext &ctx)
+    : TypeTransform(ctx), env(env) {}
 
   std::optional<Type> transform(TypeBase *type, TypePosition pos) {
     if (!type->hasTypeParameter())
@@ -552,7 +553,10 @@ struct MapTypeIntoContext: TypeTransform<MapTypeIntoContext> {
 
 Type GenericEnvironment::mapTypeIntoContext(Type type) const {
   assert(!type->hasPrimaryArchetype() && "already have a contextual type");
-  return MapTypeIntoContext(const_cast<GenericEnvironment *>(this))
+  if (!type->hasTypeParameter())
+    return type;
+  return MapTypeIntoContext(const_cast<GenericEnvironment *>(this),
+                            type->getASTContext())
       .doIt(type, TypePosition::Invariant);
 }
 
diff --git a/lib/AST/Type.cpp b/lib/AST/Type.cpp
@@ -4092,14 +4092,15 @@ Type Type::transformRec(
   class Transform : public TypeTransform<Transform> {
     llvm::function_ref<std::optional<Type>(TypeBase *)> fn;
   public:
-    explicit Transform(llvm::function_ref<std::optional<Type>(TypeBase *)> fn) : fn(fn) {}
+    explicit Transform(llvm::function_ref<std::optional<Type>(TypeBase *)> fn,
+                       ASTContext &ctx) : TypeTransform(ctx), fn(fn) {}
 
     std::optional<Type> transform(TypeBase *type, TypePosition position) {
       return fn(type);
     }
   };
 
-  return Transform(fn).doIt(*this, TypePosition::Invariant);
+  return Transform(fn, (*this)->getASTContext()).doIt(*this, TypePosition::Invariant);
 }
 
 Type Type::transformWithPosition(
@@ -4109,14 +4110,15 @@ Type Type::transformWithPosition(
   class Transform : public TypeTransform<Transform> {
     llvm::function_ref<std::optional<Type>(TypeBase *, TypePosition)> fn;
   public:
-    explicit Transform(llvm::function_ref<std::optional<Type>(TypeBase *, TypePosition)> fn) : fn(fn) {}
+    explicit Transform(llvm::function_ref<std::optional<Type>(TypeBase *, TypePosition)> fn,
+                       ASTContext &ctx) : TypeTransform(ctx), fn(fn) {}
 
     std::optional<Type> transform(TypeBase *type, TypePosition position) {
       return fn(type, position);
     }
   };
 
-  return Transform(fn).doIt(*this, pos);
+  return Transform(fn, (*this)->getASTContext()).doIt(*this, pos);
 }
 
 bool Type::findIf(llvm::function_ref<bool(Type)> pred) const {
diff --git a/lib/AST/TypeSubstitution.cpp b/lib/AST/TypeSubstitution.cpp
@@ -388,8 +388,8 @@ class TypeSubstituter : public TypeTransform<TypeSubstituter> {
   InFlightSubstitution &IFS;
 
 public:
-  TypeSubstituter(unsigned level, InFlightSubstitution &IFS)
-    : level(level), IFS(IFS) {}
+  TypeSubstituter(ASTContext &ctx, unsigned level, InFlightSubstitution &IFS)
+    : TypeTransform(ctx), level(level), IFS(IFS) {}
 
   std::optional<Type> transform(TypeBase *type, TypePosition pos);
 
@@ -566,7 +566,7 @@ Type Type::subst(InFlightSubstitution &IFS) const {
   if (IFS.isInvariant(*this))
     return *this;
 
-  TypeSubstituter transform(/*level=*/0, IFS);
+  TypeSubstituter transform((*this)->getASTContext(), /*level=*/0, IFS);
   return transform.doIt(*this, TypePosition::Invariant);
 }
 
diff --git a/lib/Sema/ConstraintSystem.cpp b/lib/Sema/ConstraintSystem.cpp
@@ -31,6 +31,7 @@
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/TypeCheckRequests.h"
+#include "swift/AST/TypeTransform.h"
 #include "swift/Basic/Assertions.h"
 #include "swift/Basic/Defer.h"
 #include "swift/Basic/Statistic.h"
@@ -1199,49 +1200,63 @@ Type ConstraintSystem::replaceInferableTypesWithTypeVars(
   return type;
 }
 
+namespace {
+
+struct TypeOpener : public TypeTransform<TypeOpener> {
+  OpenedTypeMap &replacements;
+  ConstraintLocatorBuilder locator;
+  ConstraintSystem &cs;
+
+  TypeOpener(OpenedTypeMap &replacements,
+             ConstraintLocatorBuilder locator,
+             ConstraintSystem &cs)
+      : TypeTransform<TypeOpener>(cs.getASTContext()),
+        replacements(replacements), locator(locator), cs(cs) {}
+
+  std::optional<Type> transform(TypeBase *type, TypePosition pos) {
+    if (!type->hasTypeParameter())
+      return Type(type);
+
+    return std::nullopt;
+  }
+
+  Type transformGenericTypeParamType(GenericTypeParamType *genericParam,
+                                     TypePosition pos) {
+    auto known = replacements.find(
+      cast<GenericTypeParamType>(genericParam->getCanonicalType()));
+    // FIXME: This should be an assert, however protocol generic signatures
+    // drop outer generic parameters.
+    // assert(known != replacements.end());
+    if (known == replacements.end())
+      return ErrorType::get(ctx);
+    return known->second;
+  }
+
+  Type transformPackExpansionType(PackExpansionType *expansion,
+                                  TypePosition pos) {
+    return cs.openPackExpansionType(expansion, replacements, locator);
+  }
+
+  bool shouldUnwrapVanishingTuples() const {
+    return false;
+  }
+
+  bool shouldDesugarTypeAliases() const {
+    return true;
+  }
+};
+
+}
+
 Type ConstraintSystem::openType(Type type, OpenedTypeMap &replacements,
                                 ConstraintLocatorBuilder locator) {
   assert(!type->hasUnboundGenericType());
 
   if (!type->hasTypeParameter())
     return type;
 
-  return type.transformRec([&](Type type) -> std::optional<Type> {
-      assert(!type->is<GenericFunctionType>());
-
-      // Preserve single element tuples if their element is
-      // pack expansion, otherwise it wouldn't be expanded.
-      if (auto *tuple = type->getAs<TupleType>()) {
-        if (tuple->getNumElements() == 1) {
-          const auto &elt = tuple->getElement(0);
-          if (!elt.hasName() && elt.getType()->is<PackExpansionType>()) {
-            return TupleType::get(
-                {openPackExpansionType(
-                    elt.getType()->castTo<PackExpansionType>(), replacements,
-                    locator)},
-                tuple->getASTContext());
-          }
-        }
-      }
-
-      if (auto *expansion = type->getAs<PackExpansionType>()) {
-        return openPackExpansionType(expansion, replacements, locator);
-      }
-
-      // Replace a generic type parameter with its corresponding type variable.
-      if (auto genericParam = type->getAs<GenericTypeParamType>()) {
-        auto known = replacements.find(
-          cast<GenericTypeParamType>(genericParam->getCanonicalType()));
-        // FIXME: This should be an assert, however protocol generic signatures
-        // drop outer generic parameters.
-        // assert(known != replacements.end());
-        if (known == replacements.end())
-          return ErrorType::get(getASTContext());
-        return known->second;
-      }
-
-      return std::nullopt;
-    });
+  return TypeOpener(replacements, locator, *this)
+      .doIt(type, TypePosition::Invariant);
 }
 
 Type ConstraintSystem::openPackExpansionType(PackExpansionType *expansion,
diff --git a/test/IDE/complete_call_arg.swift b/test/IDE/complete_call_arg.swift
@@ -1119,7 +1119,7 @@ func testLValueBaseTyOfSubscript() {
   var cache: [String: Codable] = [:]
   if let cached = cache[#^LVALUEBASETY^#
 
-  // LVALUEBASETY-DAG: Decl[Subscript]/CurrNominal/Flair[ArgLabels]/IsSystem: ['[']{#(position): Dictionary<String, any Codable>.Index#}[']'][#Dictionary<String, any Codable>.Element#];
+  // LVALUEBASETY-DAG: Decl[Subscript]/CurrNominal/Flair[ArgLabels]/IsSystem: ['[']{#(position): Dictionary<String, any Codable>.Index#}[']'][#(key: String, value: any Codable)#];
   // LVALUEBASETY-DAG: Decl[Subscript]/CurrNominal/Flair[ArgLabels]/IsSystem: ['[']{#(key): String#}[']'][#@lvalue (any Codable)?#];
 }
 
diff --git a/test/Macros/macro_default_argument_diagnostics.swift b/test/Macros/macro_default_argument_diagnostics.swift
@@ -55,5 +55,5 @@ func testIdentifier(notOkay: Stringified<String> = #stringify(myString)) {}
 // expected-error@+1{{only literals are permitted}}
 func testString(interpolated: Stringified<String> = #stringify("Hello \(0b10001)")) {}
 
-// expected-error@+1{{default argument value of type 'Stringified<Int>' (aka '(Int, String)') cannot be converted to type 'Int'}}
+// expected-error@+1{{default argument value of type '(Int, String)' cannot be converted to type 'Int'}}
 func testReturn(wrongType: Int = #stringify(0)) {}

Original file line number	Diff line number	Diff line change
`@@ -1119,7 +1119,7 @@ func testLValueBaseTyOfSubscript() {`
`1119`	`1119`	`var cache: [String: Codable] = [:]`
`1120`	`1120`	`if let cached = cache[#^LVALUEBASETY^#`
`1121`	`1121`
`1122`		`- // LVALUEBASETY-DAG: Decl[Subscript]/CurrNominal/Flair[ArgLabels]/IsSystem: ['[']{#(position): Dictionary<String, any Codable>.Index#}[']'][#Dictionary<String, any Codable>.Element#];`
	`1122`	`+ // LVALUEBASETY-DAG: Decl[Subscript]/CurrNominal/Flair[ArgLabels]/IsSystem: ['[']{#(position): Dictionary<String, any Codable>.Index#}[']'][#(key: String, value: any Codable)#];`
`1123`	`1123`	`// LVALUEBASETY-DAG: Decl[Subscript]/CurrNominal/Flair[ArgLabels]/IsSystem: ['[']{#(key): String#}[']'][#@lvalue (any Codable)?#];`
`1124`	`1124`	`}`
`1125`	`1125`