RequirementMachine: Preserve rewrite path for simplifying substitutions

Author: slavapestov

lib/AST/RequirementMachine/PropertyUnification.cpp

@@ -736,7 +736,6 @@ MutableTerm PropertyMap::computeConstraintTermForTypeWitness(
                                       Context, result);
   auto typeWitnessSymbol =
       Symbol::forConcreteType(typeWitnessSchema, result, Context);
-  System.simplifySubstitutions(typeWitnessSymbol);
 
   auto concreteConformanceSymbol = *(subjectType.end() - 2);
   auto assocTypeSymbol = *(subjectType.end() - 1);
@@ -746,6 +745,11 @@ MutableTerm PropertyMap::computeConstraintTermForTypeWitness(
                                              typeWitnessSymbol);
   unsigned witnessID = System.recordConcreteTypeWitness(witness);
 
+  // Simplify the substitution terms in the type witness symbol.
+  RewritePath substPath;
+  System.simplifySubstitutions(typeWitnessSymbol, &substPath);
+  substPath.invert();
+
   // If it is equal to the parent type, introduce a same-type requirement
   // between the two parameters.
   if (requirementKind == RequirementKind::SameType &&
@@ -758,13 +762,21 @@ MutableTerm PropertyMap::computeConstraintTermForTypeWitness(
       llvm::dbgs() << "^^ Type witness is the same as the concrete type\n";
     }
 
-    // Add a rule T.[concrete: C : P].[P:X] => T.[concrete: C : P].
+    // Add a rule T.[concrete: C : P] => T.[concrete: C : P].[P:X].
     MutableTerm result(key);
     result.add(concreteConformanceSymbol);
 
+    // ([concrete: C : P] => [concrete: C : P].[P:X].[concrete: C])
     path.add(RewriteStep::forSameTypeWitness(
         witnessID, /*inverse=*/true));
 
+    // [concrete: C : P].[P:X].([concrete: C] => [concrete: C.X])
+    path.append(substPath);
+
+    // T.([concrete: C : P].[P:X].[concrete: C.X] => [concrete: C : P].[P:X])
+    path.add(RewriteStep::forConcreteTypeWitness(
+        witnessID, /*inverse=*/false));
+
     return result;
   }
 
@@ -795,10 +807,16 @@ MutableTerm PropertyMap::computeConstraintTermForTypeWitness(
   //
   // Add a rule:
   //
-  // T.[concrete: C : P].[P:X].[concrete: C.X] => T.[concrete: C : P].[P:X].
+  // T.[concrete: C : P].[P:X].[concrete: C.X'] => T.[concrete: C : P].[P:X].
+  //
+  // Where C.X' is the canonical form of C.X.
   MutableTerm constraintType = subjectType;
   constraintType.add(typeWitnessSymbol);
 
+  // T.[concrete: C : P].[P:X].([concrete: C.X'] => [concrete: C.X])
+  path.append(substPath);
+
+  // T.([concrete: C : P].[P:X].[concrete: C.X] => [concrete: C : P].[P:X])
   path.add(RewriteStep::forConcreteTypeWitness(
       witnessID, /*inverse=*/false));
 

lib/AST/RequirementMachine/RewriteLoop.cpp

@@ -443,6 +443,8 @@ void RewritePathEvaluator::applyConcreteTypeWitness(const RewriteStep &step,
     llvm::errs() << "Conformance: " << witness.ConcreteConformance << "\n";
     llvm::errs() << "Assoc type: " << witness.AssocType << "\n";
     llvm::errs() << "Concrete type: " << witness.ConcreteType << "\n";
+    llvm::errs() << "End offset: " << step.EndOffset << "\n";
+    llvm::errs() << "Inverse: " << step.Inverse << "\n";
     abort();
   };
 
@@ -496,7 +498,7 @@ void RewritePathEvaluator::applySameTypeWitness(const RewriteStep &step,
 
   const auto &witness = system.getConcreteTypeWitness(step.RuleID);
   auto fail = [&]() {
-    llvm::errs() << "Bad concrete type witness term:\n";
+    llvm::errs() << "Bad same-type witness term:\n";
     llvm::errs() << term << "\n\n";
     llvm::errs() << "Conformance: " << witness.ConcreteConformance << "\n";
     llvm::errs() << "Assoc type: " << witness.AssocType << "\n";
@@ -505,26 +507,30 @@ void RewritePathEvaluator::applySameTypeWitness(const RewriteStep &step,
   };
 
 #ifndef NDEBUG
-  if ((witness.ConcreteType.getConcreteType() !=
-       witness.ConcreteConformance.getConcreteType()) ||
-      (witness.ConcreteType.getSubstitutions() !=
-       witness.ConcreteConformance.getSubstitutions())) {
+  if (witness.ConcreteType.getConcreteType() !=
+      witness.ConcreteConformance.getConcreteType()) {
     fail();
   }
 #endif
 
+  auto witnessConcreteType = Symbol::forConcreteType(
+      witness.ConcreteConformance.getConcreteType(),
+      witness.ConcreteConformance.getSubstitutions(),
+      system.getRewriteContext());
+
   if (!step.Inverse) {
     // Make sure the term takes the following form, where |V| == EndOffset:
     //
-    //    U.[concrete: C : P].[P:X].V
-    if (term.size() <= step.EndOffset + 2 ||
-        *(term.end() - step.EndOffset - 2) != witness.ConcreteConformance ||
-        *(term.end() - step.EndOffset - 1) != witness.AssocType) {
+    //    U.[concrete: C : P].[P:X].[concrete: C].V
+    if (term.size() <= step.EndOffset + 3 ||
+        *(term.end() - step.EndOffset - 3) != witness.ConcreteConformance ||
+        *(term.end() - step.EndOffset - 2) != witness.AssocType ||
+        *(term.end() - step.EndOffset - 1) != witnessConcreteType) {
       fail();
     }
 
     // Get the subterm U.[concrete: C : P].
-    MutableTerm newTerm(term.begin(), term.end() - step.EndOffset - 1);
+    MutableTerm newTerm(term.begin(), term.end() - step.EndOffset - 2);
 
     // Add the subterm V, to get U.[concrete: C : P].V.
     newTerm.append(term.end() - step.EndOffset, term.end());
@@ -545,9 +551,12 @@ void RewritePathEvaluator::applySameTypeWitness(const RewriteStep &step,
     // Add the symbol [P:X].
     newTerm.add(witness.AssocType);
 
+    // Add the symbol [concrete: C].
+    newTerm.add(witnessConcreteType);
+
     // Add the subterm V, to get
     //
-    //    U.[concrete: C : P].[P:X].V
+    //    U.[concrete: C : P].[P:X].[concrete: C].V
     newTerm.append(term.end() - step.EndOffset, term.end());
 
     term = newTerm;

lib/AST/RequirementMachine/RewriteLoop.h

@@ -249,7 +249,7 @@ class RewritePath {
   }
 
   // Horizontal composition of paths.
-  void append(RewritePath other) {
+  void append(const RewritePath &other) {
     Steps.append(other.begin(), other.end());
   }