RequirementMachine: Clean up merged associated type algorithm some more

The left and right hand side of a merged associated type candidate rule
have a common prefix except for the associated type symbol at the end.

Instead of passing two MutableTerms, we can pass a single uniqued
Term and a Symbol.

Also, we can compute the merged symbol before pushing the candidate
onto the vector. This avoids unnecessary work if the merged symbol
is equal to the right hand side's symbol.

Author: slavapestov

lib/AST/RequirementMachine/RewriteSystem.cpp

@@ -97,7 +97,11 @@ bool RewriteSystem::addRule(MutableTerm lhs, MutableTerm rhs) {
   }
 
   unsigned i = Rules.size();
-  Rules.emplace_back(Term::get(lhs, Context), Term::get(rhs, Context));
+
+  auto uniquedLHS = Term::get(lhs, Context);
+  auto uniquedRHS = Term::get(rhs, Context);
+  Rules.emplace_back(uniquedLHS, uniquedRHS);
+
   auto oldRuleID = Trie.insert(lhs.begin(), lhs.end(), i);
   if (oldRuleID) {
     llvm::errs() << "Duplicate rewrite rule!\n";
@@ -112,23 +116,7 @@ bool RewriteSystem::addRule(MutableTerm lhs, MutableTerm rhs) {
     abort();
   }
 
-  // Check if we have a rule of the form
-  //
-  //   X.[P1:T] => X.[P2:T]
-  //
-  // If so, record this rule for later. We'll try to merge the associated
-  // types in RewriteSystem::processMergedAssociatedTypes().
-  if (lhs.size() == rhs.size() &&
-      std::equal(lhs.begin(), lhs.end() - 1, rhs.begin()) &&
-      lhs.back().getKind() == Symbol::Kind::AssociatedType &&
-      rhs.back().getKind() == Symbol::Kind::AssociatedType &&
-      lhs.back().getName() == rhs.back().getName()) {
-    if (Debug.contains(DebugFlags::Merge)) {
-      llvm::dbgs() << "## Associated type merge candidate ";
-      llvm::dbgs() << lhs << " => " << rhs << "\n\n";
-    }
-    MergedAssociatedTypes.emplace_back(lhs, rhs);
-  }
+  checkMergedAssociatedType(uniquedLHS, uniquedRHS);
 
   // Tell the caller that we added a new rule.
   return true;

lib/AST/RequirementMachine/RewriteSystem.h

@@ -97,17 +97,24 @@ class RewriteSystem final {
   /// rewrite rules, used for the linear order on symbols.
   ProtocolGraph Protos;
 
+  /// Constructed from a rule of the form X.[P2:T] => X.[P1:T] by
+  /// checkMergedAssociatedType().
+  struct MergedAssociatedType {
+    /// The *right* hand side of the original rule, X.[P1:T].
+    Term rhs;
+
+    /// The associated type symbol appearing at the end of the *left*
+    /// hand side of the original rule, [P2:T].
+    Symbol lhsSymbol;
+
+    /// The merged associated type symbol, [P1&P2:T].
+    Symbol mergedSymbol;
+  };
+
   /// A list of pending terms for the associated type merging completion
-  /// heuristic.
-  ///
-  /// The pair (lhs, rhs) satisfies the following conditions:
-  /// - lhs > rhs
-  /// - all symbols but the last are pair-wise equal in lhs and rhs
-  /// - the last symbol in both lhs and rhs is an associated type symbol
-  /// - the last symbol in both lhs and rhs has the same name
-  ///
-  /// See RewriteSystem::processMergedAssociatedTypes() for details.
-  std::vector<std::pair<MutableTerm, MutableTerm>> MergedAssociatedTypes;
+  /// heuristic. Entries are added by checkMergedAssociatedType(), and
+  /// consumed in processMergedAssociatedTypes().
+  std::vector<MergedAssociatedType> MergedAssociatedTypes;
 
   /// Pairs of rules which have already been checked for overlap.
   llvm::DenseSet<std::pair<unsigned, unsigned>> CheckedOverlaps;
@@ -173,6 +180,8 @@ class RewriteSystem final {
                                             MutableTerm>> &result) const;
 
   void processMergedAssociatedTypes();
+
+  void checkMergedAssociatedType(Term lhs, Term rhs);
 };
 
 } // end namespace rewriting

lib/AST/RequirementMachine/RewriteSystemCompletion.cpp

