Merge pull request #82321 from slavapestov/rqm-fixes

slavapestov · web-flow · commit 5987bbf966d7 · 2025-06-18T23:00:38.000-04:00
RequirementMachine: Add more limits to catch runaway computation, and fix a bug
diff --git a/include/swift/AST/DiagnosticsSema.def b/include/swift/AST/DiagnosticsSema.def
@@ -3441,7 +3441,7 @@ WARNING(associated_type_override_typealias,none,
 
 ERROR(requirement_machine_completion_failed,none,
       "cannot build rewrite system for %select{generic signature|protocol}0; "
-      "%select{%error|rule count|rule length|concrete nesting}1 limit exceeded",
+      "%select{%error|rule count|rule length|concrete type nesting|concrete type size|concrete type difference}1 limit exceeded",
       (unsigned, unsigned))
 NOTE(requirement_machine_completion_rule,none,
      "failed rewrite rule is %0", (StringRef))
diff --git a/include/swift/Basic/LangOptions.h b/include/swift/Basic/LangOptions.h
@@ -580,10 +580,18 @@ namespace swift {
     /// algorithm.
     unsigned RequirementMachineMaxRuleLength = 12;
 
-    /// Maximum concrete type nesting depth for requirement machine property map
-    /// algorithm.
+    /// Maximum concrete type nesting depth (when type is viewed as a tree) for
+    /// requirement machine property map algorithm.
     unsigned RequirementMachineMaxConcreteNesting = 30;
 
+    /// Maximum concrete type size (total number of nodes in the type tree) for
+    /// requirement machine property map algorithm.
+    unsigned RequirementMachineMaxConcreteSize = 4000;
+
+    /// Maximum number of "type difference" structures for the requirement machine
+    /// property map algorithm.
+    unsigned RequirementMachineMaxTypeDifferences = 4000;
+
     /// Maximum number of attempts to make when splitting concrete equivalence
     /// classes.
     unsigned RequirementMachineMaxSplitConcreteEquivClassAttempts = 2;
diff --git a/include/swift/Option/FrontendOptions.td b/include/swift/Option/FrontendOptions.td
@@ -447,6 +447,14 @@ def requirement_machine_max_concrete_nesting : Joined<["-"], "requirement-machin
   Flags<[FrontendOption, HelpHidden, DoesNotAffectIncrementalBuild]>,
   HelpText<"Set the maximum concrete type nesting depth before giving up">;
 
+def requirement_machine_max_concrete_size : Joined<["-"], "requirement-machine-max-concrete-size=">,
+  Flags<[FrontendOption, HelpHidden, DoesNotAffectIncrementalBuild]>,
+  HelpText<"Set the maximum concrete type total size before giving up">;
+
+def requirement_machine_max_type_differences : Joined<["-"], "requirement-machine-max-type-differences=">,
+  Flags<[FrontendOption, HelpHidden, DoesNotAffectIncrementalBuild]>,
+  HelpText<"Set the maximum number of type difference structures allocated before giving up">;
+
 def requirement_machine_max_split_concrete_equiv_class_attempts : Joined<["-"], "requirement-machine-max-split-concrete-equiv-class-attempts=">,
   Flags<[FrontendOption, HelpHidden, DoesNotAffectIncrementalBuild]>,
   HelpText<"Set the maximum concrete number of attempts at splitting "
diff --git a/lib/AST/GenericSignature.cpp b/lib/AST/GenericSignature.cpp
@@ -1436,7 +1436,8 @@ RequirementSignature RequirementSignature::getPlaceholderRequirementSignature(
                        });
 
   return RequirementSignature(ctx.AllocateCopy(requirements),
-                              ArrayRef<ProtocolTypeAlias>());
+                              ArrayRef<ProtocolTypeAlias>(),
+                              errors);
 }
 
 void RequirementSignature::getRequirementsWithInverses(
diff --git a/lib/AST/RequirementMachine/HomotopyReduction.cpp b/lib/AST/RequirementMachine/HomotopyReduction.cpp
@@ -291,8 +291,8 @@ RewriteSystem::findRuleToDelete(EliminationPredicate isRedundantRuleFn) {
     {
       // If both are concrete type requirements, prefer to eliminate the
       // one with the more deeply nested type.
-      unsigned ruleNesting = rule.getNesting();
-      unsigned otherRuleNesting = otherRule.getNesting();
+      unsigned ruleNesting = rule.getNestingAndSize().first;
+      unsigned otherRuleNesting = otherRule.getNestingAndSize().first;
 
       if (ruleNesting != otherRuleNesting) {
         if (ruleNesting > otherRuleNesting)
diff --git a/lib/AST/RequirementMachine/KnuthBendix.cpp b/lib/AST/RequirementMachine/KnuthBendix.cpp
@@ -399,7 +399,8 @@ RewriteSystem::performKnuthBendix(unsigned maxRuleCount,
 
           // We don't have to consider the same pair of rules more than once,
           // since those critical pairs were already resolved.
-          if (!CheckedOverlaps.insert(std::make_pair(i, j)).second)
+          unsigned k = from - lhs.getLHS().begin();
+          if (!CheckedOverlaps.insert(std::make_tuple(i, j, k)).second)
             return;
 
           // Try to repair the confluence violation by adding a new rule.
diff --git a/lib/AST/RequirementMachine/RequirementMachine.cpp b/lib/AST/RequirementMachine/RequirementMachine.cpp
@@ -212,6 +212,8 @@ RequirementMachine::RequirementMachine(RewriteContext &ctx)
   MaxRuleCount = langOpts.RequirementMachineMaxRuleCount;
   MaxRuleLength = langOpts.RequirementMachineMaxRuleLength;
   MaxConcreteNesting = langOpts.RequirementMachineMaxConcreteNesting;
+  MaxConcreteSize = langOpts.RequirementMachineMaxConcreteSize;
+  MaxTypeDifferences = langOpts.RequirementMachineMaxTypeDifferences;
   Stats = ctx.getASTContext().Stats;
 
   if (Stats)
@@ -247,6 +249,18 @@ void RequirementMachine::checkCompletionResult(CompletionResult result) const {
       out << "Rewrite system exceeded concrete type nesting depth limit\n";
       dump(out);
     });
+
+  case CompletionResult::MaxConcreteSize:
+    ABORT([&](auto &out) {
+      out << "Rewrite system exceeded concrete type size limit\n";
+      dump(out);
+    });
+
+  case CompletionResult::MaxTypeDifferences:
+    ABORT([&](auto &out) {
+      out << "Rewrite system exceeded concrete type difference limit\n";
+      dump(out);
+    });
   }
 }
 
@@ -496,16 +510,26 @@ RequirementMachine::computeCompletion(RewriteSystem::ValidityPolicy policy) {
           return std::make_pair(CompletionResult::MaxRuleLength,
                                 ruleCount + i);
         }
-        if (newRule.getNesting() > MaxConcreteNesting + System.getDeepestInitialRule()) {
+        auto nestingAndSize = newRule.getNestingAndSize();
+        if (nestingAndSize.first > MaxConcreteNesting + System.getMaxNestingOfInitialRule()) {
           return std::make_pair(CompletionResult::MaxConcreteNesting,
                                 ruleCount + i);
         }
+        if (nestingAndSize.second > MaxConcreteSize + System.getMaxSizeOfInitialRule()) {
+          return std::make_pair(CompletionResult::MaxConcreteSize,
+                                ruleCount + i);
+        }
       }
 
       if (System.getLocalRules().size() > MaxRuleCount) {
         return std::make_pair(CompletionResult::MaxRuleCount,
                               System.getRules().size() - 1);
       }
+
+      if (System.getTypeDifferenceCount() > MaxTypeDifferences) {
+        return std::make_pair(CompletionResult::MaxTypeDifferences,
+                              System.getRules().size() - 1);
+      }
     }
   }
 
