RequirementMachine: Simplify concrete substitutions when adding new rules

Author: slavapestov

include/swift/AST/RewriteSystem.h

@@ -181,6 +181,10 @@ class Atom final {
 
   int compare(Atom other, const ProtocolGraph &protos) const;
 
+  Atom transformConcreteSubstitutions(
+      llvm::function_ref<Term(Term)> fn,
+      RewriteContext &ctx) const;
+
   Atom prependPrefixToConcreteSubstitutions(
       const MutableTerm &prefix,
       RewriteContext &ctx) const;
@@ -246,6 +250,14 @@ class Term final {
   static Term get(const MutableTerm &term, RewriteContext &ctx);
 
   void dump(llvm::raw_ostream &out) const;
+
+  friend bool operator==(Term lhs, Term rhs) {
+    return lhs.Ptr == rhs.Ptr;
+  }
+
+  friend bool operator!=(Term lhs, Term rhs) {
+    return !(lhs == rhs);
+  }
 };
 
 /// A term is a sequence of one or more atoms.
@@ -499,6 +511,8 @@ class RewriteSystem final {
   void initialize(std::vector<std::pair<MutableTerm, MutableTerm>> &&rules,
                   ProtocolGraph &&protos);
 
+  Atom simplifySubstitutionsInSuperclassOrConcreteAtom(Atom atom) const;
+
   bool addRule(MutableTerm lhs, MutableTerm rhs);
 
   bool simplify(MutableTerm &term) const;
@@ -519,6 +533,8 @@ class RewriteSystem final {
       unsigned maxIterations,
       unsigned maxDepth);
 
+  void simplifyRightHandSides();
+
   void dump(llvm::raw_ostream &out) const;
 
 private:

lib/AST/RewriteSystem.cpp

@@ -493,30 +493,34 @@ int Atom::compare(Atom other, const ProtocolGraph &graph) const {
 ///   [concrete: Foo<X1, ..., Xn>]
 ///   [superclass: Foo<X1, ..., Xn>]
 ///
-/// Return a new atom where the prefix T is prepended to each of the
+/// Return a new atom where the function fn is applied to each of the
 /// substitutions:
 ///
-///   [concrete: Foo<T.X1, ..., T.Xn>]
-///   [superclass: Foo<T.X1, ..., T.Xn>]
+///   [concrete: Foo<fn(X1), ..., fn(Xn)>]
+///   [superclass: Foo<fn(X1), ..., fn(Xn)>]
 ///
 /// Asserts if this is not a superclass or concrete type atom.
-Atom Atom::prependPrefixToConcreteSubstitutions(
-    const MutableTerm &prefix,
+Atom Atom::transformConcreteSubstitutions(
+    llvm::function_ref<Term(Term)> fn,
     RewriteContext &ctx) const {
   assert(isSuperclassOrConcreteType());
 
-  if (prefix.empty())
+  if (getSubstitutions().empty())
     return *this;
 
+  bool anyChanged = false;
   SmallVector<Term, 2> substitutions;
   for (auto term : getSubstitutions()) {
-    MutableTerm mutTerm;
-    mutTerm.append(prefix);
-    mutTerm.append(term);
+    auto newTerm = fn(term);
+    if (newTerm != term)
+      anyChanged = true;
 
-    substitutions.push_back(Term::get(mutTerm, ctx));
+    substitutions.push_back(newTerm);
   }
 
+  if (!anyChanged)
+    return *this;
+
   switch (getKind()) {
   case Kind::Superclass:
     return Atom::forSuperclass(getSuperclass(), substitutions, ctx);
@@ -534,6 +538,34 @@ Atom Atom::prependPrefixToConcreteSubstitutions(
   llvm_unreachable("Bad atom kind");
 }
 
+/// 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 {
+  if (prefix.empty())
+    return *this;
+
+  return transformConcreteSubstitutions(
+    [&](Term term) -> Term {
+      MutableTerm mutTerm;
+      mutTerm.append(prefix);
+      mutTerm.append(term);
+
+      return Term::get(mutTerm, ctx);
+    }, ctx);
+}
+
 /// Print the atom using our mnemonic representation.
 void Atom::dump(llvm::raw_ostream &out) const {
   auto dumpSubstitutions = [&]() {
@@ -976,6 +1008,18 @@ void RewriteSystem::initialize(
     addRule(rule.first, rule.second);
 }
 
+Atom RewriteSystem::simplifySubstitutionsInSuperclassOrConcreteAtom(
+    Atom atom) const {
+  return atom.transformConcreteSubstitutions(
+    [&](Term term) -> Term {
+      MutableTerm mutTerm(term);
+      if (!simplify(mutTerm))
+        return term;
+
+      return Term::get(mutTerm, Context);
+    }, Context);
+}
+
 bool RewriteSystem::addRule(MutableTerm lhs, MutableTerm rhs) {
   // Simplify the rule as much as possible with the rules we have so far.
   //
@@ -994,6 +1038,9 @@ bool RewriteSystem::addRule(MutableTerm lhs, MutableTerm rhs) {
   if (result < 0)
     std::swap(lhs, rhs);
 
+  if (lhs.back().isSuperclassOrConcreteType())
+    lhs.back() = simplifySubstitutionsInSuperclassOrConcreteAtom(lhs.back());
+
   assert(lhs.compare(rhs, Protos) > 0);
 
   if (DebugAdd) {
@@ -1457,7 +1504,12 @@ RewriteSystem::computeConfluentCompletion(unsigned maxIterations,
     processMergedAssociatedTypes();
   }
 
-  // This isn't necessary for correctness, it's just an optimization.
+  simplifyRightHandSides();
+
+  return CompletionResult::Success;
+}
+
+void RewriteSystem::simplifyRightHandSides() {
   for (auto &rule : Rules) {
     if (rule.isDeleted())
       continue;
@@ -1466,8 +1518,6 @@ RewriteSystem::computeConfluentCompletion(unsigned maxIterations,
     simplify(rhs);
     rule = Rule(rule.getLHS(), rhs);
   }
-
-  return CompletionResult::Success;
 }
 
 void RewriteSystem::dump(llvm::raw_ostream &out) const {