RequirementMachine: Completion step and depth limits also apply to equivalence class map

Author: slavapestov

lib/AST/RequirementMachine/EquivalenceClassMap.cpp

@@ -485,10 +485,19 @@ void EquivalenceClassMap::dump(llvm::raw_ostream &out) const {
 /// Build the equivalence class map from all rules of the form T.[p] => T, where
 /// [p] is a property atom.
 ///
-/// Returns true if concrete term unification performed while building the map
-/// introduced new rewrite rules; in this case, the completion procedure must
-/// run again.
-bool RewriteSystem::buildEquivalenceClassMap(EquivalenceClassMap &map) {
+/// Returns a pair consisting of a status and number of iterations executed.
+///
+/// The status is CompletionResult::MaxIterations if we exceed \p maxIterations
+/// iterations.
+///
+/// The status is CompletionResult::MaxDepth if we produce a rewrite rule whose
+/// left hand side has a length exceeding \p maxDepth.
+///
+/// Otherwise, the status is CompletionResult::Success.
+std::pair<RewriteSystem::CompletionResult, unsigned>
+RewriteSystem::buildEquivalenceClassMap(EquivalenceClassMap &map,
+                                        unsigned maxIterations,
+                                        unsigned maxDepth) {
   map.clear();
 
   std::vector<std::pair<MutableTerm, Atom>> properties;
@@ -537,14 +546,17 @@ bool RewriteSystem::buildEquivalenceClassMap(EquivalenceClassMap &map) {
   // where the left hand side is not already equivalent to the right hand side.
   unsigned addedNewRules = 0;
   for (auto pair : inducedRules) {
-    if (addRule(pair.first, pair.second))
+    if (addRule(pair.first, pair.second)) {
       ++addedNewRules;
-  }
 
-  if (auto *stats = Context.getASTContext().Stats) {
-    stats->getFrontendCounters()
-      .NumRequirementMachineUnifiedConcreteTerms += addedNewRules;
+      const auto &newRule = Rules.back();
+      if (newRule.getLHS().size() > maxDepth)
+        return std::make_pair(CompletionResult::MaxDepth, addedNewRules);
+    }
   }
 
-  return addedNewRules > 0;
+  if (Rules.size() > maxIterations)
+    return std::make_pair(CompletionResult::MaxIterations, addedNewRules);
+
+  return std::make_pair(CompletionResult::Success, addedNewRules);
 }

lib/AST/RequirementMachine/RequirementMachine.cpp

@@ -280,39 +280,38 @@ void RequirementMachine::addGenericSignature(CanGenericSignature sig) {
 /// Attempt to obtain a confluent rewrite system using the completion
 /// procedure.
 void RequirementMachine::computeCompletion(CanGenericSignature sig) {
-  unsigned remainingSteps = Context.LangOpts.RequirementMachineStepLimit;
-  bool keepGoing = true;
-
-  while (keepGoing) {
+  while (true) {
     // First, run the Knuth-Bendix algorithm to resolve overlapping rules.
     auto result = Impl->System.computeConfluentCompletion(
-        remainingSteps, Context.LangOpts.RequirementMachineDepthLimit);
-
-    assert(remainingSteps >= result.second);
-    remainingSteps -= result.second;
+        Context.LangOpts.RequirementMachineStepLimit,
+        Context.LangOpts.RequirementMachineDepthLimit);
 
     if (Context.Stats) {
       Context.Stats->getFrontendCounters()
           .NumRequirementMachineCompletionSteps += result.second;
     }
 
     // Check for failure.
-    switch (result.first) {
-    case RewriteSystem::CompletionResult::Success:
-      break;
-
-    case RewriteSystem::CompletionResult::MaxIterations:
-      llvm::errs() << "Generic signature " << sig
-                   << " exceeds maximum completion step count\n";
-      Impl->System.dump(llvm::errs());
-      abort();
-
-    case RewriteSystem::CompletionResult::MaxDepth:
-      llvm::errs() << "Generic signature " << sig
-                   << " exceeds maximum completion depth\n";
-      Impl->System.dump(llvm::errs());
-      abort();
-    }
+    auto checkCompletionResult = [&]() {
+      switch (result.first) {
+      case RewriteSystem::CompletionResult::Success:
+        break;
+
+      case RewriteSystem::CompletionResult::MaxIterations:
+        llvm::errs() << "Generic signature " << sig
+                     << " exceeds maximum completion step count\n";
+        Impl->System.dump(llvm::errs());
+        abort();
+
+      case RewriteSystem::CompletionResult::MaxDepth:
+        llvm::errs() << "Generic signature " << sig
+                     << " exceeds maximum completion depth\n";
+        Impl->System.dump(llvm::errs());
+        abort();
+      }
+    };
+
+    checkCompletionResult();
 
     // Simplify right hand sides in preparation for building the
     // equivalence class map.
@@ -321,7 +320,22 @@ void RequirementMachine::computeCompletion(CanGenericSignature sig) {
     // Build the equivalence class map, which performs concrete term
     // unification; if this added any new rules, run the completion
     // procedure again.
-    keepGoing = Impl->System.buildEquivalenceClassMap(Impl->Map);
+    result = Impl->System.buildEquivalenceClassMap(
+        Impl->Map,
+        Context.LangOpts.RequirementMachineStepLimit,
+        Context.LangOpts.RequirementMachineDepthLimit);
+
+    if (Context.Stats) {
+      Context.Stats->getFrontendCounters()
+        .NumRequirementMachineUnifiedConcreteTerms += result.second;
+    }
+
+    checkCompletionResult();
+
+    // If buildEquivalenceClassMap() added new rules, we run another
+    // round of Knuth-Bendix, and build the equivalence class map again.
+    if (result.second == 0)
+      break;
   }
 
   if (Context.LangOpts.DebugRequirementMachine) {

lib/AST/RequirementMachine/RewriteSystem.h

@@ -580,11 +580,15 @@ class RewriteSystem final {
   };
 
   std::pair<CompletionResult, unsigned>
-  computeConfluentCompletion(unsigned maxIterations, unsigned maxDepth);
+  computeConfluentCompletion(unsigned maxIterations,
+                             unsigned maxDepth);
 
   void simplifyRightHandSides();
 
-  bool buildEquivalenceClassMap(EquivalenceClassMap &map);
+  std::pair<CompletionResult, unsigned>
+  buildEquivalenceClassMap(EquivalenceClassMap &map,
+                           unsigned maxIterations,
+                           unsigned maxDepth);
 
   void dump(llvm::raw_ostream &out) const;
 

lib/AST/RequirementMachine/RewriteSystemCompletion.cpp

@@ -446,7 +446,7 @@ RewriteSystem::computeConfluentCompletion(unsigned maxIterations,
   // moving on to overlaps between rules introduced by completion.
   while (!Worklist.empty()) {
     // Check if we've already done too much work.
-    if (steps >= maxIterations)
+    if (Rules.size() > maxIterations)
       return std::make_pair(CompletionResult::MaxIterations, steps);
 
     auto next = Worklist.front();