Sema: Improve the infinite opaque return type check

include/swift/AST/InFlightSubstitution.h

@@ -23,21 +23,36 @@
 #define SWIFT_AST_INFLIGHTSUBSTITUTION_H
 
 #include "swift/AST/SubstitutionMap.h"
+#include "llvm/ADT/DenseMap.h"
 
 namespace swift {
 class SubstitutionMap;
 
 class InFlightSubstitution {
-  SubstOptions Options;
+  friend class SubstitutionMap;
+
   TypeSubstitutionFn BaselineSubstType;
   LookupConformanceFn BaselineLookupConformance;
+  SubstOptions Options;
   RecursiveTypeProperties Props;
+  unsigned RemainingCount : 15;
+  unsigned InitLimit : 1;
+  unsigned RemainingDepth : 15;
+  unsigned LimitReached : 1;
 
   struct ActivePackExpansion {
     bool isSubstExpansion = false;
     unsigned expansionIndex = 0;
   };
-  SmallVector<ActivePackExpansion, 4> ActivePackExpansions;
+  llvm::SmallVector<ActivePackExpansion, 4> ActivePackExpansions;
+  llvm::SmallDenseMap<SubstitutionMap, SubstitutionMap, 2> SubMaps;
+
+  Type projectLaneFromPackType(
+      Type substType, unsigned level);
+  ProtocolConformanceRef projectLaneFromPackConformance(
+      PackConformance *substPackConf, unsigned level);
+
+  bool checkLimits(Type ty);
 
 public:
   InFlightSubstitution(TypeSubstitutionFn substType,
@@ -145,6 +160,10 @@ class InFlightSubstitution {
 
   /// Is the given type invariant to substitution?
   bool isInvariant(Type type) const;
+
+  bool wasLimitReached() const {
+    return LimitReached;
+  }
 };
 
 /// A helper classes that provides stable storage for the query

include/swift/AST/SubstitutionMap.h

@@ -217,7 +217,7 @@ class SubstitutionMap {
 
   /// Whether to dump the full substitution map, or just a minimal useful subset
   /// (on a single line).
-  enum class DumpStyle { Minimal, Full };
+  enum class DumpStyle { Minimal, NoConformances, Full };
   /// Dump the contents of this substitution map for debugging purposes.
   void dump(llvm::raw_ostream &out, DumpStyle style = DumpStyle::Full,
             unsigned indent = 0) const;

include/swift/AST/Types.h

@@ -7077,24 +7077,16 @@ enum class OpaqueSubstitutionKind {
 /// archetypes with underlying types visible at a given resilience expansion
 /// to their underlying types.
 class ReplaceOpaqueTypesWithUnderlyingTypes {
-public:
-  using SeenDecl = std::pair<OpaqueTypeDecl *, SubstitutionMap>;
 private:
   ResilienceExpansion contextExpansion;
   llvm::PointerIntPair<const DeclContext *, 1, bool> inContextAndIsWholeModule;
-  llvm::DenseSet<SeenDecl> *seenDecls;
 
 public:
   ReplaceOpaqueTypesWithUnderlyingTypes(const DeclContext *inContext,
                                         ResilienceExpansion contextExpansion,
                                         bool isWholeModuleContext)
       : contextExpansion(contextExpansion),
-        inContextAndIsWholeModule(inContext, isWholeModuleContext),
-        seenDecls(nullptr) {}
-
-  ReplaceOpaqueTypesWithUnderlyingTypes(
-      const DeclContext *inContext, ResilienceExpansion contextExpansion,
-      bool isWholeModuleContext, llvm::DenseSet<SeenDecl> &seen);
+        inContextAndIsWholeModule(inContext, isWholeModuleContext) {}
 
   /// TypeSubstitutionFn
   Type operator()(SubstitutableType *maybeOpaqueType) const;

include/swift/Basic/LangOptions.h

@@ -612,6 +612,14 @@ namespace swift {
     /// rewrite system.
     bool EnableRequirementMachineOpaqueArchetypes = false;
 
+    /// Maximum nesting depth for type substitution operations, to prevent
+    /// runaway recursion.
+    unsigned MaxSubstitutionDepth = 50;
+
+    /// Maximum step count for type substitution operations, to prevent
+    /// runaway recursion.
+    unsigned MaxSubstitutionCount = 2000;
+
     /// Enable implicit lifetime dependence for ~Escapable return types.
     bool EnableExperimentalLifetimeDependenceInference = false;
 

include/swift/Option/FrontendOptions.td

@@ -474,6 +474,14 @@ def disable_requirement_machine_reuse : Flag<["-"], "disable-requirement-machine
 def enable_requirement_machine_opaque_archetypes : Flag<["-"], "enable-requirement-machine-opaque-archetypes">,
   HelpText<"Enable more correct opaque archetype support, which is off by default because it might fail to produce a convergent rewrite system">;
 
+def max_substitution_depth : Joined<["-"], "max-substitution-depth=">,
+  Flags<[FrontendOption, HelpHidden, DoesNotAffectIncrementalBuild]>,
+  HelpText<"Set the maximum nesting depth for type substitution operations">;
+
+def max_substitution_count : Joined<["-"], "max-substitution-count=">,
+  Flags<[FrontendOption, HelpHidden, DoesNotAffectIncrementalBuild]>,
+  HelpText<"Set the maximum step count for type substitution operations">;
+
 def dump_type_witness_systems : Flag<["-"], "dump-type-witness-systems">,
   HelpText<"Enables dumping type witness systems from associated type inference">;
 

lib/AST/ASTDumper.cpp

@@ -1255,13 +1255,19 @@ namespace {
 
     /// Print a substitution map as a child node.
     void printRec(SubstitutionMap map, Label label) {
+      printRec(map, SubstitutionMap::DumpStyle::Full, label);
+    }
+
+    /// Print a substitution map as a child node.
+    void printRec(SubstitutionMap map, SubstitutionMap::DumpStyle style,
+                  Label label) {
       SmallPtrSet<const ProtocolConformance *, 4> Dumped;
-      printRec(map, Dumped, label);
+      printRec(map, style, Dumped, label);
     }
 
     /// Print a substitution map as a child node.
-    void printRec(SubstitutionMap map, VisitedConformances &visited,
-                  Label label);
+    void printRec(SubstitutionMap map, SubstitutionMap::DumpStyle style,
+                  VisitedConformances &visited, Label label);
 
     /// Print a substitution map as a child node.
     void printRec(const ProtocolConformanceRef &conf,
@@ -5807,7 +5813,8 @@ class PrintConformance : public PrintBase {
         if (!shouldPrintDetails)
           break;
 
-        printRec(conf->getSubstitutionMap(), visited,
+        printRec(conf->getSubstitutionMap(),
+                 SubstitutionMap::DumpStyle::Full, visited,
                  Label::optional("substitutions"));
         if (auto condReqs = conf->getConditionalRequirementsIfAvailableOrCached(/*computeIfPossible=*/false)) {
           printList(*condReqs, [&](auto subReq, Label label) {
@@ -5906,7 +5913,7 @@ class PrintConformance : public PrintBase {
 
     // A minimal dump doesn't need the details about the conformances, a lot of
     // that info can be inferred from the signature.
-    if (style == SubstitutionMap::DumpStyle::Minimal)
+    if (style != SubstitutionMap::DumpStyle::Full)
       return;
 
     auto conformances = map.getConformances();
@@ -5925,13 +5932,12 @@ class PrintConformance : public PrintBase {
   }
 };
 
-void PrintBase::printRec(SubstitutionMap map, VisitedConformances &visited,
-                         Label label) {
+void PrintBase::printRec(SubstitutionMap map, SubstitutionMap::DumpStyle style,
+                         VisitedConformances &visited, Label label) {
   printRecArbitrary(
       [&](Label label) {
         PrintConformance(Writer, MemberLoading)
-            .visitSubstitutionMap(map, SubstitutionMap::DumpStyle::Full,
-                                  visited, label);
+            .visitSubstitutionMap(map, style, visited, label);
       },
       label);
 }
@@ -6354,7 +6360,9 @@ namespace {
 
       printArchetypeCommonRec(T);
       if (!T->getSubstitutions().empty()) {
-        printRec(T->getSubstitutions(), Label::optional("substitutions"));
+        printRec(T->getSubstitutions(),
+                 SubstitutionMap::DumpStyle::NoConformances,
+                 Label::optional("substitutions"));
       }
 
       printFoot();

lib/AST/SubstitutionMap.cpp

@@ -325,6 +325,13 @@ SubstitutionMap SubstitutionMap::subst(TypeSubstitutionFn subs,
 SubstitutionMap SubstitutionMap::subst(InFlightSubstitution &IFS) const {
   if (empty()) return SubstitutionMap();
 
+  // FIXME: Get this caching working with pack expansions as well.
+  if (IFS.ActivePackExpansions.empty()) {
+    auto found = IFS.SubMaps.find(*this);
+    if (found != IFS.SubMaps.end())
+      return found->second;
+  }
+
   SmallVector<Type, 4> newSubs;
   for (Type type : getReplacementTypes()) {
     newSubs.push_back(type.subst(IFS));
@@ -345,7 +352,12 @@ SubstitutionMap SubstitutionMap::subst(InFlightSubstitution &IFS) const {
   }
 
   assert(oldConformances.empty());
-  return SubstitutionMap(genericSig, newSubs, newConformances);
+  auto result = SubstitutionMap(genericSig, newSubs, newConformances);
+
+  if (IFS.ActivePackExpansions.empty())
+    (void) IFS.SubMaps.insert(std::make_pair(*this, result));
+
+  return result;
 }
 
 SubstitutionMap
@@ -624,9 +636,23 @@ SubstitutionMap swift::substOpaqueTypesWithUnderlyingTypes(
   ReplaceOpaqueTypesWithUnderlyingTypes replacer(
       context.getContext(), context.getResilienceExpansion(),
       context.isWholeModuleContext());
-  return subs.subst(replacer, replacer,
-                    SubstFlags::SubstituteOpaqueArchetypes |
-                    SubstFlags::PreservePackExpansionLevel);
+  InFlightSubstitution IFS(replacer, replacer,
+                           SubstFlags::SubstituteOpaqueArchetypes |
+                           SubstFlags::PreservePackExpansionLevel);
+
+  auto substSubs = subs.subst(IFS);
+
+  if (IFS.wasLimitReached()) {
+    ABORT([&](auto &out) {
+      out << "Possible non-terminating type substitution detected\n\n";
+      out << "Original substitution map:\n";
+      subs.dump(out, SubstitutionMap::DumpStyle::NoConformances);
+      out << "Substituted substitution map:\n";
+      substSubs.dump(out, SubstitutionMap::DumpStyle::NoConformances);
+    });
+  }
+
+  return substSubs;
 }
 
 Type OuterSubstitutions::operator()(SubstitutableType *type) const {