AST: Type substitution can skip subtrees that won't change

include/swift/AST/InFlightSubstitution.h

@@ -31,6 +31,7 @@ class InFlightSubstitution {
   SubstOptions Options;
   TypeSubstitutionFn BaselineSubstType;
   LookupConformanceFn BaselineLookupConformance;
+  RecursiveTypeProperties Props;
 
   struct ActivePackExpansion {
     bool isSubstExpansion = false;
@@ -41,10 +42,7 @@ class InFlightSubstitution {
 public:
   InFlightSubstitution(TypeSubstitutionFn substType,
                        LookupConformanceFn lookupConformance,
-                       SubstOptions options)
-    : Options(options),
-      BaselineSubstType(substType),
-      BaselineLookupConformance(lookupConformance) {}
+                       SubstOptions options);
 
   InFlightSubstitution(const InFlightSubstitution &) = delete;
   InFlightSubstitution &operator=(const InFlightSubstitution &) = delete;

include/swift/AST/LocalArchetypeRequirementCollector.h

@@ -54,6 +54,10 @@ struct MapLocalArchetypesOutOfContext {
   Type operator()(SubstitutableType *type) const;
 };
 
+Type mapLocalArchetypesOutOfContext(Type type,
+                                    GenericSignature baseGenericSig,
+                                    ArrayRef<GenericEnvironment *> capturedEnvs);
+
 struct MapIntoLocalArchetypeContext {
   GenericEnvironment *baseGenericEnv;
   ArrayRef<GenericEnvironment *> capturedEnvs;

include/swift/AST/Types.h

@@ -6577,10 +6577,6 @@ class ArchetypeType : public SubstitutableType,
   /// which the archetype conforms.
   Type getNestedTypeByName(Identifier name);
 
-  /// Return the archetype that represents the root generic parameter of its
-  /// interface type.
-  ArchetypeType *getRoot() const;
-
   /// Determine whether this is a root archetype within the environment.
   bool isRoot() const;
 
@@ -6761,19 +6757,12 @@ class LocalArchetypeType : public ArchetypeType {
   using ArchetypeType::ArchetypeType;
 
 public:
-  LocalArchetypeType *getRoot() const {
-    return cast<LocalArchetypeType>(ArchetypeType::getRoot());
-  }
-
   static bool classof(const TypeBase *type) {
     return type->getKind() == TypeKind::OpenedArchetype ||
            type->getKind() == TypeKind::ElementArchetype;
   }
 };
 BEGIN_CAN_TYPE_WRAPPER(LocalArchetypeType, ArchetypeType)
-  CanLocalArchetypeType getRoot() const {
-    return CanLocalArchetypeType(getPointer()->getRoot());
-  }
 END_CAN_TYPE_WRAPPER(LocalArchetypeType, ArchetypeType)
 
 /// An archetype that represents the dynamic type of an opened existential.
@@ -6828,12 +6817,6 @@ class OpenedArchetypeType final : public LocalArchetypeType,
   /// Retrieve the ID number of this opened existential.
   UUID getOpenedExistentialID() const;
 
-  /// Return the archetype that represents the root generic parameter of its
-  /// interface type.
-  OpenedArchetypeType *getRoot() const {
-    return cast<OpenedArchetypeType>(ArchetypeType::getRoot());
-  }
-
   static bool classof(const TypeBase *T) {
     return T->getKind() == TypeKind::OpenedArchetype;
   }
@@ -6846,9 +6829,6 @@ class OpenedArchetypeType final : public LocalArchetypeType,
                       RecursiveTypeProperties properties);
 };
 BEGIN_CAN_TYPE_WRAPPER(OpenedArchetypeType, LocalArchetypeType)
-  CanOpenedArchetypeType getRoot() const {
-    return CanOpenedArchetypeType(getPointer()->getRoot());
-  }
 END_CAN_TYPE_WRAPPER(OpenedArchetypeType, LocalArchetypeType)
 
 /// A wrapper around a shape type to use in ArchetypeTrailingObjects
@@ -6916,12 +6896,6 @@ class ElementArchetypeType final : public LocalArchetypeType,
   /// Retrieve the ID number of this opened element.
   UUID getOpenedElementID() const;
 
-  /// Return the archetype that represents the root generic parameter of its
-  /// interface type.
-  ElementArchetypeType *getRoot() const {
-    return cast<ElementArchetypeType>(ArchetypeType::getRoot());
-  }
-
   static bool classof(const TypeBase *T) {
     return T->getKind() == TypeKind::ElementArchetype;
   }
@@ -6933,9 +6907,6 @@ class ElementArchetypeType final : public LocalArchetypeType,
                        LayoutConstraint layout);
 };
 BEGIN_CAN_TYPE_WRAPPER(ElementArchetypeType, LocalArchetypeType)
