RequirementMachine: Topological order for protocol declarations

Author: slavapestov

include/swift/AST/RewriteSystem.h

@@ -27,6 +27,9 @@ namespace swift {
 
 namespace rewriting {
 
+using ProtocolOrder = std::function<int (const ProtocolDecl *,
+                                         const ProtocolDecl *)>;
+
 class Atom final {
   using Storage = llvm::PointerUnion<Identifier, GenericTypeParamType *>;
 
@@ -103,7 +106,7 @@ class Atom final {
     return Value.get<GenericTypeParamType *>();
   }
 
-  int compare(Atom other) const;
+  int compare(Atom other, ProtocolOrder compare) const;
 
   void dump(llvm::raw_ostream &out) const;
 
@@ -129,7 +132,7 @@ class Term final {
     Atoms.push_back(atom);
   }
 
-  int compare(const Term &other) const;
+  int compare(const Term &other, ProtocolOrder order) const;
 
   size_t size() const { return Atoms.size(); }
 
@@ -165,9 +168,7 @@ class Rule final {
 
 public:
   Rule(const Term &lhs, const Term &rhs)
-      : LHS(lhs), RHS(rhs), deleted(false) {
-    assert(LHS.compare(RHS) > 0);
-  }
+      : LHS(lhs), RHS(rhs), deleted(false) {}
 
   bool apply(Term &term) const {
     assert(!deleted);
@@ -196,9 +197,17 @@ class Rule final {
 
 class RewriteSystem final {
   std::vector<Rule> Rules;
+  ProtocolOrder Order;
   bool Debug = false;
 
 public:
+  explicit RewriteSystem(ProtocolOrder order) : Order(order) {}
+
+  RewriteSystem(const RewriteSystem &) = delete;
+  RewriteSystem(RewriteSystem &&) = delete;
+  RewriteSystem &operator=(const RewriteSystem &) = delete;
+  RewriteSystem &operator=(RewriteSystem &&) = delete;
+
   bool addRule(Term lhs, Term rhs);
 
   bool simplify(Term &term) const;

lib/AST/RequirementMachine.cpp

@@ -72,6 +72,13 @@ struct ProtocolGraph {
     }
   }
 
+  const ProtocolInfo &getProtocolInfo(
+      const ProtocolDecl *proto) const {
+    auto found = Info.find(proto);
+    assert(found != Info.end());
+    return found->second;
+  }
+
   void addProtocol(const ProtocolDecl *proto) {
     if (Info.count(proto) > 0)
       return;
@@ -86,7 +93,7 @@ struct ProtocolGraph {
     unsigned i = 0;
     while (i < Protocols.size()) {
       auto *proto = Protocols[i++];
-      visitRequirements(Info[proto].Requirements);
+      visitRequirements(getProtocolInfo(proto).Requirements);
     }
   }
 
@@ -99,8 +106,8 @@ struct ProtocolGraph {
         Protocols.begin(), Protocols.end(),
         [&](const ProtocolDecl *lhs,
             const ProtocolDecl *rhs) -> int {
-          const auto &lhsInfo = Info[lhs];
-          const auto &rhsInfo = Info[rhs];
+          const auto &lhsInfo = getProtocolInfo(lhs);
+          const auto &rhsInfo = getProtocolInfo(rhs);
 
           // protocol Base {} // depth 1
           // protocol Derived : Base {} // depth 2
@@ -126,7 +133,7 @@ struct ProtocolGraph {
         if (inherited == proto)
           continue;
 
-        for (auto *inheritedType : Info[inherited].AssociatedTypes) {
+        for (auto *inheritedType : getProtocolInfo(inherited).AssociatedTypes) {
           if (!visited.insert(inheritedType).second)
             continue;
 
@@ -164,12 +171,12 @@ struct ProtocolGraph {
   }
 };
 
-class RewriteSystemBuilder {
+struct RewriteSystemBuilder {
   ASTContext &Context;
 
+  ProtocolGraph Protocols;
   std::vector<std::pair<Term, Term>> Rules;
 
-public:
   RewriteSystemBuilder(ASTContext &ctx) : Context(ctx) {}
   void addGenericSignature(CanGenericSignature sig);
   void addAssociatedType(const AssociatedTypeDecl *type,
@@ -179,31 +186,25 @@ class RewriteSystemBuilder {
                                   const ProtocolDecl *proto);
   void addRequirement(const Requirement &req,
                       const ProtocolDecl *proto);
-
-  void addRulesToRewriteSystem(RewriteSystem &system);
 };
 
 } // end namespace
 
 void RewriteSystemBuilder::addGenericSignature(CanGenericSignature sig) {
-  ProtocolGraph graph;
-  graph.visitRequirements(sig->getRequirements());
-  graph.computeTransitiveClosure();
-  graph.computeLinearOrder();
-  graph.computeInheritedAssociatedTypes();
-
-  for (auto *proto : graph.Protocols) {
-    if (Context.LangOpts.DebugRequirementMachine) {
-      llvm::dbgs() << "protocol " << proto->getName() << " {\n";
-    }
+  Protocols.visitRequirements(sig->getRequirements());
+  Protocols.computeTransitiveClosure();
+  Protocols.computeLinearOrder();
+  Protocols.computeInheritedAssociatedTypes();
 
-    const auto &info = graph.Info[proto];
+  for (auto *proto : Protocols.Protocols) {
+    const auto &info = Protocols.getProtocolInfo(proto);
 
     for (auto *type : info.AssociatedTypes)
       addAssociatedType(type, proto);
 
     for (auto *inherited : info.Inherited) {
-      for (auto *inheritedType : graph.Info[inherited].AssociatedTypes) {
+      auto inheritedTypes = Protocols.getProtocolInfo(inherited).AssociatedTypes;
+      for (auto *inheritedType : inheritedTypes) {
         addInheritedAssociatedType(inheritedType, inherited, proto);
       }
     }
@@ -286,12 +287,6 @@ void RewriteSystemBuilder::addRequirement(const Requirement &req,
   }
 }
 
-void RewriteSystemBuilder::addRulesToRewriteSystem(RewriteSystem &system) {
-  for (auto rule : Rules) {
-    system.addRule(rule.first, rule.second);
-  }
-}
-
 Term swift::rewriting::getTermForType(CanType paramType,
                                       const ProtocolDecl *proto) {
   assert(paramType->isTypeParameter());
@@ -314,8 +309,22 @@ 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 {
+          auto infoLHS = Protocols.Info.find(lhs);
+          assert(infoLHS != Protocols.Info.end());
+          auto infoRHS = Protocols.Info.find(rhs);
+          assert(infoRHS != Protocols.Info.end());
+
+          return infoRHS->second.Index - infoLHS->second.Index;
+        }),
+        System(Order) {}
 };
 
 RequirementMachine::RequirementMachine(ASTContext &ctx) : Context(ctx) {
@@ -336,10 +345,13 @@ void RequirementMachine::addGenericSignature(CanGenericSignature sig) {
   RewriteSystemBuilder builder(Context);
   builder.addGenericSignature(sig);
 
-  builder.addRulesToRewriteSystem(Impl->System);
+  Impl->Protocols = builder.Protocols;
+
+  for (const auto &rule : builder.Rules)
+    Impl->System.addRule(rule.first, rule.second);
 
   // FIXME: Add command line flag
-  Impl->System.computeConfluentCompletion(1000);
+  Impl->System.computeConfluentCompletion(10000);
 
   markComplete();
 

lib/AST/RewriteSystem.cpp

@@ -18,7 +18,7 @@
 using namespace swift;
 using namespace rewriting;
 
-int Atom::compare(Atom other) const {
+int Atom::compare(Atom other, ProtocolOrder protocolOrder) const {
   auto kind = getKind();
   auto otherKind = other.getKind();
 
@@ -30,10 +30,10 @@ int Atom::compare(Atom other) const {
     return getName().compare(other.getName());
 
   case Kind::Protocol:
-    return TypeDecl::compare(getProtocol(), other.getProtocol());
+    return protocolOrder(getProtocol(), other.getProtocol());
 
   case Kind::AssociatedType: {
-    int result = TypeDecl::compare(getProtocol(), other.getProtocol());
+    int result = protocolOrder(getProtocol(), other.getProtocol());
     if (result)
       return result;
 
@@ -81,15 +81,15 @@ void Atom::dump(llvm::raw_ostream &out) const {
   llvm_unreachable("Bad atom kind");
 }
 
-int Term::compare(const Term &other) const {
+int Term::compare(const Term &other, ProtocolOrder protocolOrder) 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);
+    int result = lhs.compare(rhs, protocolOrder);
     if (result != 0)
       return result;
   }
@@ -120,7 +120,6 @@ bool Term::rewriteSubTerm(const Term &lhs, const Term &rhs) {
 
   auto oldSize = size();
 
-  assert(rhs.compare(lhs) < 0);
   assert(rhs.size() <= lhs.size());
 
   auto newIter = std::copy(rhs.begin(), rhs.end(), found);
@@ -196,7 +195,7 @@ bool RewriteSystem::addRule(Term lhs, Term rhs) {
   simplify(lhs);
   simplify(rhs);
 
-  int result = lhs.compare(rhs);
+  int result = lhs.compare(rhs, Order);
   if (result == 0)
     return false;
   if (result < 0)