RequirementMachine: Speed up term simplification with a prefix trie

slavapestov · slavapestov · commit 156fa2cc0436 · 2021-08-05T21:42:50.000-04:00
Previously RewriteSystem::simplify() would attempt to apply every
rewrite rule at every position in the original term, which was
obviously a source of overhead.

The trie itself is somewhat unoptimized; for example, with a bit of
effort it could merge a node with its only child, if nodes stored
a range of elements to compare rather than a single element.
diff --git a/lib/AST/RequirementMachine/RewriteContext.cpp b/lib/AST/RequirementMachine/RewriteContext.cpp
@@ -23,7 +23,9 @@ RewriteContext::RewriteContext(ASTContext &ctx)
     : Context(ctx),
       Stats(ctx.Stats),
       SymbolHistogram(Symbol::NumKinds),
-      TermHistogram(4, /*Start=*/1) {
+      TermHistogram(4, /*Start=*/1),
+      RuleTrieHistogram(16, /*Start=*/1),
+      RuleTrieRootHistogram(16)  {
 }
 
 Term RewriteContext::getTermForType(CanType paramType,
@@ -289,5 +291,9 @@ RewriteContext::~RewriteContext() {
     SymbolHistogram.dump(llvm::dbgs(), Symbol::Kinds);
     llvm::dbgs() << "\n* Term length:\n";
     TermHistogram.dump(llvm::dbgs());
+    llvm::dbgs() << "\n* Rule trie fanout:\n";
+    RuleTrieHistogram.dump(llvm::dbgs());
+    llvm::dbgs() << "\n* Rule trie root fanout:\n";
+    RuleTrieRootHistogram.dump(llvm::dbgs());
   }
-}
+}
diff --git a/lib/AST/RequirementMachine/RewriteContext.h b/lib/AST/RequirementMachine/RewriteContext.h
@@ -58,6 +58,8 @@ class RewriteContext final {
   /// Histograms.
   Histogram SymbolHistogram;
   Histogram TermHistogram;
+  Histogram RuleTrieHistogram;
+  Histogram RuleTrieRootHistogram;
 
   explicit RewriteContext(ASTContext &ctx);
 
diff --git a/lib/AST/RequirementMachine/RewriteSystem.cpp b/lib/AST/RequirementMachine/RewriteSystem.cpp
@@ -31,6 +31,11 @@ RewriteSystem::RewriteSystem(RewriteContext &ctx)
   DebugCompletion = false;
 }
 
