RequirementMachine: Enforce that rewrite rules preserve the term's 'domain'

Intuitively, the first token of a rewrite rule determines if the rule
applies to a protocol requirement signature 'Self', or a generic
parameter in the top-level generic signature. A rewrite rule can never
take a type starting with the protocol 'Self' to a type starting with a
generic parameter, or vice versa.

Enforce this by defining the notion of a 'domain', which is the set of
protocols to which the first atom in a term applies to. This set can be
empty (if we have a generic parameter atom), or it may contain more than
one element (if we have an associated type atom for a merged associated
type).

Author: slavapestov

lib/AST/RequirementMachine/RewriteSystem.cpp

@@ -139,11 +139,16 @@ const ProtocolDecl *Atom::getProtocol() const {
   return Ptr->Proto;
 }
 
-/// Get the list of protocols associated with an associated type atom.
+/// Get the list of protocols associated with a protocol or associated type
+/// atom. Note that if this is a protocol atom, the return value will have
+/// exactly one element.
 ArrayRef<const ProtocolDecl *> Atom::getProtocols() const {
-  assert(getKind() == Kind::AssociatedType);
   auto protos = Ptr->getProtocols();
-  assert(!protos.empty());
+  if (protos.empty()) {
+    assert(getKind() == Kind::Protocol);
+    return {&Ptr->Proto, 1};
+  }
+  assert(getKind() == Kind::AssociatedType);
   return protos;
 }
 
@@ -747,6 +752,35 @@ void Term::Storage::Profile(llvm::FoldingSetNodeID &id) const {
     id.AddPointer(atom.getOpaquePointer());
 }
 
+/// Returns the "domain" of this term by looking at the first atom.
+///
+/// - If the first atom is a protocol atom [P], the domain is P.
+/// - If the first atom is an associated type atom [P1&...&Pn],
+///   the domain is {P1, ..., Pn}.
+/// - If the first atom is a generic parameter atom, the domain is
+///   the empty set {}.
+/// - Anything else will assert.
+ArrayRef<const ProtocolDecl *> MutableTerm::getRootProtocols() const {
+  auto atom = *begin();
+
+  switch (atom.getKind()) {
+  case Atom::Kind::Protocol:
+  case Atom::Kind::AssociatedType:
+    return atom.getProtocols();
+
+  case Atom::Kind::GenericParam:
+    return ArrayRef<const ProtocolDecl *>();
+
+  case Atom::Kind::Name:
+  case Atom::Kind::Layout:
+  case Atom::Kind::Superclass:
+  case Atom::Kind::ConcreteType:
+    break;
+  }
+
+  llvm_unreachable("Bad root atom");
+}
+
 /// Linear order on terms.
 ///
 /// First we compare length, then perform a lexicographic comparison
@@ -1253,7 +1287,8 @@ void RewriteSystem::simplifyRightHandSides() {
 
 #define ASSERT_RULE(expr) \
   if (!(expr)) { \
-    llvm::errs() << "&&& Malformed rewrite rule: " << rule << "\n\n"; \
+    llvm::errs() << "&&& Malformed rewrite rule: " << rule << "\n"; \
+    llvm::errs() << "&&& " << #expr << "\n\n"; \
     dump(llvm::errs()); \
     assert(expr); \
   }
@@ -1298,6 +1333,11 @@ void RewriteSystem::simplifyRightHandSides() {
         ASSERT_RULE(atom.getKind() != Atom::Kind::Protocol);
       }
     }
+
+    auto lhsDomain = lhs.getRootProtocols();
+    auto rhsDomain = rhs.getRootProtocols();
+
+    ASSERT_RULE(lhsDomain == rhsDomain);
   }
 
 #undef ASSERT_RULE

lib/AST/RequirementMachine/RewriteSystem.h

@@ -340,6 +340,8 @@ class MutableTerm final {
 
   size_t size() const { return Atoms.size(); }
 
+  ArrayRef<const ProtocolDecl *> getRootProtocols() const;
+
   decltype(Atoms)::const_iterator begin() const { return Atoms.begin(); }
   decltype(Atoms)::const_iterator end() const { return Atoms.end(); }