[CS] Limit the number of chained @dynamicMemberLookup lookups

Set an upper bound on the number of chained lookups we attempt to
avoid spinning while trying to recursively apply the same dynamic
member lookup to itself.

rdar://157288911

include/swift/AST/DiagnosticsSema.def

@@ -1735,6 +1735,11 @@ ERROR(dynamic_member_lookup_candidate_inaccessible,none,
       "enclosing type",
       (ValueDecl *))
 
+ERROR(too_many_dynamic_member_lookups,none,
+      "could not find member %0; exceeded the maximum number of nested "
+      "dynamic member lookups",
+      (DeclNameRef))
+
 ERROR(string_index_not_integer,none,
       "String must not be indexed with %0, it has variable size elements",
       (Type))

include/swift/Basic/LangOptions.h

@@ -952,6 +952,10 @@ namespace swift {
     /// (It's arbitrary, but will keep the compiler from taking too much time.)
     unsigned SwitchCheckingInvocationThreshold = 200000;
 
+    /// The maximum number of `@dynamicMemberLookup`s that can be chained to
+    /// resolve a member reference.
+    unsigned DynamicMemberLookupDepthLimit = 100;
+
     /// If true, the time it takes to type-check each function will be dumped
     /// to llvm::errs().
     bool DebugTimeFunctionBodies = false;

include/swift/Option/FrontendOptions.td

@@ -895,6 +895,12 @@ def disable_invalid_ephemeralness_as_error :
 def switch_checking_invocation_threshold_EQ : Joined<["-"],
     "switch-checking-invocation-threshold=">;
 
+def dynamic_member_lookup_depth_limit_EQ
+    : Joined<["-"], "dynamic-member-lookup-depth-limit=">,
+      HelpText<
+          "The maximum number of dynamic member lookups that can be chained "
+          "to resolve a member reference">;
+
 def enable_new_operator_lookup :
   Flag<["-"], "enable-new-operator-lookup">,
   HelpText<"Enable the new operator decl and precedencegroup lookup behavior">;

include/swift/Sema/CSFix.h

@@ -488,6 +488,9 @@ enum class FixKind : uint8_t {
   /// the type it's attempting to bind to.
   AllowInlineArrayLiteralCountMismatch,
 
+  /// Reached the limit of @dynamicMemberLookup depth.
+  TooManyDynamicMemberLookups,
+
   /// Ignore that a conformance is isolated but is not allowed to be.
   IgnoreIsolatedConformance,
 };
@@ -3881,6 +3884,33 @@ public:
   }
 };
 
