AST: Introduce SubstFlags::SubstituteLocalArchetypes

Author: slavapestov

include/swift/AST/InFlightSubstitution.h

@@ -138,6 +138,10 @@ class InFlightSubstitution {
     return Options.contains(SubstFlags::SubstituteOpaqueArchetypes);
   }
 
+  bool shouldSubstituteLocalArchetypes() const {
+    return Options.contains(SubstFlags::SubstituteLocalArchetypes);
+  }
+
   /// Is the given type invariant to substitution?
   bool isInvariant(Type type) const;
 };

include/swift/AST/Type.h

@@ -117,20 +117,24 @@ class MakeAbstractConformanceForGenericType {
 
 /// Flags that can be passed when substituting into a type.
 enum class SubstFlags {
-  /// Allow substitutions to recurse into SILFunctionTypes.
-  /// Normally, SILType::subst() should be used for lowered
-  /// types, however in special cases where the substitution
-  /// is just changing between contextual and interface type
-  /// representations, using Type::subst() is allowed.
+  /// Allow substitutions to recurse into SILFunctionTypes. Normally,
+  /// SILType::subst() should be used for lowered types, however in special
+  /// cases where the substitution is just changing between contextual and
+  /// interface type representations, using Type::subst() is allowed.
   AllowLoweredTypes = 0x01,
   /// Map member types to their desugared witness type.
   DesugarMemberTypes = 0x02,
-  /// Substitute types involving opaque type archetypes.
+  /// Allow opaque archetypes to themselves be the subject of substitution,
+  /// used when erasing them to their underlying types. otherwise, we
+  /// recursively substitute their substitutions, instead, preserving the
+  /// opaque archetype.
   SubstituteOpaqueArchetypes = 0x04,
+  /// Allow local archetypes to themselves be the subject of substitution.
+  SubstituteLocalArchetypes = 0x08,
   /// Don't increase pack expansion level for free pack references.
   /// Do not introduce new usages of this flag.
   /// FIXME: Remove this.
-  PreservePackExpansionLevel = 0x08,
+  PreservePackExpansionLevel = 0x10,
 };
 
 /// Options for performing substitutions into a type.

include/swift/SIL/SILCloner.h

@@ -57,6 +57,16 @@ struct SubstitutionMapWithLocalArchetypes {
 
     return ProtocolConformanceRef(proto);
   }
+
+  void dump(llvm::raw_ostream &out) const {
+    if (SubsMap)
+      SubsMap->dump(out);
+    for (auto pair : LocalArchetypeSubs) {
+      out << "---\n";
+      pair.first->dump(out);
+      pair.second->dump(out);
+    }
+  }
 };
 
 /// SILCloner - Abstract SIL visitor which knows how to clone instructions and
