RequirementMachine: Move the protocol graph into the rewrite system

Author: slavapestov

include/swift/AST/ProtocolGraph.h

@@ -77,6 +77,9 @@ struct ProtocolGraph {
 
   void computeInheritedAssociatedTypes();
 
+  int compareProtocols(const ProtocolDecl *lhs,
+                       const ProtocolDecl *rhs) const;
+
 private:
   unsigned computeProtocolDepth(const ProtocolDecl *proto);
 };

include/swift/AST/RewriteSystem.h

@@ -16,6 +16,7 @@
 #include "swift/AST/Decl.h"
 #include "swift/AST/Identifier.h"
 #include "swift/AST/LayoutConstraint.h"
+#include "swift/AST/ProtocolGraph.h"
 #include "swift/AST/Types.h"
 #include "llvm/ADT/PointerUnion.h"
 #include "llvm/ADT/SmallVector.h"
@@ -30,9 +31,6 @@ namespace swift {
 
 namespace rewriting {
 
-using ProtocolOrder = std::function<int (const ProtocolDecl *,
-                                         const ProtocolDecl *)>;
-
 class Atom final {
   using Storage = llvm::PointerUnion<Identifier,
                                      GenericTypeParamType *,
@@ -142,7 +140,7 @@ class Atom final {
     return Value.get<LayoutConstraint>();
   }
 
-  int compare(Atom other, ProtocolOrder compare) const;
+  int compare(Atom other, const ProtocolGraph &protos) const;
 
   void dump(llvm::raw_ostream &out) const;
 
@@ -174,7 +172,7 @@ class Term final {
     Atoms.push_back(atom);
   }
 
-  int compare(const Term &other, ProtocolOrder order) const;
+  int compare(const Term &other, const ProtocolGraph &protos) const;
 
   size_t size() const { return Atoms.size(); }
 
@@ -239,22 +237,22 @@ class Rule final {
   }
 
   int compare(const Rule &other,
-              ProtocolOrder protocolOrder) const {
-    return LHS.compare(other.LHS, protocolOrder);
+              const ProtocolGraph &protos) const {
+    return LHS.compare(other.LHS, protos);
   }
 
   void dump(llvm::raw_ostream &out) const;
 };
 
 class RewriteSystem final {
   std::vector<Rule> Rules;
-  ProtocolOrder Order;
+  ProtocolGraph Protos;
 
   unsigned DebugSimplify : 1;
   unsigned DebugAdd : 1;
 
 public:
-  explicit RewriteSystem(ProtocolOrder order) : Order(order) {
+  explicit RewriteSystem() {
     DebugSimplify = false;
     DebugAdd = false;
   }
@@ -264,6 +262,11 @@ class RewriteSystem final {
   RewriteSystem &operator=(const RewriteSystem &) = delete;
   RewriteSystem &operator=(RewriteSystem &&) = delete;
 
+  const ProtocolGraph &getProtocols() const { return Protos; }
+
+  void initialize(std::vector<std::pair<Term, Term>> &&rules,
+                  ProtocolGraph &&protos);
+
   bool addRule(Term lhs, Term rhs);
 
   bool simplify(Term &term) const;

lib/AST/ProtocolGraph.cpp

@@ -130,4 +130,12 @@ unsigned ProtocolGraph::computeProtocolDepth(const ProtocolDecl *proto) {
 
   info.Depth = depth;
   return depth;
-}
+}
+
+int ProtocolGraph::compareProtocols(const ProtocolDecl *lhs,
+                                    const ProtocolDecl *rhs) const {
+  const auto &infoLHS = getProtocolInfo(lhs);
+  const auto &infoRHS = getProtocolInfo(rhs);
+
+  return infoLHS.Index - infoRHS.Index;
+}

lib/AST/RequirementMachine.cpp

@@ -174,20 +174,10 @@ Term swift::rewriting::getTermForType(CanType paramType,
 }
 
 struct RequirementMachine::Implementation {
-  ProtocolGraph Protocols;
-  ProtocolOrder Order;
   RewriteSystem System;
   bool Complete = false;
 
-  Implementation()
-      : Order([&](const ProtocolDecl *lhs,
-                  const ProtocolDecl *rhs) -> int {
-          const auto &infoLHS = Protocols.getProtocolInfo(lhs);
-          const auto &infoRHS = Protocols.getProtocolInfo(rhs);
-
-          return infoLHS.Index - infoRHS.Index;
-        }),
-        System(Order) {}
+  Implementation() {}
 };
 
 RequirementMachine::RequirementMachine(ASTContext &ctx) : Context(ctx) {
@@ -212,15 +202,8 @@ void RequirementMachine::addGenericSignature(CanGenericSignature sig) {
   RewriteSystemBuilder builder(Context);
   builder.addGenericSignature(sig);
 
-  Impl->Protocols = builder.Protocols;
-
-  std::sort(builder.Rules.begin(), builder.Rules.end(),
-            [&](std::pair<Term, Term> lhs,
-                std::pair<Term, Term> rhs) -> int {
-              return lhs.first.compare(rhs.first, Impl->Order) < 0;
-            });
-  for (const auto &rule : builder.Rules)
-    Impl->System.addRule(rule.first, rule.second);
+  Impl->System.initialize(std::move(builder.Rules),
+                          std::move(builder.Protocols));
 
   // FIXME: Add command line flag
   auto result = Impl->System.computeConfluentCompletion(1000, 10);

lib/AST/RewriteSystem.cpp

@@ -19,7 +19,7 @@
 using namespace swift;
 using namespace rewriting;
 
-int Atom::compare(Atom other, ProtocolOrder protocolOrder) const {
+int Atom::compare(Atom other, const ProtocolGraph &graph) const {
   auto kind = getKind();
   auto otherKind = other.getKind();
 
@@ -31,17 +31,18 @@ int Atom::compare(Atom other, ProtocolOrder protocolOrder) const {
     return getName().compare(other.getName());
 
   case Kind::Protocol:
-    return protocolOrder(getProtocol(), other.getProtocol());
+    return graph.compareProtocols(getProtocol(), other.getProtocol());
 
   case Kind::AssociatedType: {
     auto protos = getProtocols();
     auto otherProtos = other.getProtocols();
 
+    // Atoms with more protocols are 'smaller' than those with fewer.
     if (protos.size() != otherProtos.size())
-      return otherProtos.size() > protos.size() ? -1 : 1;
+      return protos.size() > otherProtos.size() ? -1 : 1;
 
     for (unsigned i : indices(protos)) {
-      int result = protocolOrder(protos[i], otherProtos[i]);
+      int result = graph.compareProtocols(protos[i], otherProtos[i]);
       if (result)
         return result;
     }
@@ -108,15 +109,15 @@ void Atom::dump(llvm::raw_ostream &out) const {
   llvm_unreachable("Bad atom kind");
 }
 
-int Term::compare(const Term &other, ProtocolOrder protocolOrder) const {
+int Term::compare(const Term &other, const ProtocolGraph &graph) const {
   if (size() != other.size())
     return size() < other.size() ? -1 : 1;
 
   for (unsigned i = 0, e = size(); i < e; ++i) {
     auto lhs = (*this)[i];
     auto rhs = other[i];
 
-    int result = lhs.compare(rhs, protocolOrder);
+    int result = lhs.compare(rhs, graph);
     if (result != 0) {
       assert(lhs != rhs);
       return result;
@@ -222,17 +223,30 @@ void Rule::dump(llvm::raw_ostream &out) const {
     out << " [deleted]";
 }
 
+void RewriteSystem::initialize(std::vector<std::pair<Term, Term>> &&rules,
+                               ProtocolGraph &&graph) {
+  Protos = graph;
+
+  std::sort(rules.begin(), rules.end(),
+            [&](std::pair<Term, Term> lhs,
+                std::pair<Term, Term> rhs) -> int {
+              return lhs.first.compare(rhs.first, graph) < 0;
+            });
+  for (const auto &rule : rules)
+    addRule(rule.first, rule.second);
+}
+
 bool RewriteSystem::addRule(Term lhs, Term rhs) {
   simplify(lhs);
   simplify(rhs);
 
-  int result = lhs.compare(rhs, Order);
+  int result = lhs.compare(rhs, Protos);
   if (result == 0)
     return false;
   if (result < 0)
     std::swap(lhs, rhs);
 
-  assert(lhs.compare(rhs, Order) > 0);
+  assert(lhs.compare(rhs, Protos) > 0);
 
   if (DebugAdd) {
     llvm::dbgs() << "# Adding rule ";