+class TooManyDynamicMemberLookups : public ConstraintFix {
+  DeclNameRef Name;
+
+  TooManyDynamicMemberLookups(ConstraintSystem &cs, DeclNameRef name,
+                              ConstraintLocator *locator)
+      : ConstraintFix(cs, FixKind::TooManyDynamicMemberLookups, locator),
+        Name(name) {}
+
+public:
+  std::string getName() const override {
+    return "too many dynamic member lookups";
+  }
+
+  bool diagnose(const Solution &solution, bool asNote = false) const override;
+
+  bool diagnoseForAmbiguity(CommonFixesArray commonFixes) const override {
+    return diagnose(*commonFixes.front().first);
+  }
+
+  static TooManyDynamicMemberLookups *
+  create(ConstraintSystem &cs, DeclNameRef name, ConstraintLocator *locator);
+
+  static bool classof(const ConstraintFix *fix) {
+    return fix->getKind() == FixKind::TooManyDynamicMemberLookups;
+  }
+};
+
 class IgnoreIsolatedConformance : public ConstraintFix {
   ProtocolConformance *conformance;
 

lib/Frontend/CompilerInvocation.cpp

@@ -1897,6 +1897,8 @@ static bool ParseTypeCheckerArgs(TypeCheckerOptions &Opts, ArgList &Args,
                              Opts.WarnLongExpressionTypeChecking);
   setUnsignedIntegerArgument(OPT_solver_expression_time_threshold_EQ,
                              Opts.ExpressionTimeoutThreshold);
+  setUnsignedIntegerArgument(OPT_dynamic_member_lookup_depth_limit_EQ,
+                             Opts.DynamicMemberLookupDepthLimit);
   setUnsignedIntegerArgument(OPT_switch_checking_invocation_threshold_EQ,
                              Opts.SwitchCheckingInvocationThreshold);
   setUnsignedIntegerArgument(OPT_debug_constraints_attempt,

lib/Sema/CSDiagnostics.cpp

@@ -9628,6 +9628,12 @@ bool IncorrectInlineArrayLiteralCount::diagnoseAsError() {
   return true;
 }
 
+bool TooManyDynamicMemberLookupsFailure::diagnoseAsError() {
+  emitDiagnostic(diag::too_many_dynamic_member_lookups, Name)
+      .highlight(getSourceRange());
+  return true;
+}
+
 bool DisallowedIsolatedConformance::diagnoseAsError() {
   emitDiagnostic(diag::isolated_conformance_with_sendable_simple,
                  conformance->getType(),

lib/Sema/CSDiagnostics.h

@@ -3310,6 +3310,17 @@ public:
   bool diagnoseAsError() override;
 };
 
+class TooManyDynamicMemberLookupsFailure final : public FailureDiagnostic {
+  DeclNameRef Name;
+
+public:
+  TooManyDynamicMemberLookupsFailure(const Solution &solution, DeclNameRef name,
+                                     ConstraintLocator *locator)
+      : FailureDiagnostic(solution, locator), Name(name) {}
+
+  bool diagnoseAsError() override;
+};
+
 /// Diagnose when an isolated conformance is used in a place where one cannot
 /// be, e.g., due to a Sendable or SendableMetatype requirement on the
 /// corresponding type parameter.

lib/Sema/CSFix.cpp

@@ -2796,6 +2796,18 @@ bool AllowInlineArrayLiteralCountMismatch::diagnose(const Solution &solution,
   return failure.diagnose(asNote);
 }
 
+TooManyDynamicMemberLookups *
+TooManyDynamicMemberLookups::create(ConstraintSystem &cs, DeclNameRef name,
+                                    ConstraintLocator *locator) {
+  return new (cs.getAllocator()) TooManyDynamicMemberLookups(cs, name, locator);
+}
+
+bool TooManyDynamicMemberLookups::diagnose(const Solution &solution,
+                                           bool asNote) const {
+  TooManyDynamicMemberLookupsFailure failure(solution, Name, getLocator());
+  return failure.diagnose(asNote);
+}
+
 IgnoreIsolatedConformance *
 IgnoreIsolatedConformance::create(ConstraintSystem &cs,
                                   ConstraintLocator *locator,

lib/Sema/CSSimplify.cpp

@@ -16034,6 +16034,7 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyFixConstraint(
   case FixKind::IgnoreOutOfPlaceThenStmt:
   case FixKind::IgnoreMissingEachKeyword:
   case FixKind::AllowInlineArrayLiteralCountMismatch:
+  case FixKind::TooManyDynamicMemberLookups:
   case FixKind::IgnoreIsolatedConformance:
     llvm_unreachable("handled elsewhere");
   }

lib/Sema/TypeOfReference.cpp

@@ -2191,11 +2191,25 @@ void ConstraintSystem::bindOverloadType(const SelectedOverload &overload,
                            // FIXME: Should propagate name-as-written through.
                            : DeclNameRef(choice.getName());
 
-    addValueMemberConstraint(
-        LValueType::get(rootTy), memberName, memberTy, useDC,
-        isSubscriptRef ? FunctionRefInfo::doubleBaseNameApply()
-                       : FunctionRefInfo::unappliedBaseName(),
-        /*outerAlternatives=*/{}, keyPathLoc);
+    // Check the current depth of applied dynamic member lookups, if we've
+    // exceeded the limit then record a fix and set a hole for the member.
+    unsigned lookupDepth = [&]() {
+      auto path = keyPathLoc->getPath();
+      auto iter = path.begin();
+      (void)keyPathLoc->findFirst<LocatorPathElt::KeyPathDynamicMember>(iter);
+      return path.end() - iter;
+    }();
+    if (lookupDepth > ctx.TypeCheckerOpts.DynamicMemberLookupDepthLimit) {
+      (void)recordFix(TooManyDynamicMemberLookups::create(
+          *this, DeclNameRef(choice.getName()), locator));
+      recordTypeVariablesAsHoles(memberTy);
+    } else {
+      addValueMemberConstraint(
+          LValueType::get(rootTy), memberName, memberTy, useDC,
+          isSubscriptRef ? FunctionRefInfo::doubleBaseNameApply()
+                         : FunctionRefInfo::unappliedBaseName(),
+          /*outerAlternatives=*/{}, keyPathLoc);
+    }
 
     // In case of subscript things are more complicated comparing to "dot"
     // syntax, because we have to get "applicable function" constraint

test/Frontend/dynamic_member_lookup_limit.swift

@@ -0,0 +1,21 @@
+// RUN: %target-typecheck-verify-swift -dynamic-member-lookup-depth-limit=2
+
+@dynamicMemberLookup
+struct Lens<T> {
+  init() {}
+  subscript<U>(dynamicMember kp: KeyPath<T, U>) -> U {
+    fatalError()
+  }
+}
+
+struct S {
+  var x: Int
+}
+
+_ = \Lens<S>.x // Fine
+_ = \Lens<Lens<S>>.x // Also fine
+_ = \Lens<Lens<Lens<S>>>.x // expected-error {{could not find member 'x'; exceeded the maximum number of nested dynamic member lookups}}
+
+_ = Lens<S>().x // Fine
+_ = Lens<Lens<S>>().x // Also fine
+_ = Lens<Lens<Lens<S>>>().x // expected-error {{could not find member 'x'; exceeded the maximum number of nested dynamic member lookups}}

test/attr/attr_dynamic_member_lookup.swift

@@ -978,3 +978,33 @@ struct WithSendable {
     get { false }
   }
 }
+
+// Make sure we enforce a limit on the number of chained dynamic member lookups.
+@dynamicMemberLookup
+struct SelfRecursiveLookup<T> {
+  init(_: () -> T) {}
+  init(_: () -> KeyPath<Self, T>) {}
+  subscript<U>(dynamicMember kp: KeyPath<T, U>) -> SelfRecursiveLookup<U> {}
+}
+let selfRecurse1 = SelfRecursiveLookup { selfRecurse1.e }
+// expected-error@-1 {{could not find member 'e'; exceeded the maximum number of nested dynamic member lookups}}
+
+let selfRecurse2 = SelfRecursiveLookup { selfRecurse2[a: 0] }
+// expected-error@-1 {{could not find member 'subscript'; exceeded the maximum number of nested dynamic member lookups}}
+
+let selfRecurse3 = SelfRecursiveLookup { \.e }
+// expected-error@-1 {{could not find member 'e'; exceeded the maximum number of nested dynamic member lookups}}
+
+let selfRecurse4 = SelfRecursiveLookup { \.[a: 0] }
+// expected-error@-1 {{could not find member 'subscript'; exceeded the maximum number of nested dynamic member lookups}}
+
+extension SelfRecursiveLookup where T == SelfRecursiveLookup<SelfRecursiveLookup<SelfRecursiveLookup<Int>>> {
+  var terminator: T { fatalError() }
+  subscript(terminator terminator: Int) -> T { fatalError() }
+}
+
+let selfRecurse5 = SelfRecursiveLookup { selfRecurse5.terminator }
+let selfRecurse6 = SelfRecursiveLookup { selfRecurse6[terminator: 0] }
+
+let selfRecurse7 = SelfRecursiveLookup { \.terminator }
+let selfRecurse8 = SelfRecursiveLookup { \.[terminator: 0] }

validation-test/IDE/crashers/fd52bd37cd5c96.swift

@@ -1,13 +0,0 @@
-// {"kind":"complete","original":"ea806f48","signature":"swift::Type::transformRec(llvm::function_ref<std::__1::optional<swift::Type> (swift::TypeBase*)>) const"}
-// The issue here is that the solver attempts to recursively apply the same
-// dynamic member lookup until eventually it overflows the stack. Make sure
-// we either timeout or crash.
-// RUN: not %{python} %S/../../../test/Inputs/timeout.py 60 \
-// RUN:   %target-swift-ide-test -code-completion -batch-code-completion -skip-filecheck -code-completion-diagnostics -source-filename %s || \
-// RUN:   not --crash %target-swift-ide-test -code-completion -batch-code-completion -skip-filecheck -code-completion-diagnostics -source-filename %s
-@dynamicMemberLookup
-struct a<b{
-c: () -> b^subscript<d>(dynamicMember e: WritableKeyPath<b, d>) a<d> }
-let binding = a
-{ buffer #^^#??
-binding.0

validation-test/IDE/crashers_fixed/fd52bd37cd5c96.swift

@@ -0,0 +1,8 @@
+// {"kind":"complete","original":"ea806f48","signature":"swift::Type::transformRec(llvm::function_ref<std::__1::optional<swift::Type> (swift::TypeBase*)>) const"}
+// RUN: %target-swift-ide-test -code-completion -batch-code-completion -skip-filecheck -code-completion-diagnostics -source-filename %s
+@dynamicMemberLookup
+struct a<b{
+c: () -> b^subscript<d>(dynamicMember e: WritableKeyPath<b, d>) a<d> }
+let binding = a
+{ buffer #^^#??
+binding.0

validation-test/compiler_crashers_2_fixed/a5fbf4336df96867.swift

@@ -1,9 +1,5 @@
 // {"kind":"typecheck","signature":"swift::TypeTransform<swift::Type::transformRec(llvm::function_ref<std::__1::optional<swift::Type> (swift::TypeBase*)>) const::Transform>::doIt(swift::Type, swift::TypePosition)"}
-// The issue here is that the solver attempts to recursively apply the same
-// dynamic member lookup until eventually it overflows the stack. Make sure
-// we either timeout or crash.
-// RUN: not %{python} %S/../../test/Inputs/timeout.py 60 %target-swift-frontend -typecheck %s || \
-// RUN: not --crash %target-swift-frontend -typecheck %s
+// RUN: not %target-swift-frontend -typecheck %s
 @dynamicMemberLookup struct S<T> {
   init(_: () -> T) {}
   subscript<U>(dynamicMember d: WritableKeyPath<T, U>) -> S<U> {}