AST: Store the ASTContext inside the TypeTransform

Author: slavapestov

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) \
@@ -684,14 +689,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 +1009,8 @@ case TypeKind::Id:
     return SubstitutionMap::get(sig, newSubs, LookUpConformanceInModule());
   }
 
+  bool shouldUnwrapVanishingTuples() const { return true; }
+
   CanType transformSILField(CanType fieldTy, TypePosition pos) {
     return doIt(fieldTy, pos)->getCanonicalType();
   }

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);
 }
 

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 {

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);
 }