-  CanElementArchetypeType getRoot() const {
-    return CanElementArchetypeType(getPointer()->getRoot());
-  }
 END_CAN_TYPE_WRAPPER(ElementArchetypeType, LocalArchetypeType)
 
 template<typename Type>

include/swift/SIL/SILCloner.h

@@ -60,9 +60,8 @@ struct SubstitutionMapWithLocalArchetypes {
 
         // Map the local archetype to an interface type in the new generic
         // signature.
-        MapLocalArchetypesOutOfContext mapOutOfContext(BaseGenericSig,
-                                                       CapturedEnvs);
-        auto interfaceTy = mapOutOfContext(local);
+        auto interfaceTy = mapLocalArchetypesOutOfContext(
+            local, BaseGenericSig, CapturedEnvs);
 
         // Map this interface type into the new generic environment to get
         // a primary archetype.

include/swift/SIL/SILModule.h

@@ -460,6 +460,18 @@ class SILModule {
   SILValue getLocalGenericEnvironmentDef(GenericEnvironment *genericEnv,
                                          SILFunction *inFunction);
 
+  /// Returns the instruction which defines the given local generic environment,
+  /// e.g. an open_existential_addr.
+  ///
+  /// In contrast to getLocalGenericEnvironmentDef, it is required that all local
+  /// generic environments are resolved.
+  SingleValueInstruction *
+  getLocalGenericEnvironmentDefInst(GenericEnvironment *genericEnv,
+                                    SILFunction *inFunction) {
+    return dyn_cast<SingleValueInstruction>(
+        getLocalGenericEnvironmentDef(genericEnv, inFunction));
+  }
+
   /// Returns the instruction which defines the given root local archetype,
   /// e.g. an open_existential_addr.
   ///

lib/AST/ASTMangler.cpp

@@ -3796,12 +3796,8 @@ void ASTMangler::appendClosureComponents(CanType Ty, unsigned discriminator,
   appendContext(parentContext, base, StringRef());
 
   auto Sig = parentContext->getGenericSignatureOfContext();
-
-  Ty = Ty.subst(MapLocalArchetypesOutOfContext(Sig, capturedEnvs),
-                MakeAbstractConformanceForGenericType(),
-                SubstFlags::PreservePackExpansionLevel |
-                SubstFlags::SubstitutePrimaryArchetypes |
-                SubstFlags::SubstituteLocalArchetypes)->getCanonicalType();
+  Ty = mapLocalArchetypesOutOfContext(Ty, Sig, capturedEnvs)
+      ->getCanonicalType();
 
   appendType(Ty, Sig);
   appendOperator(isImplicit ? "fu" : "fU", Index(discriminator));

lib/AST/LocalArchetypeRequirementCollector.cpp

@@ -203,6 +203,17 @@ Type MapLocalArchetypesOutOfContext::operator()(SubstitutableType *type) const {
   return getInterfaceType(archetypeTy->getInterfaceType(), genericEnv);
 }
 
+Type swift::mapLocalArchetypesOutOfContext(
+    Type type,
+    GenericSignature baseGenericSig,
+    ArrayRef<GenericEnvironment *> capturedEnvs) {
+  return type.subst(MapLocalArchetypesOutOfContext(baseGenericSig, capturedEnvs),
+                    MakeAbstractConformanceForGenericType(),
+                    SubstFlags::PreservePackExpansionLevel |
+                    SubstFlags::SubstitutePrimaryArchetypes |
+                    SubstFlags::SubstituteLocalArchetypes);
+}
+
 static Type mapIntoLocalContext(GenericTypeParamType *param, unsigned baseDepth,
                                 ArrayRef<GenericEnvironment *> capturedEnvs) {
   assert(!param->isParameterPack());

lib/AST/ParameterPack.cpp

@@ -17,6 +17,7 @@
 
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/Decl.h"
+#include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/GenericParamList.h"
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/Type.h"
@@ -161,7 +162,14 @@ void TypeBase::getTypeParameterPacks(
       if (paramTy->isParameterPack())
         rootParameterPacks.push_back(paramTy);
     } else if (auto *archetypeTy = t->getAs<PackArchetypeType>()) {
-      rootParameterPacks.push_back(archetypeTy->getRoot());
+      if (archetypeTy->isRoot()) {
+        rootParameterPacks.push_back(archetypeTy);
+      } else {
+        auto *genericEnv = archetypeTy->getGenericEnvironment();
+        auto paramTy = archetypeTy->getInterfaceType()->getRootGenericParam();
+        rootParameterPacks.push_back(
+            genericEnv->mapTypeIntoContext(paramTy));
+      }
     }
 
     return false;

lib/AST/Type.cpp

@@ -3378,17 +3378,6 @@ ArchetypeType::ArchetypeType(TypeKind Kind,
                           getSubclassTrailingObjects<ProtocolDecl *>());
 }
 
-ArchetypeType *ArchetypeType::getRoot() const {
-  if (isRoot()) {
-    return const_cast<ArchetypeType *>(this);
-  }
-
-  auto gp = InterfaceType->getRootGenericParam();
-  assert(gp && "Missing root generic parameter?");
-  return getGenericEnvironment()->mapTypeIntoContext(
-      Type(gp))->castTo<ArchetypeType>();
-}
-
 bool ArchetypeType::isRoot() const {
   return getInterfaceType()->is<GenericTypeParamType>();
 }

lib/AST/TypeSubstitution.cpp

@@ -359,6 +359,40 @@ static Type substGenericFunctionType(GenericFunctionType *genericFnType,
                                   fnType->getResult(), fnType->getExtInfo());
 }
 
+InFlightSubstitution::InFlightSubstitution(TypeSubstitutionFn substType,
+                                           LookupConformanceFn lookupConformance,
+                                           SubstOptions options)
+  : Options(options),
+    BaselineSubstType(substType),
+    BaselineLookupConformance(lookupConformance) {
+  // FIXME: Don't substitute type parameters if one of the special flags is set.
+  Props |= RecursiveTypeProperties::HasTypeParameter;
+
+  // If none of the special flags are set, we substitute type parameters and
+  // primary archetypes only.
+  if (!Options.contains(SubstFlags::SubstitutePrimaryArchetypes) &&
+      !Options.contains(SubstFlags::SubstituteLocalArchetypes) &&
+      !Options.contains(SubstFlags::SubstituteOpaqueArchetypes)) {
+    Props |= RecursiveTypeProperties::HasPrimaryArchetype;
+  }
+
+  if (Options.contains(SubstFlags::SubstitutePrimaryArchetypes))
+    Props |= RecursiveTypeProperties::HasPrimaryArchetype;
+
+  if (Options.contains(SubstFlags::SubstituteLocalArchetypes)) {
+    Props |= RecursiveTypeProperties::HasOpenedExistential;
+    Props |= RecursiveTypeProperties::HasElementArchetype;
+  }
+
+  if (Options.contains(SubstFlags::SubstituteOpaqueArchetypes))
+    Props |= RecursiveTypeProperties::HasOpaqueArchetype;
+}
+
+bool InFlightSubstitution::isInvariant(Type derivedType) const {
+  // If none of the bits are set, the type won't be changed by substitution.
+  return !(derivedType->getRecursiveProperties().getBits() & Props.getBits());
+}
+
 void InFlightSubstitution::expandPackExpansionShape(Type origShape,
     llvm::function_ref<void(Type substComponentShape)> handleComponent) {
 
@@ -470,16 +504,6 @@ InFlightSubstitution::lookupConformance(CanType dependentType,
   return substPackPatterns[index];
 }
 
-bool InFlightSubstitution::isInvariant(Type derivedType) const {
-  if (derivedType->hasPrimaryArchetype() || derivedType->hasTypeParameter())
-    return false;
-  if (shouldSubstituteLocalArchetypes() && derivedType->hasLocalArchetype())
-    return false;
-  if (shouldSubstituteOpaqueArchetypes() && derivedType->hasOpaqueArchetype())
-    return false;
-  return true;
-}
-
 namespace {
 
 class TypeSubstituter : public TypeTransform<TypeSubstituter> {
@@ -522,6 +546,9 @@ class TypeSubstituter : public TypeTransform<TypeSubstituter> {
 
 std::optional<Type>
 TypeSubstituter::transform(TypeBase *type, TypePosition position) {
+  if (IFS.isInvariant(type))
+    return Type(type);
+
   return std::nullopt;
 }
 

lib/SIL/IR/SILFunctionType.cpp

@@ -1967,14 +1967,11 @@ lowerCaptureContextParameters(TypeConverter &TC, SILDeclRef function,
   auto mapTypeOutOfContext = [&](Type t) -> CanType {
     LLVM_DEBUG(llvm::dbgs() << "-- capture with contextual type " << t << "\n");
 
-    t = t.subst(MapLocalArchetypesOutOfContext(origGenericSig, capturedEnvs),
-                MakeAbstractConformanceForGenericType(),
-                SubstFlags::PreservePackExpansionLevel |
-                SubstFlags::SubstitutePrimaryArchetypes |
-                SubstFlags::SubstituteLocalArchetypes);
-
-    LLVM_DEBUG(llvm::dbgs() << "-- maps to " << t->getCanonicalType() << "\n");
-    return t->getCanonicalType();
+    auto result = mapLocalArchetypesOutOfContext(t, origGenericSig, capturedEnvs)
+        ->getCanonicalType();
+
+    LLVM_DEBUG(llvm::dbgs() << "-- maps to " << result << "\n");
+    return result;
   };
 
   for (auto capture : loweredCaptures.getCaptures()) {
@@ -3004,26 +3001,17 @@ CanSILFunctionType swift::buildSILFunctionThunkType(
                                      interfaceSubs);
   }
 
-  MapLocalArchetypesOutOfContext mapOutOfContext(baseGenericSig, capturedEnvs);
   auto substFormalTypeIntoThunkContext =
       [&](CanType t) -> CanType {
     return genericEnv->mapTypeIntoContext(
-        t.subst(mapOutOfContext,
-                MakeAbstractConformanceForGenericType(),
-                SubstFlags::PreservePackExpansionLevel |
-                SubstFlags::SubstitutePrimaryArchetypes |
-                SubstFlags::SubstituteLocalArchetypes))
+        mapLocalArchetypesOutOfContext(t, baseGenericSig, capturedEnvs))
                ->getCanonicalType();
   };
   auto substLoweredTypeIntoThunkContext =
       [&](CanSILFunctionType t) -> CanSILFunctionType {
     return cast<SILFunctionType>(
         genericEnv->mapTypeIntoContext(
-          Type(t).subst(mapOutOfContext,
-                        MakeAbstractConformanceForGenericType(),
-                        SubstFlags::PreservePackExpansionLevel |
-                        SubstFlags::SubstitutePrimaryArchetypes |
-                        SubstFlags::SubstituteLocalArchetypes))
+          mapLocalArchetypesOutOfContext(t, baseGenericSig, capturedEnvs))
               ->getCanonicalType());
   };
 

lib/SIL/IR/TypeLowering.cpp

@@ -3905,12 +3905,8 @@ getAnyFunctionRefInterfaceType(TypeConverter &TC,
 
   if (funcType->hasArchetype()) {
     assert(isa<FunctionType>(funcType));
-    auto substType = Type(funcType).subst(
-        MapLocalArchetypesOutOfContext(sig.baseGenericSig, sig.capturedEnvs),
-        MakeAbstractConformanceForGenericType(),
-        SubstFlags::PreservePackExpansionLevel |
-        SubstFlags::SubstitutePrimaryArchetypes |
-        SubstFlags::SubstituteLocalArchetypes);
+    auto substType = mapLocalArchetypesOutOfContext(
+        funcType, sig.baseGenericSig, sig.capturedEnvs);
     funcType = cast<FunctionType>(substType->getCanonicalType());
   }
 
@@ -4964,15 +4960,10 @@ TypeConverter::getInterfaceBoxTypeForCapture(ValueDecl *captured,
   SmallSetVector<GenericEnvironment *, 2> boxCapturedEnvs;
   findCapturedEnvironments(loweredContextType, boxCapturedEnvs);
 
-  MapLocalArchetypesOutOfContext mapOutOfContext(baseGenericSig,
-                                                 boxCapturedEnvs.getArrayRef());
-
-  auto loweredInterfaceType = loweredContextType.subst(
-      mapOutOfContext,
-      MakeAbstractConformanceForGenericType(),
-      SubstFlags::PreservePackExpansionLevel |
-      SubstFlags::SubstitutePrimaryArchetypes |
-      SubstFlags::SubstituteLocalArchetypes)->getCanonicalType();
+  auto loweredInterfaceType =
+      mapLocalArchetypesOutOfContext(loweredContextType, baseGenericSig,
+                                     boxCapturedEnvs.getArrayRef())
+      ->getCanonicalType();
 
   // If the type is not dependent at all, we can form a concrete box layout.
   // We don't need to capture the generic environment.