RequirementMachine: Fix up concrete type atoms when computing overlaps

Author: slavapestov

include/swift/AST/RewriteSystem.h

@@ -125,6 +125,11 @@ class Atom final {
 public:
   Kind getKind() const;
 
+  bool isSuperclassOrConcreteType() const {
+    auto kind = getKind();
+    return (kind == Kind::Superclass || kind == Kind::ConcreteType);
+  }
+
   Identifier getName() const;
 
   const ProtocolDecl *getProtocol() const;
@@ -176,6 +181,10 @@ class Atom final {
 
   int compare(Atom other, const ProtocolGraph &protos) const;
 
+  Atom prependPrefixToConcreteSubstitutions(
+      const MutableTerm &prefix,
+      RewriteContext &ctx) const;
+
   void dump(llvm::raw_ostream &out) const;
 
   friend bool operator==(Atom lhs, Atom rhs) {
@@ -276,6 +285,10 @@ class MutableTerm final {
     Atoms.push_back(atom);
   }
 
+  void append(Term other) {
+    Atoms.append(other.begin(), other.end());
+  }
+
   void append(const MutableTerm &other) {
     Atoms.append(other.begin(), other.end());
   }

lib/AST/RewriteSystem.cpp

@@ -488,6 +488,52 @@ int Atom::compare(Atom other, const ProtocolGraph &graph) const {
   return result;
 }
 
+/// For a superclass or concrete type atom
+///
+///   [concrete: Foo<X1, ..., Xn>]
+///   [superclass: Foo<X1, ..., Xn>]
+///
+/// Return a new atom where the prefix T is prepended to each of the
+/// substitutions:
+///
+///   [concrete: Foo<T.X1, ..., T.Xn>]
+///   [superclass: Foo<T.X1, ..., T.Xn>]
+///
+/// Asserts if this is not a superclass or concrete type atom.
+Atom Atom::prependPrefixToConcreteSubstitutions(
+    const MutableTerm &prefix,
+    RewriteContext &ctx) const {
+  assert(isSuperclassOrConcreteType());
+
+  if (prefix.empty())
+    return *this;
+
+  SmallVector<Term, 2> substitutions;
+  for (auto term : getSubstitutions()) {
+    MutableTerm mutTerm;
+    mutTerm.append(prefix);
+    mutTerm.append(term);
+
+    substitutions.push_back(Term::get(mutTerm, ctx));
+  }
+
+  switch (getKind()) {
+  case Kind::Superclass:
+    return Atom::forSuperclass(getSuperclass(), substitutions, ctx);
+  case Kind::ConcreteType:
+    return Atom::forConcreteType(getConcreteType(), substitutions, ctx);
+
+  case Kind::GenericParam:
+  case Kind::Name:
+  case Kind::Protocol:
+  case Kind::AssociatedType:
+  case Kind::Layout:
+    break;
+  }
+
+  llvm_unreachable("Bad atom kind");
+}
+
 /// Print the atom using our mnemonic representation.
 void Atom::dump(llvm::raw_ostream &out) const {
   auto dumpSubstitutions = [&]() {
@@ -1286,6 +1332,11 @@ RewriteSystem::computeCriticalPair(const Rule &lhs, const Rule &rhs) const {
   case OverlapKind::Second: {
     // lhs == TU -> X, rhs == UV -> Y.
 
+    if (v.back().isSuperclassOrConcreteType()) {
+      v.back() = v.back().prependPrefixToConcreteSubstitutions(
+          t, Context);
+    }
+
     // Compute the term XV.
     MutableTerm xv;
     xv.append(lhs.getRHS());