diff --git a/lib/AST/RequirementMachine/RequirementMachine.h b/lib/AST/RequirementMachine/RequirementMachine.h
@@ -68,6 +68,8 @@ class RequirementMachine final {
   unsigned MaxRuleCount;
   unsigned MaxRuleLength;
   unsigned MaxConcreteNesting;
+  unsigned MaxConcreteSize;
+  unsigned MaxTypeDifferences;
 
   UnifiedStatsReporter *Stats;
 
diff --git a/lib/AST/RequirementMachine/RewriteSystem.cpp b/lib/AST/RequirementMachine/RewriteSystem.cpp
@@ -34,7 +34,8 @@ RewriteSystem::RewriteSystem(RewriteContext &ctx)
   Frozen = 0;
   RecordLoops = 0;
   LongestInitialRule = 0;
-  DeepestInitialRule = 0;
+  MaxNestingOfInitialRule = 0;
+  MaxSizeOfInitialRule = 0;
 }
 
 RewriteSystem::~RewriteSystem() {
@@ -90,7 +91,12 @@ void RewriteSystem::initialize(
 
   for (const auto &rule : getLocalRules()) {
     LongestInitialRule = std::max(LongestInitialRule, rule.getDepth());
-    DeepestInitialRule = std::max(DeepestInitialRule, rule.getNesting());
+
+    auto nestingAndSize = rule.getNestingAndSize();
+    MaxNestingOfInitialRule = std::max(MaxNestingOfInitialRule,
+                                      nestingAndSize.first);
+    MaxSizeOfInitialRule = std::max(MaxSizeOfInitialRule,
+                                    nestingAndSize.second);
   }
 }
 
diff --git a/lib/AST/RequirementMachine/RewriteSystem.h b/lib/AST/RequirementMachine/RewriteSystem.h
@@ -50,7 +50,13 @@ enum class CompletionResult {
   MaxRuleLength,
 
   /// Maximum concrete type nesting depth exceeded.
-  MaxConcreteNesting
+  MaxConcreteNesting,
+
+  /// Maximum concrete type size exceeded.
+  MaxConcreteSize,
+
+  /// Maximum type difference count exceeded.
+  MaxTypeDifferences,
 };
 
 /// A term rewrite system for working with types in a generic signature.
@@ -107,13 +113,16 @@ class RewriteSystem final {
   /// identities among rewrite rules discovered while resolving critical pairs.
   unsigned RecordLoops : 1;
 
-  /// The length of the longest initial rule, used for the MaxRuleLength
-  /// completion non-termination heuristic.
+  /// The length of the longest initial rule, for the MaxRuleLength limit.
   unsigned LongestInitialRule : 16;
 
-  /// The most deeply nested concrete type appearing in an initial rule, used
-  /// for the MaxConcreteNesting completion non-termination heuristic.
-  unsigned DeepestInitialRule : 16;
+  /// The most deeply nested concrete type appearing in an initial rule,
+  /// for the MaxConcreteNesting limit.
+  unsigned MaxNestingOfInitialRule : 16;
+
+  /// The largest concrete type by total tree node count that appears in an
+  /// initial rule, for the MaxConcreteSize limit.
+  unsigned MaxSizeOfInitialRule : 16;
 
 public:
   explicit RewriteSystem(RewriteContext &ctx);
@@ -143,8 +152,12 @@ class RewriteSystem final {
     return LongestInitialRule;
   }
 
-  unsigned getDeepestInitialRule() const {
-    return DeepestInitialRule;
+  unsigned getMaxNestingOfInitialRule() const {
+    return MaxNestingOfInitialRule;
+  }
+
+  unsigned getMaxSizeOfInitialRule() const {
+    return MaxSizeOfInitialRule;
   }
 
   ArrayRef<const ProtocolDecl *> getProtocols() const {
@@ -206,7 +219,7 @@ class RewriteSystem final {
   //////////////////////////////////////////////////////////////////////////////
 
   /// Pairs of rules which have already been checked for overlap.
-  llvm::DenseSet<std::pair<unsigned, unsigned>> CheckedOverlaps;
+  llvm::DenseSet<std::tuple<unsigned, unsigned, unsigned>> CheckedOverlaps;
 
   std::pair<CompletionResult, unsigned>
   performKnuthBendix(unsigned maxRuleCount, unsigned maxRuleLength);
@@ -311,6 +324,10 @@ class RewriteSystem final {
                              std::optional<unsigned> &lhsDifferenceID,
                              std::optional<unsigned> &rhsDifferenceID);
 
+  unsigned getTypeDifferenceCount() const {
+    return Differences.size();
+  }
+
   const TypeDifference &getTypeDifference(unsigned index) const;
 
   void processTypeDifference(const TypeDifference &difference,
diff --git a/lib/AST/RequirementMachine/Rule.cpp b/lib/AST/RequirementMachine/Rule.cpp
@@ -221,7 +221,12 @@ bool Rule::isDerivedFromConcreteProtocolTypeAliasRule() const {
   return true;
 }
 
-/// Returns the length of the left hand side.
+/// Returns the maximum among the length of the left-hand side,
+/// and the length of any substitution terms that appear in a
+/// property symbol at the end of the left-hand side.
+///
+/// This is a measure of the complexity of the rule, which stops
+/// completion from running forever.
 unsigned Rule::getDepth() const {
   auto result = LHS.size();
 
@@ -234,37 +239,41 @@ unsigned Rule::getDepth() const {
   return result;
 }
 
-/// Returns the nesting depth of the concrete symbol at the end of the
-/// left hand side, or 0 if there isn't one.
-unsigned Rule::getNesting() const {
+/// Returns the complexity of the concrete type in the property symbol
+/// at the end of the left-hand side, if there is one.
+///
+/// This is a measure of the complexity of the rule, which stops
+/// completion from running forever.
+std::pair<unsigned, unsigned>
+Rule::getNestingAndSize() const {
   if (LHS.back().hasSubstitutions()) {
     auto type = LHS.back().getConcreteType();
 
     struct Walker : TypeWalker {
       unsigned Nesting = 0;
       unsigned MaxNesting = 0;
+      unsigned Size = 0;
 
       Action walkToTypePre(Type ty) override {
+        ++Size;
         ++Nesting;
-        MaxNesting = std::max(Nesting, MaxNesting);
-
+        MaxNesting = std::max(MaxNesting, Nesting);
         return Action::Continue;
       }
 
       Action walkToTypePost(Type ty) override {
         --Nesting;
-
         return Action::Continue;
       }
     };
 
     Walker walker;
     type.walk(walker);
 
-    return walker.MaxNesting;
+    return std::make_pair(walker.MaxNesting, walker.Size);
   }
 
-  return 0;
+  return std::make_pair(0, 0);
 }
 
 /// Linear order on rules; compares LHS followed by RHS.
diff --git a/lib/AST/RequirementMachine/Rule.h b/lib/AST/RequirementMachine/Rule.h
@@ -224,7 +224,7 @@ class Rule final {
 
   unsigned getDepth() const;
 
-  unsigned getNesting() const;
+  std::pair<unsigned, unsigned> getNestingAndSize() const;
 
   std::optional<int> compare(const Rule &other, RewriteContext &ctx) const;
 
diff --git a/lib/Frontend/CompilerInvocation.cpp b/lib/Frontend/CompilerInvocation.cpp
@@ -1769,6 +1769,28 @@ static bool ParseLangArgs(LangOptions &Opts, ArgList &Args,
     }
   }
 
+  if (const Arg *A = Args.getLastArg(OPT_requirement_machine_max_concrete_size)) {
+    unsigned limit;
+    if (StringRef(A->getValue()).getAsInteger(10, limit)) {
+      Diags.diagnose(SourceLoc(), diag::error_invalid_arg_value,
+                     A->getAsString(Args), A->getValue());
+      HadError = true;
+    } else {
+      Opts.RequirementMachineMaxConcreteSize = limit;
+    }
+  }
+
+  if (const Arg *A = Args.getLastArg(OPT_requirement_machine_max_type_differences)) {
+    unsigned limit;
+    if (StringRef(A->getValue()).getAsInteger(10, limit)) {
+      Diags.diagnose(SourceLoc(), diag::error_invalid_arg_value,
+                     A->getAsString(Args), A->getValue());
+      HadError = true;
+    } else {
+      Opts.RequirementMachineMaxTypeDifferences = limit;
+    }
+  }
+
   if (const Arg *A = Args.getLastArg(OPT_requirement_machine_max_split_concrete_equiv_class_attempts)) {
     unsigned limit;
     if (StringRef(A->getValue()).getAsInteger(10, limit)) {
diff --git a/test/Constraints/same_types.swift b/test/Constraints/same_types.swift
@@ -284,7 +284,7 @@ protocol P2 {
 // CHECK-LABEL: same_types.(file).structuralSameTypeRecursive1@
 // CHECK-NEXT: Generic signature: <T, U>
 
-// expected-error@+2 {{cannot build rewrite system for generic signature; concrete nesting limit exceeded}}
+// expected-error@+2 {{cannot build rewrite system for generic signature; concrete type nesting limit exceeded}}
 // expected-note@+1 {{τ_0_0.[P2:Assoc1].[concrete: ((((((((((((((((((((((((((((((((τ_0_0.[P2:Assoc1], τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1), τ_0_1)] => τ_0_0.[P2:Assoc1]}}
 func structuralSameTypeRecursive1<T: P2, U>(_: T, _: U)
   where T.Assoc1 == Tuple2<T.Assoc1, U>
diff --git a/test/Generics/infinite_concrete_type.swift b/test/Generics/infinite_concrete_type.swift
@@ -2,8 +2,8 @@
 
 class G<T> {}
 
-protocol P1 { // expected-error {{cannot build rewrite system for protocol; concrete nesting limit exceeded}}
-// expected-note@-1 {{failed rewrite rule is [P1:A].[concrete: G<G<G<G<G<G<G<G<G<G<G<G<G<G<G<G<G<G<G<G<G<G<G<G<G<G<G<G<G<G<G<G<[P1].A>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>] => [P1:A]}}
+protocol P1 { // expected-error {{cannot build rewrite system for protocol; concrete type difference limit exceeded}}
+// expected-note@-1 {{failed rewrite rule is [P1:B].[superclass: G<G<G<G<G<G<G<G<G<G<G<G<G<[P1].A>>>>>>>>>>>>>] => [P1:B]}}
   associatedtype A where A == G<B>
   associatedtype B where B == G<A>
 }
diff --git a/test/Generics/non_confluent.swift b/test/Generics/non_confluent.swift
diff --git a/test/Generics/rdar90402519.swift b/test/Generics/rdar90402519.swift
diff --git a/test/attr/attr_specialize.swift b/test/attr/attr_specialize.swift
diff --git a/test/decl/protocol/req/recursion.swift b/test/decl/protocol/req/recursion.swift
diff --git a/validation-test/compiler_crashers_2_fixed/b93e9e54c72f13c9.swift b/validation-test/compiler_crashers_2_fixed/b93e9e54c72f13c9.swift

Original file line number	Diff line number	Diff line change
`@@ -1436,7 +1436,8 @@ RequirementSignature RequirementSignature::getPlaceholderRequirementSignature(`
`1436`	`1436`	`});`
`1437`	`1437`
`1438`	`1438`	`return RequirementSignature(ctx.AllocateCopy(requirements),`
`1439`		`- ArrayRef<ProtocolTypeAlias>());`
	`1439`	`+ ArrayRef<ProtocolTypeAlias>(),`
	`1440`	`+ errors);`
`1440`	`1441`	`}`
`1441`	`1442`
`1442`	`1443`	`void RequirementSignature::getRequirementsWithInverses(`