RequirementMachine: Fully canonicalize substitutions in concrete type rules

Previously we did this when adding new concrete type rules,
but we don't have a complete rewrite system at that point yet,
so there was no guarantee concrete substitution terms would
be canonical.

Now, perform simplification in a post-pass after completion,
at the same time as simplifying rule right hand sides.

Rewrite loops are recorded relating the original rule with the
simplified substitutions.

Author: slavapestov

lib/AST/RequirementMachine/KnuthBendix.cpp

@@ -596,7 +596,7 @@ RewriteSystem::computeConfluentCompletion(unsigned maxIterations,
       }
     }
 
-    simplifyRewriteSystem();
+    simplifyLeftHandSides();
 
     assert(resolvedCriticalPairs.size() == resolvedPaths.size());
 
@@ -629,6 +629,8 @@ RewriteSystem::computeConfluentCompletion(unsigned maxIterations,
     resolvedPaths.clear();
     resolvedLoops.clear();
 
+    simplifyRightHandSidesAndSubstitutions();
+
     // If the added rules merged any associated types, process the merges now
     // before we continue with the completion procedure. This is important
     // to perform incrementally since merging is required to repair confluence

lib/AST/RequirementMachine/RewriteSystem.cpp

@@ -251,14 +251,13 @@ bool RewriteSystem::simplify(MutableTerm &term, RewritePath *path) const {
 
 /// Simplify terms appearing in the substitutions of the last symbol of \p term,
 /// which must be a superclass or concrete type symbol.
-void RewriteSystem::simplifySubstitutions(MutableTerm &term,
+bool RewriteSystem::simplifySubstitutions(Symbol &symbol,
                                           RewritePath &path) const {
-  auto symbol = term.back();
   assert(symbol.hasSubstitutions());
 
   auto substitutions = symbol.getSubstitutions();
   if (substitutions.empty())
-    return;
+    return false;
 
   // Save the original rewrite path length so that we can reset if if we don't
   // find anything to simplify.
@@ -307,11 +306,12 @@ void RewriteSystem::simplifySubstitutions(MutableTerm &term,
 #endif
 
     path.resize(oldSize);
-    return;
+    return false;
   }
 
   // Build the new symbol with simplified substitutions.
-  term.back() = symbol.withConcreteSubstitutions(newSubstitutions, Context);
+  symbol = symbol.withConcreteSubstitutions(newSubstitutions, Context);
+  return true;
 }
 
 /// Adds a rewrite rule, returning true if the new rule was non-trivial.
@@ -341,12 +341,6 @@ bool RewriteSystem::addRule(MutableTerm lhs, MutableTerm rhs,
   RewritePath lhsPath;
   RewritePath rhsPath;
 
-  if (lhs.back().hasSubstitutions())
-    simplifySubstitutions(lhs, lhsPath);
-
-  if (rhs.back().hasSubstitutions())
-    simplifySubstitutions(rhs, rhsPath);
-
   simplify(lhs, &lhsPath);
   simplify(rhs, &rhsPath);
 
@@ -460,13 +454,11 @@ bool RewriteSystem::addExplicitRule(MutableTerm lhs, MutableTerm rhs) {
   return added;
 }
 
-/// Delete any rules whose left hand sides can be reduced by other rules,
-/// and reduce the right hand sides of all remaining rules as much as
-/// possible.
+/// Delete any rules whose left hand sides can be reduced by other rules.
 ///
 /// Must be run after the completion procedure, since the deletion of
 /// rules is only valid to perform if the rewrite system is confluent.
-void RewriteSystem::simplifyRewriteSystem() {
+void RewriteSystem::simplifyLeftHandSides() {
   assert(Complete);
 
   for (unsigned ruleID = 0, e = Rules.size(); ruleID < e; ++ruleID) {
@@ -500,8 +492,19 @@ void RewriteSystem::simplifyRewriteSystem() {
         break;
       }
     }
+  }
+}
+
+/// Reduce the right hand sides of all remaining rules as much as
+/// possible.
+///
+/// Must be run after the completion procedure, since the deletion of
+/// rules is only valid to perform if the rewrite system is confluent.
+void RewriteSystem::simplifyRightHandSidesAndSubstitutions() {
+  assert(Complete);
 
-    // If the rule was deleted above, skip the rest.
+  for (unsigned ruleID = 0, e = Rules.size(); ruleID < e; ++ruleID) {
+    auto &rule = getRule(ruleID);
     if (rule.isSimplified())
       continue;
 
@@ -511,6 +514,8 @@ void RewriteSystem::simplifyRewriteSystem() {
     if (!simplify(rhs, &rhsPath))
       continue;
 
+    auto lhs = rule.getLHS();
+
     // We're adding a new rule, so the old rule won't apply anymore.
     rule.markSimplified();
 
@@ -544,6 +549,41 @@ void RewriteSystem::simplifyRewriteSystem() {
 
     recordRewriteLoop(MutableTerm(lhs), loop);
   }
+
+  // Finally try to simplify substitutions in superclass, concrete type and
+  // concrete conformance symbols.
+  for (unsigned ruleID = 0, e = Rules.size(); ruleID < e; ++ruleID) {
+    auto &rule = getRule(ruleID);
+    if (rule.isSimplified())
+      continue;
+
+    auto lhs = rule.getLHS();
+    auto symbol = lhs.back();
+    if (!symbol.hasSubstitutions())
+      continue;
+
+    RewritePath path;
+
+    // (1) First, apply the original rule to produce the original lhs.
+    path.add(RewriteStep::forRewriteRule(/*startOffset=*/0, /*endOffset=*/0,
+                                         ruleID, /*inverse=*/true));
+
+    // (2) Now, simplify the substitutions to get the new lhs.
+    if (!simplifySubstitutions(symbol, path))
+      continue;
+
+    // We're either going to add a new rule or record an identity, so
+    // mark the old rule as simplified.
+    rule.markSimplified();
+
+    MutableTerm newLHS(lhs.begin(), lhs.end() - 1);
+    newLHS.add(symbol);
+
+    // Invert the path to get a path from the new lhs to the old rhs.
+    path.invert();
+
+    addRule(newLHS, MutableTerm(rule.getRHS()), &path);
+  }
 }
 
 void RewriteSystem::verifyRewriteRules(ValidityPolicy policy) const {

lib/AST/RequirementMachine/RewriteSystem.h

@@ -278,7 +278,7 @@ class RewriteSystem final {
 
   bool simplify(MutableTerm &term, RewritePath *path=nullptr) const;
 
-  void simplifySubstitutions(MutableTerm &term, RewritePath &path) const;
+  bool simplifySubstitutions(Symbol &symbol, RewritePath &path) const;
 
   //////////////////////////////////////////////////////////////////////////////
   ///
@@ -290,7 +290,9 @@ class RewriteSystem final {
   computeConfluentCompletion(unsigned maxIterations,
                              unsigned maxDepth);
 
-  void simplifyRewriteSystem();
+  void simplifyLeftHandSides();
+
+  void simplifyRightHandSidesAndSubstitutions();
 
   enum ValidityPolicy {
     AllowInvalidRequirements,

test/Generics/canonical_concrete_substitutions_in_protocol.swift

@@ -0,0 +1,23 @@
+// RUN: %target-swift-frontend -typecheck %s -debug-generic-signatures -requirement-machine-protocol-signatures=on 2>&1 | %FileCheck %s
+
+struct G<T> {}
+
+// CHECK-LABEL: canonical_concrete_substitutions_in_protocol.(file).P@
+// CHECK-NEXT: Requirement signature: <Self where Self.A == G<Self.B>, Self.B == Self.T.X, Self.T : Q>
+
+protocol P {
+  associatedtype A where A == G<T.X>
+  associatedtype B where B == T.X
+  associatedtype T : Q
+}
+
+protocol Q {
+  associatedtype X
+}
+
+protocol QQ : Q {}
+
+protocol R {
+  associatedtype A
+  associatedtype C: QQ where C.X == G<A>
+}

test/Generics/unify_concrete_types_1.swift

@@ -28,5 +28,5 @@ struct MergeTest<G : P1 & P2> {
 // CHECK:  [P1:X] => { concrete_type: [concrete: Foo<τ_0_0, τ_0_1> with <[P1:Y1], [P1:Z1]>] }
 // CHECK:  [P2:X] => { concrete_type: [concrete: Foo<τ_0_0, τ_0_1> with <[P2:Y2], [P2:Z2]>] }
 // CHECK:  τ_0_0 => { conforms_to: [P1 P2] }
-// CHECK:  τ_0_0.[P1&P2:X] => { concrete_type: [concrete: Foo<τ_0_0, τ_0_1> with <τ_0_0.[P2:Y2], τ_0_0.[P2:Z2]>] }
+// CHECK:  τ_0_0.[P1&P2:X] => { concrete_type: [concrete: Foo<τ_0_0, τ_0_1> with <τ_0_0.[P1:Y1], τ_0_0.[P1:Z1]>] }
 // CHECK: }

test/Generics/unify_superclass_types_2.swift

@@ -33,6 +33,8 @@ func unifySuperclassTest<T : P1 & P2>(_: T) {
 // CHECK:      - τ_0_0.[P1&P2:X].[superclass: Generic<τ_0_0, String, τ_0_1> with <τ_0_0.[P2:A2], τ_0_0.[P2:B2]>] => τ_0_0.[P1&P2:X]
 // CHECK-NEXT: - τ_0_0.[P1&P2:X].[layout: _NativeClass] => τ_0_0.[P1&P2:X]
 // CHECK-NEXT: - τ_0_0.[P1&P2:X].[superclass: Generic<Int, τ_0_0, τ_0_1> with <τ_0_0.[P1:A1], τ_0_0.[P1:B1]>] => τ_0_0.[P1&P2:X]
+// CHECK-NEXT: - τ_0_0.X => τ_0_0.[P1&P2:X]
+// CHECK-NEXT: - τ_0_0.[P2:X] => τ_0_0.[P1&P2:X]
 // CHECK-NEXT: - τ_0_0.[P2:A2].[concrete: Int] => τ_0_0.[P2:A2]
 // CHECK-NEXT: - τ_0_0.[P1:A1].[concrete: String] => τ_0_0.[P1:A1]
 // CHECK-NEXT: - τ_0_0.[P2:B2] => τ_0_0.[P1:B1]
@@ -43,7 +45,7 @@ func unifySuperclassTest<T : P1 & P2>(_: T) {
 // CHECK-NEXT:   [P1:X] => { layout: _NativeClass superclass: [superclass: Generic<Int, τ_0_0, τ_0_1> with <[P1:A1], [P1:B1]>] }
 // CHECK-NEXT:   [P2:X] => { layout: _NativeClass superclass: [superclass: Generic<τ_0_0, String, τ_0_1> with <[P2:A2], [P2:B2]>] }
 // CHECK-NEXT:   τ_0_0 => { conforms_to: [P1 P2] }
-// CHECK-NEXT:   τ_0_0.[P1&P2:X] => { layout: _NativeClass superclass: [superclass: Generic<Int, τ_0_0, τ_0_1> with <τ_0_0.[P1:A1], τ_0_0.[P1:B1]>] }
+// CHECK-NEXT:   τ_0_0.[P1&P2:X] => { layout: _NativeClass superclass: [superclass: Generic<τ_0_0, String, τ_0_1> with <τ_0_0.[P2:A2], τ_0_0.[P1:B1]>] }
 // CHECK-NEXT:   τ_0_0.[P2:A2] => { concrete_type: [concrete: Int] }
 // CHECK-NEXT:   τ_0_0.[P1:A1] => { concrete_type: [concrete: String] }
 // CHECK-NEXT: }

test/Generics/unify_superclass_types_3.swift

@@ -35,10 +35,13 @@ func unifySuperclassTest<T : P1 & P2>(_: T) {
 // CHECK:      - τ_0_0.[P1&P2:X].[superclass: Generic<τ_0_0, String, τ_0_1> with <τ_0_0.[P2:A2], τ_0_0.[P2:B2]>] => τ_0_0.[P1&P2:X]
 // CHECK-NEXT: - τ_0_0.[P1&P2:X].[layout: _NativeClass] => τ_0_0.[P1&P2:X]
 // CHECK-NEXT: - τ_0_0.[P1&P2:X].[superclass: Derived<τ_0_0, τ_0_1> with <τ_0_0.[P1:A1], τ_0_0.[P1:B1]>] => τ_0_0.[P1&P2:X]
+// CHECK-NEXT: - τ_0_0.X => τ_0_0.[P1&P2:X]
+// CHECK-NEXT: - τ_0_0.[P2:X] => τ_0_0.[P1&P2:X]
 // CHECK-NEXT: - τ_0_0.[P2:A2].[concrete: Int] => τ_0_0.[P2:A2]
 // CHECK-NEXT: - τ_0_0.[P1:A1].[concrete: String] => τ_0_0.[P1:A1]
 // CHECK-NEXT: - τ_0_0.[P2:B2] => τ_0_0.[P1:B1]
 // CHECK-NEXT: - τ_0_0.B2 => τ_0_0.[P1:B1]
+// CHECK-NEXT: - τ_0_0.[P1&P2:X].[superclass: Generic<τ_0_0, String, τ_0_1> with <τ_0_0.[P2:A2], τ_0_0.[P1:B1]>] => τ_0_0.[P1&P2:X]
 // CHECK-NEXT: }
 // CHECK-NEXT: Rewrite loops: {
 // CHECK:      }

test/Generics/unify_superclass_types_4.swift

@@ -38,6 +38,8 @@ func unifySuperclassTest<T : P1 & P2>(_: T) {
 // CHECK:      - τ_0_0.[P1&P2:X].[superclass: Derived<τ_0_0> with <τ_0_0.[P2:A2]>] => τ_0_0.[P1&P2:X]
 // CHECK-NEXT: - τ_0_0.[P1&P2:X].[layout: _NativeClass] => τ_0_0.[P1&P2:X]
 // CHECK-NEXT: - τ_0_0.[P1&P2:X].[superclass: Base<τ_0_0> with <τ_0_0.[P1:A1]>] => τ_0_0.[P1&P2:X]
+// CHECK-NEXT: - τ_0_0.X => τ_0_0.[P1&P2:X]
+// CHECK-NEXT: - τ_0_0.[P2:X] => τ_0_0.[P1&P2:X]
 // CHECK-NEXT: - τ_0_0.[P2:A2].[Q:T] => τ_0_0.[P1:A1]
 // CHECK-NEXT: }
 // CHECK-NEXT: Rewrite loops: {