@@ -231,8 +241,7 @@ class SILCloner : protected SILInstructionVisitor<ImplClass> {
     for (auto substConf : substSubs.getConformances()) {
       if (substConf.isInvalid()) {
         llvm::errs() << "Invalid conformance in SIL cloner:\n";
-        if (Functor.SubsMap)
-          Functor.SubsMap->dump(llvm::errs());
+        Functor.dump(llvm::errs());
         llvm::errs() << "\nsubstitution map:\n";
         Subs.dump(llvm::errs());
         llvm::errs() << "\n";
@@ -363,11 +372,10 @@ class SILCloner : protected SILInstructionVisitor<ImplClass> {
 #ifndef NDEBUG
     if (substConf.isInvalid()) {
       llvm::errs() << "Invalid conformance in SIL cloner:\n";
-      if (Functor.SubsMap)
-        Functor.SubsMap->dump(llvm::errs());
+      Functor.dump(llvm::errs());
       llvm::errs() << "\nconformance:\n";
       conformance.dump(llvm::errs());
-      llvm::errs() << "original type:\n";
+      llvm::errs() << "\noriginal type:\n";
       ty.dump(llvm::errs());
       abort();
     }
@@ -460,8 +468,12 @@ class SILCloner : protected SILInstructionVisitor<ImplClass> {
 
   SILType remapType(SILType Ty) {
     if (Functor.SubsMap || Ty.hasLocalArchetype()) {
+      SubstOptions options(std::nullopt);
+      if (!Functor.LocalArchetypeSubs.empty())
+        options |= SubstFlags::SubstituteLocalArchetypes;
+
       Ty = Ty.subst(Builder.getModule(), Functor, Functor,
-                    CanGenericSignature());
+                    CanGenericSignature(), options);
     }
 
     if (asImpl().shouldSubstOpaqueArchetypes()) {
@@ -480,8 +492,13 @@ class SILCloner : protected SILInstructionVisitor<ImplClass> {
   }
 
   CanType remapASTType(CanType ty) {
-    if (Functor.SubsMap || ty->hasLocalArchetype())
-      ty = ty.subst(Functor, Functor)->getCanonicalType();
+    if (Functor.SubsMap || ty->hasLocalArchetype()) {
+      SubstOptions options(std::nullopt);
+      if (!Functor.LocalArchetypeSubs.empty())
+        options |= SubstFlags::SubstituteLocalArchetypes;
+
+      ty = ty.subst(Functor, Functor, options)->getCanonicalType();
+    }
 
     if (asImpl().shouldSubstOpaqueArchetypes()) {
       auto context = getBuilder().getTypeExpansionContext();
@@ -500,9 +517,13 @@ class SILCloner : protected SILInstructionVisitor<ImplClass> {
 
   ProtocolConformanceRef remapConformance(Type Ty, ProtocolConformanceRef C) {
     if (Functor.SubsMap || Ty->hasLocalArchetype()) {
-      C = C.subst(Ty, Functor, Functor);
+      SubstOptions options(std::nullopt);
+      if (!Functor.LocalArchetypeSubs.empty())
+        options |= SubstFlags::SubstituteLocalArchetypes;
+
+      C = C.subst(Ty, Functor, Functor, options);
       if (asImpl().shouldSubstOpaqueArchetypes())
-        Ty = Ty.subst(Functor, Functor);
+        Ty = Ty.subst(Functor, Functor, options);
     }
 
     if (asImpl().shouldSubstOpaqueArchetypes()) {
@@ -520,8 +541,13 @@ class SILCloner : protected SILInstructionVisitor<ImplClass> {
 
   SubstitutionMap remapSubstitutionMap(SubstitutionMap Subs) {
     // If we have local archetypes to substitute, do so now.
-    if (Subs.getRecursiveProperties().hasLocalArchetype() || Functor.SubsMap)
-      Subs = Subs.subst(Functor, Functor);
+    if (Subs.getRecursiveProperties().hasLocalArchetype() || Functor.SubsMap) {
+      SubstOptions options(std::nullopt);
+      if (!Functor.LocalArchetypeSubs.empty())
+        options |= SubstFlags::SubstituteLocalArchetypes;
+
+      Subs = Subs.subst(Functor, Functor, options);
+    }
 
     if (asImpl().shouldSubstOpaqueArchetypes()) {
       auto context = getBuilder().getTypeExpansionContext();

lib/AST/ASTMangler.cpp

@@ -3800,7 +3800,8 @@ void ASTMangler::appendClosureComponents(CanType Ty, unsigned discriminator,
 
   Ty = Ty.subst(MapLocalArchetypesOutOfContext(Sig, capturedEnvs),
                 MakeAbstractConformanceForGenericType(),
-                SubstFlags::PreservePackExpansionLevel)->getCanonicalType();
+                SubstFlags::PreservePackExpansionLevel |
+                SubstFlags::SubstituteLocalArchetypes)->getCanonicalType();
 
   appendType(Ty, Sig);
   appendOperator(isImplicit ? "fu" : "fU", Index(discriminator));

lib/AST/GenericEnvironment.cpp

@@ -717,7 +717,8 @@ GenericEnvironment::mapElementTypeIntoPackContext(Type type) const {
     },
     MakeAbstractConformanceForGenericType(),
     SubstFlags::AllowLoweredTypes |
-    SubstFlags::PreservePackExpansionLevel);
+    SubstFlags::PreservePackExpansionLevel |
+    SubstFlags::SubstituteLocalArchetypes);
 
   auto shapeClass = elementEnv->getOpenedElementShapeClass();
 

lib/AST/TypeSubstitution.cpp

@@ -472,10 +472,13 @@ InFlightSubstitution::lookupConformance(CanType dependentType,
 }
 
 bool InFlightSubstitution::isInvariant(Type derivedType) const {
-  return !derivedType->hasArchetype()
-      && !derivedType->hasTypeParameter()
-      && (!shouldSubstituteOpaqueArchetypes()
-          || !derivedType->hasOpaqueArchetype());
+  if (derivedType->hasPrimaryArchetype() || derivedType->hasTypeParameter())
+    return false;
+  if (shouldSubstituteLocalArchetypes() && derivedType->hasLocalArchetype())
+    return false;
+  if (shouldSubstituteOpaqueArchetypes() && derivedType->hasOpaqueArchetype())
+    return false;
+  return true;
 }
 
 namespace {
@@ -524,6 +527,12 @@ TypeSubstituter::transform(TypeBase *type, TypePosition position) {
       && isa<OpaqueTypeArchetypeType>(substOrig))
     return std::nullopt;
 
+  // Local types can't normally be directly substituted unless we
+  // specifically were asked to substitute them.
+  if (!IFS.shouldSubstituteLocalArchetypes()
+      && isa<LocalArchetypeType>(substOrig))
+    return std::nullopt;
+
   // If we have a substitution for this type, use it.
   if (auto known = IFS.substType(substOrig, level)) {
     if (IFS.shouldSubstituteOpaqueArchetypes() &&

lib/SIL/IR/SILFunctionType.cpp

@@ -1969,7 +1969,8 @@ lowerCaptureContextParameters(TypeConverter &TC, SILDeclRef function,
 
     t = t.subst(MapLocalArchetypesOutOfContext(origGenericSig, capturedEnvs),
                 MakeAbstractConformanceForGenericType(),
-                SubstFlags::PreservePackExpansionLevel);
+                SubstFlags::PreservePackExpansionLevel |
+                SubstFlags::SubstituteLocalArchetypes);
 
     LLVM_DEBUG(llvm::dbgs() << "-- maps to " << t->getCanonicalType() << "\n");
     return t->getCanonicalType();
@@ -2958,7 +2959,8 @@ buildThunkSignature(SILFunction *fn,
     auto thunkInterfaceType = Type(localArchetype).subst(
         mapOutOfContext,
         MakeAbstractConformanceForGenericType(),
-        SubstFlags::PreservePackExpansionLevel);
+        SubstFlags::PreservePackExpansionLevel |
+        SubstFlags::SubstituteLocalArchetypes);
     auto thunkArchetype = genericEnv->mapTypeIntoContext(
         thunkInterfaceType);
     contextLocalArchetypes.insert(std::make_pair(localArchetype,
@@ -3026,25 +3028,27 @@ CanSILFunctionType swift::buildSILFunctionThunkType(
   // Does the thunk type involve a local archetype type?
   SmallVector<CanLocalArchetypeType, 8> localArchetypes;
   auto archetypeVisitor = [&](CanType t) {
-    if (auto archetypeTy = dyn_cast<ArchetypeType>(t)) {
-      if (auto local = dyn_cast<LocalArchetypeType>(archetypeTy)) {
-        auto root = local.getRoot();
-        if (llvm::find(localArchetypes, root) == localArchetypes.end())
-          localArchetypes.push_back(root);
-      }
+    if (auto local = dyn_cast<LocalArchetypeType>(t)) {
+      auto root = local.getRoot();
+      if (llvm::find(localArchetypes, root) == localArchetypes.end())
+        localArchetypes.push_back(root);
     }
   };
 
+  if (expectedType->hasLocalArchetype())
+    expectedType.visit(archetypeVisitor);
+  if (sourceType->hasLocalArchetype())
+    sourceType.visit(archetypeVisitor);
+
   // Use the generic signature from the context if the thunk involves
   // generic parameters.
   CanGenericSignature genericSig;
   SubstitutionMap contextSubs;
   llvm::DenseMap<ArchetypeType*, Type> contextLocalArchetypes;
 
-  if (expectedType->hasArchetype() || sourceType->hasArchetype()) {
-    expectedType.visit(archetypeVisitor);
-    sourceType.visit(archetypeVisitor);
-
+  if (!localArchetypes.empty() ||
+      expectedType->hasPrimaryArchetype() ||
+      sourceType->hasPrimaryArchetype()) {
     genericSig = buildThunkSignature(fn,
                                      localArchetypes,
                                      genericEnv,
@@ -3079,13 +3083,16 @@ CanSILFunctionType swift::buildSILFunctionThunkType(
   // opened existential with the new archetype.
   auto substFormalTypeIntoThunkContext =
       [&](CanType t) -> CanType {
-    return t.subst(substTypeHelper, substConformanceHelper)
+    return t.subst(substTypeHelper, substConformanceHelper,
+                   SubstFlags::SubstituteLocalArchetypes)
                ->getCanonicalType();
   };
   auto substLoweredTypeIntoThunkContext =
       [&](CanSILFunctionType t) -> CanSILFunctionType {
     return SILType::getPrimitiveObjectType(t)
-             .subst(fn->getModule(), substTypeHelper, substConformanceHelper)
+             .subst(fn->getModule(), substTypeHelper, substConformanceHelper,
+                    CanGenericSignature(),
+                    SubstFlags::SubstituteLocalArchetypes)
              .castTo<SILFunctionType>();
   };
 

lib/SIL/IR/TypeLowering.cpp

@@ -3909,7 +3909,8 @@ getAnyFunctionRefInterfaceType(TypeConverter &TC,
     auto substType = Type(funcType).subst(
         MapLocalArchetypesOutOfContext(sig.baseGenericSig, sig.capturedEnvs),
         MakeAbstractConformanceForGenericType(),
-        SubstFlags::PreservePackExpansionLevel);
+        SubstFlags::PreservePackExpansionLevel |
+        SubstFlags::SubstituteLocalArchetypes);
     funcType = cast<FunctionType>(substType->getCanonicalType());
   }
 
@@ -4991,7 +4992,8 @@ TypeConverter::getInterfaceBoxTypeForCapture(ValueDecl *captured,
       mapOutOfContext,
       MakeAbstractConformanceForGenericType(),
       SubstFlags::PreservePackExpansionLevel |
-      SubstFlags::AllowLoweredTypes)->getCanonicalType();
+      SubstFlags::AllowLoweredTypes |
+      SubstFlags::SubstituteLocalArchetypes)->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.

lib/SIL/Verifier/SILVerifier.cpp

@@ -6310,7 +6310,8 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
                                      substTypes,
                                      LookUpConformanceInModule(),
                                      CanGenericSignature(),
-                                     SubstFlags::PreservePackExpansionLevel);
+                                     SubstFlags::PreservePackExpansionLevel |
+                                     SubstFlags::SubstituteLocalArchetypes);
         requireSameType(indexedElementSILType, substTargetElementSILType,
                         "lanewise-substituted pack element type didn't "
                         "match expected element type");

lib/SILGen/ResultPlan.cpp

@@ -145,7 +145,8 @@ mapTypeOutOfOpenedExistentialContext(CanType t) {
 
         return params[index];
       },
-      MakeAbstractConformanceForGenericType());
+      MakeAbstractConformanceForGenericType(),
+      SubstFlags::SubstituteLocalArchetypes);
 
   return std::make_pair(mappedTy->getCanonicalType(), mappedSubs);
 }