RequirementMachine: Parametrize Trie by value type and add support for longest-suffix match

The PropertyMap wants to use a try to map to PropertyBags, and it
needs longest-suffix rather than shortest-prefix matching.

Implement both by making Trie into a template class with two
parameters; the ValueType, and the MatchKind.

Note that while the MatchKind encodes the longest vs shortest
match part, matching on the prefix vs suffix of a term is up
to the caller, since the find() and insert() methods of Trie
take a pair of iterators, so simply passing in begin()/end() vs
rbegin()/rend() selects the direction.

Author: slavapestov

lib/AST/RequirementMachine/RewriteSystem.cpp

@@ -32,8 +32,8 @@ RewriteSystem::RewriteSystem(RewriteContext &ctx)
 }
 
 RewriteSystem::~RewriteSystem() {
-  RuleTrie.updateHistograms(Context.RuleTrieHistogram,
-                            Context.RuleTrieRootHistogram);
+  Trie.updateHistograms(Context.RuleTrieHistogram,
+                        Context.RuleTrieRootHistogram);
 }
 
 void Rule::dump(llvm::raw_ostream &out) const {
@@ -99,7 +99,7 @@ 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);
+  auto oldRuleID = Trie.insert(lhs.begin(), lhs.end(), i);
   if (oldRuleID) {
     llvm::errs() << "Duplicate rewrite rule!\n";
     const auto &oldRule = Rules[*oldRuleID];
@@ -167,7 +167,7 @@ bool RewriteSystem::simplify(MutableTerm &term) const {
     auto from = term.begin();
     auto end = term.end();
     while (from < end) {
-      auto ruleID = RuleTrie.find(from, end);
+      auto ruleID = Trie.find(from, end);
       if (ruleID) {
         const auto &rule = Rules[*ruleID];
         if (!rule.isDeleted()) {

lib/AST/RequirementMachine/RewriteSystem.h

@@ -98,8 +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;
+  /// A prefix trie of rule left hand sides to optimize lookup. The value
+  /// type is an index into the Rules array defined above.
+  Trie<unsigned, MatchKind::Shortest> Trie;
 
   /// The graph of all protocols transitively referenced via our set of
   /// rewrite rules, used for the linear order on symbols.

lib/AST/RequirementMachine/Trie.h

@@ -20,12 +20,18 @@ namespace swift {
 
 namespace rewriting {
 
+enum class MatchKind {
+  Shortest,
+  Longest
+};
+
+template<typename ValueType, MatchKind Kind>
 class Trie {
 public:
   struct Node;
 
   struct Entry {
-    Optional<unsigned> RuleID;
+    Optional<ValueType> Value;
     Node *Children = nullptr;
   };
 
@@ -47,8 +53,10 @@ class Trie {
     rootStats.add(Root.Entries.size());
   }
 
-  /// The destructor deletes all nodes.
-  ~Trie() {
+  /// Delete all entries from the trie.
+  void clear() {
+    Root.Entries.clear();
+
     for (auto iter = Nodes.rbegin(); iter != Nodes.rend(); ++iter) {
       auto *node = *iter;
       delete node;
@@ -57,12 +65,16 @@ class Trie {
     Nodes.clear();
   }
 
+  ~Trie() {
+    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) {
+  Optional<ValueType> insert(Iter begin, Iter end, ValueType value) {
     assert(begin != end);
     auto *node = &Root;
 
@@ -71,10 +83,10 @@ class Trie {
       ++begin;
 
       if (begin == end) {
-        if (entry.RuleID)
-          return entry.RuleID;
+        if (entry.Value)
+          return entry.Value;
 
-        entry.RuleID = ruleID;
+        entry.Value = value;
         return None;
       }
 
@@ -87,26 +99,38 @@ class Trie {
     }
   }
 
-  /// Find the shortest prefix of the range given by [begin,end).
+  /// Find the shortest or longest prefix of the range given by [begin,end),
+  /// depending on whether the Kind template parameter was bound to
+  /// MatchKind::Shortest or MatchKind::Longest.
   template<typename Iter>
-  Optional<unsigned>
+  Optional<ValueType>
   find(Iter begin, Iter end) const {
     assert(begin != end);
     auto *node = &Root;
 
+    Optional<ValueType> bestMatch = None;
+
     while (true) {
       auto found = node->Entries.find(*begin);
       ++begin;
 
       if (found == node->Entries.end())
-        return None;
+        return bestMatch;
 
       const auto &entry = found->second;
-      if (begin == end || entry.RuleID)
-        return entry.RuleID;
+
+      if (entry.Value) {
+        if (Kind == MatchKind::Shortest)
+          return entry.Value;
+
+        bestMatch = entry.Value;
+      }
+
+      if (begin == end)
+        return bestMatch;
 
       if (entry.Children == nullptr)
-        return None;
+        return bestMatch;
 
       node = entry.Children;
     }