@@ -206,49 +206,28 @@ void RewriteSystem::processMergedAssociatedTypes() {
 
   unsigned i = 0;
 
-  // Chase the end of the vector; calls to RewriteSystem::addRule()
-  // can theoretically add new elements below.
+  // Chase the end of the vector, since addRule() might add new elements below.
   while (i < MergedAssociatedTypes.size()) {
-    auto pair = MergedAssociatedTypes[i++];
-    const auto &lhs = pair.first;
-    const auto &rhs = pair.second;
+    // Copy the entry out, since addRule() might add new elements below.
+    auto entry = MergedAssociatedTypes[i++];
 
-    // If we have X.[P2:T] => Y.[P1:T], add a new pair of rules:
-    // X.[P1:T] => X.[P1&P2:T]
-    // X.[P2:T] => X.[P1&P2:T]
     if (Debug.contains(DebugFlags::Merge)) {
-      llvm::dbgs() << "## Processing associated type merge candidate ";
-      llvm::dbgs() << lhs << " => " << rhs << "\n";
-    }
-
-    auto mergedSymbol = Context.mergeAssociatedTypes(lhs.back(), rhs.back(),
-                                                     Protos);
-    if (Debug.contains(DebugFlags::Merge)) {
-      llvm::dbgs() << "### Merged symbol " << mergedSymbol << "\n";
+      llvm::dbgs() << "## Processing associated type merge with ";
+      llvm::dbgs() << entry.rhs << ", ";
+      llvm::dbgs() << entry.lhsSymbol << ", ";
+      llvm::dbgs() << entry.mergedSymbol << "\n";
     }
 
-    // We must have mergedSymbol <= rhs < lhs, therefore mergedSymbol != lhs.
-    assert(lhs.back() != mergedSymbol &&
-           "Left hand side should not already end with merged symbol?");
-    assert(mergedSymbol.compare(rhs.back(), Protos) <= 0);
-    assert(rhs.back().compare(lhs.back(), Protos) < 0);
-
-    // If the merge didn't actually produce a new symbol, there is nothing else
-    // to do.
-    if (rhs.back() == mergedSymbol) {
-      if (Debug.contains(DebugFlags::Merge)) {
-        llvm::dbgs() << "### Skipping\n";
-      }
-
-      continue;
-    }
+    // If we have X.[P2:T] => Y.[P1:T], add a new rule:
+    // X.[P1:T] => X.[P1&P2:T]
+    MutableTerm lhs(entry.rhs);
 
     // Build the term X.[P1&P2:T].
-    MutableTerm mergedTerm = lhs;
-    mergedTerm.back() = mergedSymbol;
+    MutableTerm rhs(entry.rhs);
+    rhs.back() = entry.mergedSymbol;
 
     // Add the rule X.[P1:T] => X.[P1&P2:T].
-    addRule(rhs, mergedTerm);
+    addRule(lhs, rhs);
 
     // Collect new rules here so that we're not adding rules while traversing
     // the trie.
@@ -260,8 +239,8 @@ void RewriteSystem::processMergedAssociatedTypes() {
       const auto &otherLHS = otherRule.getLHS();
       if (otherLHS.size() == 2 &&
           otherLHS[1].getKind() == Symbol::Kind::Protocol) {
-        if (otherLHS[0] == lhs.back() ||
-            otherLHS[0] == rhs.back()) {
+        if (otherLHS[0] == entry.lhsSymbol ||
+            otherLHS[0] == entry.rhs.back()) {
           // We have a rule of the form
           //
           //   [P1:T].[Q] => [P1:T]
@@ -280,11 +259,11 @@ void RewriteSystem::processMergedAssociatedTypes() {
           //   [P1&P2:T].[Q] => [P1&P2:T]
           //
           MutableTerm newLHS;
-          newLHS.add(mergedSymbol);
+          newLHS.add(entry.mergedSymbol);
           newLHS.add(otherLHS[1]);
 
           MutableTerm newRHS;
-          newRHS.add(mergedSymbol);
+          newRHS.add(entry.mergedSymbol);
 
           inducedRules.emplace_back(newLHS, newRHS);
         }
@@ -294,8 +273,8 @@ void RewriteSystem::processMergedAssociatedTypes() {
     // Visit rhs first to preserve the ordering of protocol requirements in the
     // the property map. This is just for aesthetic purposes in the debug dump,
     // it doesn't change behavior.
-    Trie.findAll(rhs.back(), visitRule);
-    Trie.findAll(lhs.back(), visitRule);
+    Trie.findAll(entry.rhs.back(), visitRule);
+    Trie.findAll(entry.lhsSymbol, visitRule);
 
     // Now add the new rules.
     for (const auto &pair : inducedRules)
@@ -305,6 +284,49 @@ void RewriteSystem::processMergedAssociatedTypes() {
   MergedAssociatedTypes.clear();
 }
 
+/// Check if we have a rule of the form
+///
+///   X.[P1:T] => X.[P2:T]
+///
+/// If so, record this rule for later. We'll try to merge the associated
+/// types in RewriteSystem::processMergedAssociatedTypes().
+void RewriteSystem::checkMergedAssociatedType(Term lhs, Term rhs) {
+  if (lhs.size() == rhs.size() &&
+      std::equal(lhs.begin(), lhs.end() - 1, rhs.begin()) &&
+      lhs.back().getKind() == Symbol::Kind::AssociatedType &&
+      rhs.back().getKind() == Symbol::Kind::AssociatedType &&
+      lhs.back().getName() == rhs.back().getName()) {
+    if (Debug.contains(DebugFlags::Merge)) {
+      llvm::dbgs() << "## Associated type merge candidate ";
+      llvm::dbgs() << lhs << " => " << rhs << "\n\n";
+    }
+
+    auto mergedSymbol = Context.mergeAssociatedTypes(lhs.back(), rhs.back(),
+                                                     Protos);
+    if (Debug.contains(DebugFlags::Merge)) {
+      llvm::dbgs() << "### Merged symbol " << mergedSymbol << "\n";
+    }
+
+    // We must have mergedSymbol <= rhs < lhs, therefore mergedSymbol != lhs.
+    assert(lhs.back() != mergedSymbol &&
+           "Left hand side should not already end with merged symbol?");
+    assert(mergedSymbol.compare(rhs.back(), Protos) <= 0);
+    assert(rhs.back().compare(lhs.back(), Protos) < 0);
+
+    // If the merge didn't actually produce a new symbol, there is nothing else
+    // to do.
+    if (rhs.back() == mergedSymbol) {
+      if (Debug.contains(DebugFlags::Merge)) {
+        llvm::dbgs() << "### Skipping\n";
+      }
+
+      return;
+    }
+
+    MergedAssociatedTypes.push_back({rhs, lhs.back(), mergedSymbol});
+  }
+}
+
 /// Compute a critical pair from the left hand sides of two rewrite rules,
 /// where \p rhs begins at \p from, which must be an iterator pointing
 /// into \p lhs.