Merge pull request #40971 from CodaFi/hitting-the-stacks

Add a Recursion Breaker to Opaque Type Substitution

Author: CodaFi

include/swift/AST/Types.h

@@ -5608,15 +5608,21 @@ enum class OpaqueSubstitutionKind {
 /// archetypes with underlying types visible at a given resilience expansion
 /// to their underlying types.
 class ReplaceOpaqueTypesWithUnderlyingTypes {
-public:
-  const DeclContext *inContext;
   ResilienceExpansion contextExpansion;
-  bool isContextWholeModule;
+  llvm::PointerIntPair<const DeclContext *, 1, bool> inContextAndIsWholeModule;
+  llvm::SmallPtrSetImpl<OpaqueTypeDecl *> *seenDecls;
+
+public:
   ReplaceOpaqueTypesWithUnderlyingTypes(const DeclContext *inContext,
                                         ResilienceExpansion contextExpansion,
                                         bool isWholeModuleContext)
-      : inContext(inContext), contextExpansion(contextExpansion),
-        isContextWholeModule(isWholeModuleContext) {}
+      : contextExpansion(contextExpansion),
+        inContextAndIsWholeModule(inContext, isWholeModuleContext),
+        seenDecls(nullptr) {}
+
+  ReplaceOpaqueTypesWithUnderlyingTypes(
+      const DeclContext *inContext, ResilienceExpansion contextExpansion,
+      bool isWholeModuleContext, llvm::SmallPtrSetImpl<OpaqueTypeDecl *> &seen);
 
   /// TypeSubstitutionFn
   Type operator()(SubstitutableType *maybeOpaqueType) const;
@@ -5632,6 +5638,13 @@ class ReplaceOpaqueTypesWithUnderlyingTypes {
   static OpaqueSubstitutionKind
   shouldPerformSubstitution(OpaqueTypeDecl *opaque, ModuleDecl *contextModule,
                             ResilienceExpansion contextExpansion);
+
+private:
+  const DeclContext *getContext() const {
+    return inContextAndIsWholeModule.getPointer();
+  }
+
+  bool isWholeModule() const { return inContextAndIsWholeModule.getInt(); }
 };
 
 /// An archetype that represents the dynamic type of an opened existential.

lib/AST/Type.cpp

@@ -3381,6 +3381,7 @@ getArchetypeAndRootOpaqueArchetype(Type maybeOpaqueType) {
 OpaqueSubstitutionKind
 ReplaceOpaqueTypesWithUnderlyingTypes::shouldPerformSubstitution(
     OpaqueTypeDecl *opaque) const {
+  const auto *inContext = getContext();
   auto inModule = inContext ? inContext->getParentModule()
                             : opaque->getParentModule();
   return shouldPerformSubstitution(opaque, inModule, contextExpansion);
@@ -3424,12 +3425,11 @@ ReplaceOpaqueTypesWithUnderlyingTypes::shouldPerformSubstitution(
   return OpaqueSubstitutionKind::SubstituteNonResilientModule;
 }
 
-static Type
-substOpaqueTypesWithUnderlyingTypes(Type ty, const DeclContext *inContext,
-                                    ResilienceExpansion contextExpansion,
-                                    bool isWholeModuleContext) {
+static Type substOpaqueTypesWithUnderlyingTypesRec(
+    Type ty, const DeclContext *inContext, ResilienceExpansion contextExpansion,
+    bool isWholeModuleContext, SmallPtrSetImpl<OpaqueTypeDecl *> &decls) {
   ReplaceOpaqueTypesWithUnderlyingTypes replacer(inContext, contextExpansion,
-                                                 isWholeModuleContext);
+                                                 isWholeModuleContext, decls);
   return ty.subst(replacer, replacer, SubstFlags::SubstituteOpaqueArchetypes);
 }
 
@@ -3489,6 +3489,13 @@ static bool canSubstituteTypeInto(Type ty, const DeclContext *dc,
   llvm_unreachable("invalid subsitution kind");
 }
 
+ReplaceOpaqueTypesWithUnderlyingTypes::ReplaceOpaqueTypesWithUnderlyingTypes(
+    const DeclContext *inContext, ResilienceExpansion contextExpansion,
+    bool isWholeModuleContext, llvm::SmallPtrSetImpl<OpaqueTypeDecl *> &seen)
+    : contextExpansion(contextExpansion),
+      inContextAndIsWholeModule(inContext, isWholeModuleContext),
+      seenDecls(&seen) {}
+
 Type ReplaceOpaqueTypesWithUnderlyingTypes::
 operator()(SubstitutableType *maybeOpaqueType) const {
   auto archetypeAndRoot = getArchetypeAndRootOpaqueArchetype(maybeOpaqueType);
@@ -3516,8 +3523,8 @@ operator()(SubstitutableType *maybeOpaqueType) const {
 
   // Check that we are allowed to substitute the underlying type into the
   // context.
-  auto inContext = this->inContext;
-  auto isContextWholeModule = this->isContextWholeModule;
+  auto inContext = this->getContext();
+  auto isContextWholeModule = this->isWholeModule();
   if (inContext &&
       partialSubstTy.findIf(
           [inContext, substitutionKind, isContextWholeModule](Type t) -> bool {
@@ -3536,18 +3543,39 @@ operator()(SubstitutableType *maybeOpaqueType) const {
 
   // If the type changed, but still contains opaque types, recur.
   if (!substTy->isEqual(maybeOpaqueType) && substTy->hasOpaqueArchetype()) {
-    return ::substOpaqueTypesWithUnderlyingTypes(
-        substTy, inContext, contextExpansion, isContextWholeModule);
+    if (auto *alreadySeen = this->seenDecls) {
+      // Detect substitution loops. If we find one, just bounce the original
+      // type back to the caller. This substitution will fail at runtime
+      // instead.
+      if (!alreadySeen->insert(opaqueRoot->getDecl()).second) {
+        return maybeOpaqueType;
+      }
+
+      auto res = ::substOpaqueTypesWithUnderlyingTypesRec(
+          substTy, inContext, contextExpansion, isContextWholeModule,
+          *alreadySeen);
+      alreadySeen->erase(opaqueRoot->getDecl());
+      return res;
+    } else {
+      // We're the top of the stack for the recursion check. Allocate a set of
+      // opaque result type decls we've already seen for the rest of the check.
+      SmallPtrSet<OpaqueTypeDecl *, 8> seenDecls;
+      seenDecls.insert(opaqueRoot->getDecl());
+      return ::substOpaqueTypesWithUnderlyingTypesRec(
+          substTy, inContext, contextExpansion, isContextWholeModule,
+          seenDecls);
+    }
   }
 
   return substTy;
 }
 
-static ProtocolConformanceRef substOpaqueTypesWithUnderlyingTypes(
+static ProtocolConformanceRef substOpaqueTypesWithUnderlyingTypesRec(
     ProtocolConformanceRef ref, Type origType, const DeclContext *inContext,
-    ResilienceExpansion contextExpansion, bool isWholeModuleContext) {
+    ResilienceExpansion contextExpansion, bool isWholeModuleContext,
+    SmallPtrSetImpl<OpaqueTypeDecl *> &decls) {
   ReplaceOpaqueTypesWithUnderlyingTypes replacer(inContext, contextExpansion,
-                                                 isWholeModuleContext);
+                                                 isWholeModuleContext, decls);
   return ref.subst(origType, replacer, replacer,
                    SubstFlags::SubstituteOpaqueArchetypes);
 }
@@ -3575,6 +3603,7 @@ operator()(CanType maybeOpaqueType, Type replacementType,
     // SIL type lowering may have already substituted away the opaque type, in
     // which case we'll end up "substituting" the same type.
     if (maybeOpaqueType->isEqual(replacementType)) {
+      const auto *inContext = getContext();
       assert(inContext && "Need context for already-substituted opaque types");
       return inContext->getParentModule()
                       ->lookupConformance(replacementType, protocol);
@@ -3604,8 +3633,8 @@ operator()(CanType maybeOpaqueType, Type replacementType,
 
   // Check that we are allowed to substitute the underlying type into the
   // context.
-  auto inContext = this->inContext;
-  auto isContextWholeModule = this->isContextWholeModule;
+  auto inContext = this->getContext();
+  auto isContextWholeModule = this->isWholeModule();
   if (partialSubstTy.findIf(
           [inContext, substitutionKind, isContextWholeModule](Type t) -> bool {
             if (!canSubstituteTypeInto(t, inContext, substitutionKind,
@@ -3628,8 +3657,28 @@ operator()(CanType maybeOpaqueType, Type replacementType,
 
   // If the type still contains opaque types, recur.
   if (substTy->hasOpaqueArchetype()) {
-    return ::substOpaqueTypesWithUnderlyingTypes(
-        substRef, substTy, inContext, contextExpansion, isContextWholeModule);
+    if (auto *alreadySeen = this->seenDecls) {
+      // Detect substitution loops. If we find one, just bounce the original
+      // type back to the caller. This substitution will fail at runtime
+      // instead.
+      if (!alreadySeen->insert(opaqueRoot->getDecl()).second) {
+        return abstractRef;
+      }
+
+      auto res = ::substOpaqueTypesWithUnderlyingTypesRec(
+          substRef, substTy, inContext, contextExpansion, isContextWholeModule,
+          *alreadySeen);
+      alreadySeen->erase(opaqueRoot->getDecl());
+      return res;
+    } else {
+      // We're the top of the stack for the recursion check. Allocate a set of
+      // opaque result type decls we've already seen for the rest of the check.
+      SmallPtrSet<OpaqueTypeDecl *, 8> seenDecls;
+      seenDecls.insert(opaqueRoot->getDecl());
+      return ::substOpaqueTypesWithUnderlyingTypesRec(
+          substRef, substTy, inContext, contextExpansion, isContextWholeModule,
+          seenDecls);
+    }
   }
   return substRef;
 }

validation-test/compiler_crashers_2_fixed/rdar87121502.swift

@@ -0,0 +1,16 @@
+// RUN: %target-swift-frontend -emit-ir -target %target-cpu-apple-macosx10.15 %s
+// REQUIRES: OS=macosx
+
+protocol P {}
+
+func one() -> some P {
+    return two()
+}
+
+func two() -> some P {
+    return three()
+}
+
+func three() -> some P {
+    return one()
+}