Prevent infinite recursion with conformance enumeration.

Author: DougGregor

lib/Sema/TypeCheckProtocol.cpp

@@ -7163,17 +7163,22 @@ void TypeChecker::inferDefaultWitnesses(ProtocolDecl *proto) {
 
 bool swift::forEachConformance(
     SubstitutionMap subs,
-    llvm::function_ref<bool(ProtocolConformanceRef)> body) {
+    llvm::function_ref<bool(ProtocolConformanceRef)> body,
+    VisitedConformances *visitedConformances) {
   if (!subs)
     return false;
 
+  VisitedConformances visited;
+  if (!visitedConformances)
+    visitedConformances = &visited;
+
   for (auto type: subs.getReplacementTypes()) {
-    if (forEachConformance(type, body))
+    if (forEachConformance(type, body, visitedConformances))
       return true;
   }
 
   for (auto conformance: subs.getConformances()) {
-    if (forEachConformance(conformance, body))
+    if (forEachConformance(conformance, body, visitedConformances))
       return true;
   }
 
@@ -7182,46 +7187,70 @@ bool swift::forEachConformance(
 
 bool swift::forEachConformance(
     ProtocolConformanceRef conformance,
-    llvm::function_ref<bool(ProtocolConformanceRef)> body) {
-  // Visit this conformance.
-  if (body(conformance))
-    return true;
+    llvm::function_ref<bool(ProtocolConformanceRef)> body,
+    VisitedConformances *visitedConformances) {
+  // Make sure we can store visited conformances.
+  VisitedConformances visited;
+  if (!visitedConformances)
+    visitedConformances = &visited;
 
   if (conformance.isInvalid() || conformance.isAbstract())
     return false;
 
   if (conformance.isPack()) {
     auto pack = conformance.getPack()->getPatternConformances();
     for (auto conformance : pack) {
-      if (forEachConformance(conformance, body))
+      if (forEachConformance(conformance, body, visitedConformances))
         return true;
     }
 
     return false;
   }
 
-  // Check the substitution make within this conformance.
+  // Extract the concrete conformance.
   auto concrete = conformance.getConcrete();
-  if (forEachConformance(concrete->getSubstitutionMap(), body))
+
+  // Prevent recursion.
+  if (!visitedConformances->insert(concrete).second)
+    return false;
+
+  // Visit this conformance.
+  if (body(conformance))
     return true;
 
+  // Check the substitution map within this conformance.
+  if (forEachConformance(concrete->getSubstitutionMap(), body,
+                         visitedConformances))
+    return true;
 
   return false;
 }
 
 bool swift::forEachConformance(
-    Type type, llvm::function_ref<bool(ProtocolConformanceRef)> body) {
+    Type type, llvm::function_ref<bool(ProtocolConformanceRef)> body,
+    VisitedConformances *visitedConformances) {
+  // Make sure we can store visited conformances.
+  VisitedConformances visited;
+  if (!visitedConformances)
+    visitedConformances = &visited;
+
+  // Prevent recursion.
+  if (!visitedConformances->insert(type.getPointer()).second)
+    return false;
+
   return type.findIf([&](Type type) {
     if (auto typeAlias = dyn_cast<TypeAliasType>(type.getPointer())) {
-      if (forEachConformance(typeAlias->getSubstitutionMap(), body))
+      if (forEachConformance(typeAlias->getSubstitutionMap(), body,
+                             visitedConformances))
         return true;
 
       return false;
     }
 
     if (auto opaqueArchetype =
             dyn_cast<OpaqueTypeArchetypeType>(type.getPointer())) {
-      if (forEachConformance(opaqueArchetype->getSubstitutions(), body))
+      if (forEachConformance(opaqueArchetype->getSubstitutions(), body,
+                             visitedConformances))
         return true;
 
       return false;
@@ -7234,7 +7263,7 @@ bool swift::forEachConformance(
 
     if (auto boundGeneric = dyn_cast<BoundGenericType>(type.getPointer())) {
       auto subs = boundGeneric->getContextSubstitutionMap();
-      if (forEachConformance(subs, body))
+      if (forEachConformance(subs, body, visitedConformances))
         return true;
 
       return false;
@@ -7246,19 +7275,25 @@ bool swift::forEachConformance(
 
 bool swift::forEachConformance(
     ConcreteDeclRef declRef,
-    llvm::function_ref<bool(ProtocolConformanceRef)> body) {
+    llvm::function_ref<bool(ProtocolConformanceRef)> body,
+    VisitedConformances *visitedConformances) {
   if (!declRef)
     return false;
 
+  // Make sure we can store visited conformances.
+  VisitedConformances visited;
+  if (!visitedConformances)
+    visitedConformances = &visited;
+
   Type type = declRef.getDecl()->getInterfaceType();
   if (auto subs = declRef.getSubstitutions()) {
-    if (forEachConformance(subs, body))
+    if (forEachConformance(subs, body, visitedConformances))
       return true;
 
     type = type.subst(subs);
   }
 
-  if (forEachConformance(type, body))
+  if (forEachConformance(type, body, visitedConformances))
     return true;
 
   return false;

lib/Sema/TypeCheckProtocol.h

@@ -240,12 +240,15 @@ bool witnessHasImplementsAttrForRequiredName(ValueDecl *witness,
 bool witnessHasImplementsAttrForExactRequirement(ValueDecl *witness,
                                                  ValueDecl *requirement);
 
+using VisitedConformances = llvm::SmallPtrSet<void *, 16>;
+
 /// Visit each conformance within the given type.
 ///
 /// If `body` returns true for any conformance, this function stops the
 /// traversal and returns true.
 bool forEachConformance(
-    Type type, llvm::function_ref<bool(ProtocolConformanceRef)> body);
+    Type type, llvm::function_ref<bool(ProtocolConformanceRef)> body,
+    VisitedConformances *visitedConformances = nullptr);
 
 /// Visit each conformance within the given conformance (including the given
 /// one).
@@ -254,23 +257,26 @@ bool forEachConformance(
 /// traversal and returns true.
 bool forEachConformance(
     ProtocolConformanceRef conformance,
-    llvm::function_ref<bool(ProtocolConformanceRef)> body);
+    llvm::function_ref<bool(ProtocolConformanceRef)> body,
+    VisitedConformances *visitedConformances = nullptr);
 
 /// Visit each conformance within the given substitution map.
 ///
 /// If `body` returns true for any conformance, this function stops the
 /// traversal and returns true.
 bool forEachConformance(
     SubstitutionMap subs,
-    llvm::function_ref<bool(ProtocolConformanceRef)> body);
+    llvm::function_ref<bool(ProtocolConformanceRef)> body,
+    VisitedConformances *visitedConformances = nullptr);
 
 /// Visit each conformance within the given declaration reference.
 ///
 /// If `body` returns true for any conformance, this function stops the
 /// traversal and returns true.
 bool forEachConformance(
     ConcreteDeclRef declRef,
-    llvm::function_ref<bool(ProtocolConformanceRef)> body);
+    llvm::function_ref<bool(ProtocolConformanceRef)> body,
+    VisitedConformances *visitedConformances = nullptr);
 
 }