+RewriteSystem::~RewriteSystem() {
+  RuleTrie.updateHistograms(Context.RuleTrieHistogram,
+                            Context.RuleTrieRootHistogram);
+}
+
 void Rule::dump(llvm::raw_ostream &out) const {
   out << LHS << " => " << RHS;
   if (deleted)
@@ -94,6 +99,19 @@ bool RewriteSystem::addRule(MutableTerm lhs, MutableTerm rhs) {
 
   unsigned i = Rules.size();
   Rules.emplace_back(Term::get(lhs, Context), Term::get(rhs, Context));
+  auto oldRuleID = RuleTrie.insert(lhs.begin(), lhs.end(), i);
+  if (oldRuleID) {
+    llvm::errs() << "Duplicate rewrite rule!\n";
+    const auto &oldRule = Rules[*oldRuleID];
+    llvm::errs() << "Old rule #" << *oldRuleID << ": ";
+    oldRule.dump(llvm::errs());
+    llvm::errs() << "\nTrying to replay what happened when I simplified this term:\n";
+    DebugSimplify = true;
+    MutableTerm term = lhs;
+    simplify(lhs);
+
+    abort();
+  }
 
   // Check if we have a rule of the form
   //
@@ -145,22 +163,34 @@ bool RewriteSystem::simplify(MutableTerm &term) const {
 
   while (true) {
     bool tryAgain = false;
-    for (const auto &rule : Rules) {
-      if (rule.isDeleted())
-        continue;
 
-      if (DebugSimplify) {
-        llvm::dbgs() << "== Rule " << rule << "\n";
-      }
-
-      if (rule.apply(term)) {
-        if (DebugSimplify) {
-          llvm::dbgs() << "=== Result " << term << "\n";
+    auto from = term.begin();
+    auto end = term.end();
+    while (from < end) {
+      auto ruleID = RuleTrie.find(from, end);
+      if (ruleID) {
+        const auto &rule = Rules[*ruleID];
+        if (!rule.isDeleted()) {
+          if (DebugSimplify) {
+            llvm::dbgs() << "== Rule #" << *ruleID << ": " << rule << "\n";
+          }
+
+          auto to = from + rule.getLHS().size();
+          assert(std::equal(from, to, rule.getLHS().begin()));
+
+          term.rewriteSubTerm(from, to, rule.getRHS());
+
+          if (DebugSimplify) {
+            llvm::dbgs() << "=== Result " << term << "\n";
+          }
+
+          changed = true;
+          tryAgain = true;
+          break;
         }
-
-        changed = true;
-        tryAgain = true;
       }
+
+      ++from;
     }
 
     if (!tryAgain)
diff --git a/lib/AST/RequirementMachine/RewriteSystem.h b/lib/AST/RequirementMachine/RewriteSystem.h
@@ -18,6 +18,7 @@
 #include "ProtocolGraph.h"
 #include "Symbol.h"
 #include "Term.h"
+#include "Trie.h"
 
 namespace llvm {
   class raw_ostream;
@@ -47,10 +48,6 @@ class Rule final {
   const Term &getLHS() const { return LHS; }
   const Term &getRHS() const { return RHS; }
 
-  bool apply(MutableTerm &term) const {
-    return term.rewriteSubTerm(LHS, RHS);
-  }
-
   OverlapKind checkForOverlap(const Rule &other,
                               MutableTerm &t,
                               MutableTerm &v) const {
@@ -101,6 +98,9 @@ class RewriteSystem final {
   /// as rules introduced by the completion procedure.
   std::vector<Rule> Rules;
 
+  /// A prefix trie of rule left hand sides to optimize lookup.
+  Trie RuleTrie;
+
   /// The graph of all protocols transitively referenced via our set of
   /// rewrite rules, used for the linear order on symbols.
   ProtocolGraph Protos;
@@ -129,6 +129,7 @@ class RewriteSystem final {
 
 public:
   explicit RewriteSystem(RewriteContext &ctx);
+  ~RewriteSystem();
 
   RewriteSystem(const RewriteSystem &) = delete;
   RewriteSystem(RewriteSystem &&) = delete;
diff --git a/lib/AST/RequirementMachine/Term.cpp b/lib/AST/RequirementMachine/Term.cpp
@@ -147,46 +147,21 @@ int MutableTerm::compare(const MutableTerm &other,
   return 0;
 }
 
-/// Find the start of \p other in this term, returning end() if
-/// \p other does not occur as a subterm of this term.
-decltype(MutableTerm::Symbols)::const_iterator
-MutableTerm::findSubTerm(Term other) const {
-  if (other.size() > size())
-    return end();
-
-  return std::search(begin(), end(), other.begin(), other.end());
-}
-
-/// Non-const variant of the above.
-decltype(MutableTerm::Symbols)::iterator
-MutableTerm::findSubTerm(Term other) {
-  if (other.size() > size())
-    return end();
-
-  return std::search(begin(), end(), other.begin(), other.end());
-}
-
-/// Replace the first occurrence of \p lhs in this term with
-/// \p rhs. Note that \p rhs must precede \p lhs in the linear
-/// order on terms. Returns true if the term contained \p lhs;
-/// otherwise returns false, in which case the term remains
-/// unchanged.
-bool MutableTerm::rewriteSubTerm(Term lhs, Term rhs) {
-  // Find the start of lhs in this term.
-  auto found = findSubTerm(lhs);
-
-  // This term cannot be reduced using this rule.
-  if (found == end())
-    return false;
-
+/// Replace the subterm in the range [from,to) with \p rhs.
+///
+/// Note that \p rhs must precede [from,to) in the linear
+/// order on terms.
+void MutableTerm::rewriteSubTerm(
+    decltype(MutableTerm::Symbols)::iterator from,
+    decltype(MutableTerm::Symbols)::iterator to,
+    Term rhs) {
   auto oldSize = size();
-
-  assert(rhs.size() <= lhs.size());
+  unsigned lhsLength = (unsigned)(to - from);
+  assert(rhs.size() <= lhsLength);
 
   // Overwrite the occurrence of the left hand side with the
   // right hand side.
-  auto newIter = std::copy(rhs.begin(), rhs.end(), found);
-  auto oldIter = found + lhs.size();
+  auto newIter = std::copy(rhs.begin(), rhs.end(), from);
 
   // If the right hand side is shorter than the left hand side,
   // then newIter will point to a location before oldIter, eg
@@ -199,16 +174,15 @@ bool MutableTerm::rewriteSubTerm(Term lhs, Term rhs) {
   //
   // Shift everything over to close the gap (by one location,
   // in this case).
-  if (newIter != oldIter) {
-    auto newEnd = std::copy(oldIter, end(), newIter);
+  if (newIter != to) {
+    auto newEnd = std::copy(to, end(), newIter);
 
     // Now, we've moved the gap to the end of the term; close
     // it by shortening the term.
     Symbols.erase(newEnd, end());
   }
 
-  assert(size() == oldSize - lhs.size() + rhs.size());
-  return true;
+  assert(size() == oldSize - lhsLength + rhs.size());
 }
 
 void MutableTerm::dump(llvm::raw_ostream &out) const {
diff --git a/lib/AST/RequirementMachine/Term.h b/lib/AST/RequirementMachine/Term.h
@@ -191,11 +191,9 @@ class MutableTerm final {
     return Symbols[index];
   }
 
-  decltype(Symbols)::const_iterator findSubTerm(Term other) const;
-
-  decltype(Symbols)::iterator findSubTerm(Term other);
-
-  bool rewriteSubTerm(Term lhs, Term rhs);
+  void rewriteSubTerm(decltype(Symbols)::iterator from,
+                      decltype(Symbols)::iterator to,
+                      Term rhs);
 
   void dump(llvm::raw_ostream &out) const;
 
diff --git a/lib/AST/RequirementMachine/Trie.h b/lib/AST/RequirementMachine/Trie.h
@@ -0,0 +1,120 @@
+//===--- Trie.h - Trie with terms as keys ---------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2021 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_RQM_TRIE_H
+#define SWIFT_RQM_TRIE_H
+
+#include "llvm/ADT/DenseMap.h"
+#include "Histogram.h"
+
+namespace swift {
+
+namespace rewriting {
+
+class Trie {
+public:
+  struct Node;
+
+  struct Entry {
+    Optional<unsigned> RuleID;
+    Node *Children = nullptr;
+  };
+
+  struct Node {
+    llvm::SmallDenseMap<Symbol, Entry, 1> Entries;
+  };
+
+private:
+  /// We never delete nodes, except for when the entire trie is torn down.
+  std::vector<Node *> Nodes;
+
+  /// The root is stored directly.
+  Node Root;
+
+public:
+  void updateHistograms(Histogram &stats, Histogram &rootStats) const {
+    for (const auto &node : Nodes)
+      stats.add(node->Entries.size());
+    rootStats.add(Root.Entries.size());
+  }
+
+  /// The destructor deletes all nodes.
+  ~Trie() {
+    for (auto iter = Nodes.rbegin(); iter != Nodes.rend(); ++iter) {
+      auto *node = *iter;
+      delete node;
+    }
+
+    Nodes.clear();
+  }
+
+  /// Insert an entry with the key given by the range [begin, end).
+  /// Returns the old value if the trie already had an entry for this key;
+  /// this is actually an invariant violation, but we can produce a better
+  /// assertion further up the stack.
+  template<typename Iter>
+  Optional<unsigned> insert(Iter begin, Iter end, unsigned ruleID) {
+    assert(begin != end);
+    auto *node = &Root;
+
+    while (true) {
+      auto &entry = node->Entries[*begin];
+      ++begin;
+
+      if (begin == end) {
+        if (entry.RuleID)
+          return entry.RuleID;
+
+        entry.RuleID = ruleID;
+        return None;
+      }
+
+      if (entry.Children == nullptr) {
+        entry.Children = new Node();
+        Nodes.push_back(entry.Children);
+      }
+
+      node = entry.Children;
+    }
+  }
+
+  /// Find the shortest prefix of the range given by [begin,end).
+  template<typename Iter>
+  Optional<unsigned>
+  find(Iter begin, Iter end) const {
+    assert(begin != end);
+    auto *node = &Root;
+
+    while (true) {
+      auto found = node->Entries.find(*begin);
+      ++begin;
+
+      if (found == node->Entries.end())
+        return None;
+
+      const auto &entry = found->second;
+      if (begin == end || entry.RuleID)
+        return entry.RuleID;
+
+      if (entry.Children == nullptr)
+        return None;
+
+      node = entry.Children;
+    }
+  }
+};
+
+}  // end namespace rewriting
+
+}  // end namespace swift
+
+#endif
diff --git a/test/Generics/unify_superclass_types_2.swift b/test/Generics/unify_superclass_types_2.swift
@@ -31,8 +31,8 @@ func unifySuperclassTest<T : P1 & P2>(_: T) {
 // CHECK-LABEL: Requirement machine for <τ_0_0 where τ_0_0 : P1, τ_0_0 : P2>
 // CHECK-NEXT: Rewrite system: {
 // CHECK:      - τ_0_0.[P1&P2:X].[superclass: Generic<τ_0_0, String, τ_0_1> with <τ_0_0.[P2:A2], τ_0_0.[P2:B2]>] => τ_0_0.[P1&P2:X]
-// CHECK-NEXT: - τ_0_0.[P1&P2:X].[superclass: Generic<Int, τ_0_0, τ_0_1> with <τ_0_0.[P1:A1], τ_0_0.[P1:B1]>] => τ_0_0.[P1&P2:X]
 // CHECK-NEXT: - τ_0_0.[P1&P2:X].[layout: _NativeClass] => τ_0_0.[P1&P2:X]
+// CHECK-NEXT: - τ_0_0.[P1&P2:X].[superclass: Generic<Int, τ_0_0, τ_0_1> with <τ_0_0.[P1:A1], τ_0_0.[P1:B1]>] => τ_0_0.[P1&P2:X]
 // CHECK-NEXT: - τ_0_0.[P2:A2].[concrete: Int] => τ_0_0.[P2:A2]
 // CHECK-NEXT: - τ_0_0.[P1:A1].[concrete: String] => τ_0_0.[P1:A1]
 // CHECK-NEXT: - τ_0_0.[P2:B2] => τ_0_0.[P1:B1]
diff --git a/test/Generics/unify_superclass_types_3.swift b/test/Generics/unify_superclass_types_3.swift
@@ -33,8 +33,8 @@ func unifySuperclassTest<T : P1 & P2>(_: T) {
 // CHECK-LABEL: Requirement machine for <τ_0_0 where τ_0_0 : P1, τ_0_0 : P2>
 // CHECK-NEXT: Rewrite system: {
 // CHECK:      - τ_0_0.[P1&P2:X].[superclass: Generic<τ_0_0, String, τ_0_1> with <τ_0_0.[P2:A2], τ_0_0.[P2:B2]>] => τ_0_0.[P1&P2:X]
-// CHECK-NEXT: - τ_0_0.[P1&P2:X].[superclass: Derived<τ_0_0, τ_0_1> with <τ_0_0.[P1:A1], τ_0_0.[P1:B1]>] => τ_0_0.[P1&P2:X]
 // CHECK-NEXT: - τ_0_0.[P1&P2:X].[layout: _NativeClass] => τ_0_0.[P1&P2:X]
+// CHECK-NEXT: - τ_0_0.[P1&P2:X].[superclass: Derived<τ_0_0, τ_0_1> with <τ_0_0.[P1:A1], τ_0_0.[P1:B1]>] => τ_0_0.[P1&P2:X]
 // CHECK-NEXT: - τ_0_0.[P2:A2].[concrete: Int] => τ_0_0.[P2:A2]
 // CHECK-NEXT: - τ_0_0.[P1:A1].[concrete: String] => τ_0_0.[P1:A1]
 // CHECK-NEXT: - τ_0_0.[P2:B2] => τ_0_0.[P1:B1]
diff --git a/test/Generics/unify_superclass_types_4.swift b/test/Generics/unify_superclass_types_4.swift
@@ -36,8 +36,8 @@ func unifySuperclassTest<T : P1 & P2>(_: T) {
 // CHECK-LABEL: Requirement machine for <τ_0_0 where τ_0_0 : P1, τ_0_0 : P2>
 // CHECK-NEXT: Rewrite system: {
 // CHECK:      - τ_0_0.[P1&P2:X].[superclass: Derived<τ_0_0> with <τ_0_0.[P2:A2]>] => τ_0_0.[P1&P2:X]
-// CHECK-NEXT: - τ_0_0.[P1&P2:X].[superclass: Base<τ_0_0> with <τ_0_0.[P1:A1]>] => τ_0_0.[P1&P2:X]
 // CHECK-NEXT: - τ_0_0.[P1&P2:X].[layout: _NativeClass] => τ_0_0.[P1&P2:X]
+// CHECK-NEXT: - τ_0_0.[P1&P2:X].[superclass: Base<τ_0_0> with <τ_0_0.[P1:A1]>] => τ_0_0.[P1&P2:X]
 // CHECK-NEXT: - τ_0_0.[P2:A2].[Q:T] => τ_0_0.[P1:A1]
 // CHECK-NEXT: }
 // CHECK-NEXT: Property map: {
