AST: Add -max-substitution-count= and -max-substitution-depth= frontend flags

Author: slavapestov

include/swift/AST/InFlightSubstitution.h

@@ -35,7 +35,8 @@ class InFlightSubstitution {
   LookupConformanceFn BaselineLookupConformance;
   SubstOptions Options;
   RecursiveTypeProperties Props;
-  unsigned RemainingCount : 16;
+  unsigned RemainingCount : 15;
+  unsigned InitLimit : 1;
   unsigned RemainingDepth : 15;
   unsigned LimitReached : 1;
 
@@ -51,7 +52,7 @@ class InFlightSubstitution {
   ProtocolConformanceRef projectLaneFromPackConformance(
       PackConformance *substPackConf, unsigned level);
 
-  bool checkLimits();
+  bool checkLimits(Type ty);
 
 public:
   InFlightSubstitution(TypeSubstitutionFn substType,

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/TypeSubstitution.cpp

@@ -127,11 +127,10 @@ InFlightSubstitution::InFlightSubstitution(TypeSubstitutionFn substType,
     BaselineLookupConformance(lookupConformance),
     Options(options) {
 
+  InitLimit = 0;
   LimitReached = 0;
-
-  // FIXME: Make these configurable
-  RemainingCount = 1000;
-  RemainingDepth = 20;
+  RemainingDepth = 0;
+  RemainingCount = 0;
 
   // FIXME: Don't substitute type parameters if one of the special flags is set.
   Props |= RecursiveTypeProperties::HasTypeParameter;
@@ -240,7 +239,17 @@ Type InFlightSubstitution::projectLaneFromPackType(Type substType,
   }
 }
 
-bool InFlightSubstitution::checkLimits() {
+/// Note that the type is just used to recover an ASTContext.
+bool InFlightSubstitution::checkLimits(Type ty) {
+  if (!InitLimit) {
+    auto &ctx = ty->getASTContext();
+
+    InitLimit = 1;
+    RemainingCount = ctx.LangOpts.MaxSubstitutionCount;
+    RemainingDepth = ctx.LangOpts.MaxSubstitutionDepth;
+    return false;
+  }
+
   if (RemainingCount == 0 || RemainingDepth == 0) {
     LimitReached = true;
     return true;
@@ -263,7 +272,7 @@ Type InFlightSubstitution::substType(SubstitutableType *origType,
   if (shouldSubstituteOpaqueArchetypes() &&
       substType->hasOpaqueArchetype() &&
       !substType->isEqual(origType)) {
-    if (checkLimits()) {
+    if (checkLimits(origType)) {
       return ErrorType::get(substType);
     }
 
@@ -304,7 +313,7 @@ InFlightSubstitution::lookupConformance(Type dependentType,
   if (shouldSubstituteOpaqueArchetypes() && substConfRef &&
       substConfRef.getType()->hasOpaqueArchetype() &&
       !substConfRef.getType()->isEqual(dependentType)) {
-    if (checkLimits()) {
+    if (checkLimits(dependentType)) {
       return ProtocolConformanceRef::forInvalid();
     }
     --RemainingCount;

lib/Frontend/CompilerInvocation.cpp

@@ -1758,6 +1758,11 @@ static bool ParseLangArgs(LangOptions &Opts, ArgList &Args,
   if (Args.hasArg(OPT_enable_requirement_machine_opaque_archetypes))
     Opts.EnableRequirementMachineOpaqueArchetypes = true;
 
+  setUnsignedIntegerArgument(OPT_max_substitution_depth,
+                             Opts.MaxSubstitutionDepth);
+  setUnsignedIntegerArgument(OPT_max_substitution_count,
+                             Opts.MaxSubstitutionCount);
+
   if (Args.hasArg(OPT_enable_experimental_lifetime_dependence_inference))
     Opts.EnableExperimentalLifetimeDependenceInference = true;
   if (Args.hasArg(OPT_disable_experimental_lifetime_dependence_inference))