RequirementMachine: New way to tie off type witness recursion

If two equivalence classes have the same concrete type and the
concrete type does not have any concrete substitutions, they are
essentially equivalent. Take advantage of this by introducing
same-type requirements to existing type witnesses where possible.

Author: slavapestov

lib/AST/RequirementMachine/EquivalenceClassMap.cpp

@@ -565,6 +565,7 @@ EquivalenceClassMap::getOrCreateEquivalenceClass(const MutableTerm &key) {
 
 void EquivalenceClassMap::clear() {
   Map.clear();
+  ConcreteTypeInDomainMap.clear();
 }
 
 /// Record a protocol conformance, layout or superclass constraint on the given
@@ -578,6 +579,37 @@ void EquivalenceClassMap::addProperty(
                           inducedRules, DebugConcreteUnification);
 }
 
+/// For each fully-concrete type, find the shortest term having that concrete type.
+/// This is later used by computeConstraintTermForTypeWitness().
+void EquivalenceClassMap::computeConcreteTypeInDomainMap() {
+  for (const auto &equivClass : Map) {
+    if (!equivClass->isConcreteType())
+      continue;
+
+    auto concreteType = equivClass->ConcreteType->getConcreteType();
+    if (concreteType->hasTypeParameter())
+      continue;
+
+    assert(equivClass->ConcreteType->getSubstitutions().empty());
+
+    auto domain = equivClass->Key.getRootProtocols();
+    auto concreteTypeKey = std::make_pair(concreteType, domain);
+
+    auto found = ConcreteTypeInDomainMap.find(concreteTypeKey);
+    if (found != ConcreteTypeInDomainMap.end()) {
+      const auto &otherTerm = found->second;
+      assert(equivClass->Key.compare(otherTerm, Protos) > 0 &&
+             "Out-of-order keys?");
+      continue;
+    }
+
+    auto inserted = ConcreteTypeInDomainMap.insert(
+        std::make_pair(concreteTypeKey, equivClass->Key));
+    assert(inserted.second);
+    (void) inserted;
+  }
+}
+
 void EquivalenceClassMap::concretizeNestedTypesFromConcreteParents(
     SmallVectorImpl<std::pair<MutableTerm, MutableTerm>> &inducedRules) const {
   for (const auto &equivClass : Map) {
@@ -709,44 +741,13 @@ void EquivalenceClassMap::concretizeNestedTypesFromConcreteParent(
                      << " of " << concreteType << " is " << typeWitness << "\n";
       }
 
-      auto nestedType = Atom::forAssociatedType(proto, assocType->getName(),
-                                                Context);
-
       MutableTerm subjectType = key;
-      subjectType.add(nestedType);
+      subjectType.add(Atom::forAssociatedType(proto, assocType->getName(),
+                                              Context));
 
-      MutableTerm constraintType;
-
-      if (concreteType == typeWitness) {
-        if (DebugConcretizeNestedTypes) {
-          llvm::dbgs() << "^^ Type witness is the same as the concrete type\n";
-        }
-
-        // Add a rule T.[P:A] => T.
-        constraintType = key;
-
-      } else if (typeWitness->isTypeParameter()) {
-        // The type witness is a type parameter of the form τ_0_n.X.Y...Z,
-        // where 'n' is an index into the substitution array.
-        //
-        // Add a rule T => S.X.Y...Z, where S is the nth substitution term.
-        constraintType = getRelativeTermForType(typeWitness, substitutions,
-                                                Context);
-
-      } else {
-        // The type witness is a concrete type.
-        constraintType = subjectType;
-
-        SmallVector<Term, 3> result;
-        auto typeWitnessSchema =
-            remapConcreteSubstitutionSchema(typeWitness, substitutions,
-                                            Context, result);
-
-        // Add a rule T.[P:A].[concrete: Foo.A] => T.[P:A].
-        constraintType.add(
-            Atom::forConcreteType(
-                typeWitnessSchema, result, Context));
-      }
+      MutableTerm constraintType = computeConstraintTermForTypeWitness(
+          key, concreteType, typeWitness, subjectType,
+          substitutions);
 
       inducedRules.emplace_back(subjectType, constraintType);
       if (DebugConcretizeNestedTypes) {
@@ -757,6 +758,97 @@ void EquivalenceClassMap::concretizeNestedTypesFromConcreteParent(
   }
 }
 
+/// Given the key of an equivalence class known to have \p concreteType,
+/// together with a \p typeWitness from a conformance on that concrete
+/// type, return the right hand side of a rewrite rule to relate
+/// \p subjectType with a term representing the type witness.
+///
+/// Suppose the key is T and the subject type is T.[P:A].
+///
+/// If the type witness is an abstract type U, this produces a rewrite
+/// rule
+///
+///     T.[P:A] => U
+///
+/// If the type witness is a concrete type Foo, this produces a rewrite
+/// rule
+///
+///     T.[P:A].[concrete: Foo] => T.[P:A]
+///
+/// However, this also tries to tie off recursion first, using two
+/// heuristics:
+///
+/// 1) If the type witness is the same exact type as the concrete type,
+///    we instead produce a rewrite rule:
+///
+///        T.[P:A] => T
+///
+/// 2) If the type witness is fully concrete and we've already seen a
+///    shorter term V in the same domain with the same concrete type,
+///    we produce a rewrite rule:
+///
+///        T.[P:A] => V
+MutableTerm EquivalenceClassMap::computeConstraintTermForTypeWitness(
+    const MutableTerm &key,
+    CanType concreteType,
+    CanType typeWitness,
+    const MutableTerm &subjectType,
+    ArrayRef<Term> substitutions) const {
+  if (!typeWitness->hasTypeParameter()) {
+    // Check if we have a shorter representative we can use.
+    auto domain = key.getRootProtocols();
+    auto concreteTypeKey = std::make_pair(typeWitness, domain);
+
+    auto found = ConcreteTypeInDomainMap.find(concreteTypeKey);
+    if (found != ConcreteTypeInDomainMap.end()) {
+      if (found->second != subjectType) {
+        if (DebugConcretizeNestedTypes) {
+          llvm::dbgs() << "^^ Type witness can re-use equivalence class of "
+                       << found->second << "\n";
+        }
+        return found->second;
+      }
+    }
+  }
+
+  if (concreteType == typeWitness) {
+    // FIXME: This is wrong for the superclass case!
+    //
+    // FIXME: ConcreteTypeInDomainMap should support substitutions so
+    // that we can remove this.
+
+    if (DebugConcretizeNestedTypes) {
+      llvm::dbgs() << "^^ Type witness is the same as the concrete type\n";
+    }
+
+    // Add a rule T.[P:A] => T.
+    return key;
+  }
+
+  if (typeWitness->isTypeParameter()) {
+    // The type witness is a type parameter of the form τ_0_n.X.Y...Z,
+    // where 'n' is an index into the substitution array.
+    //
+    // Add a rule T => S.X.Y...Z, where S is the nth substitution term.
+    return getRelativeTermForType(typeWitness, substitutions, Context);
+  }
+
+  // The type witness is a concrete type.
+  MutableTerm constraintType = subjectType;
+
+  SmallVector<Term, 3> result;
+  auto typeWitnessSchema =
+      remapConcreteSubstitutionSchema(typeWitness, substitutions,
+                                      Context, result);
+
+  // Add a rule T.[P:A].[concrete: Foo.A] => T.[P:A].
+  constraintType.add(
+      Atom::forConcreteType(
+          typeWitnessSchema, result, Context));
+
+  return constraintType;
+}
+
 void EquivalenceClassMap::dump(llvm::raw_ostream &out) const {
   out << "Equivalence class map: {\n";
   for (const auto &equivClass : Map) {
@@ -827,8 +919,13 @@ RewriteSystem::buildEquivalenceClassMap(EquivalenceClassMap &map,
     map.addProperty(pair.first, pair.second, inducedRules);
   }
 
-  // We also need to merge concrete type rules with conformance rules, by
-  // concretizing the associated type witnesses of the concrete type.
+  // We collect equivalence classes with fully concrete types so that we can
+  // re-use them to tie off recursion in the next step.
+  map.computeConcreteTypeInDomainMap();
+
+  // Now, we merge concrete type rules with conformance rules, by adding
+  // relations between associated type members of type parameters with
+  // the concrete type witnesses in the concrete type's conformance.
   map.concretizeNestedTypesFromConcreteParents(inducedRules);
 
   // Some of the induced rules might be trivial; only count the induced rules

lib/AST/RequirementMachine/EquivalenceClassMap.h

@@ -111,6 +111,10 @@ class EquivalenceClass {
 class EquivalenceClassMap {
   RewriteContext &Context;
   std::vector<std::unique_ptr<EquivalenceClass>> Map;
+
+  using ConcreteTypeInDomain = std::pair<CanType, ArrayRef<const ProtocolDecl *>>;
+  llvm::DenseMap<ConcreteTypeInDomain, MutableTerm> ConcreteTypeInDomainMap;
+
   const ProtocolGraph &Protos;
   unsigned DebugConcreteUnification : 1;
   unsigned DebugConcretizeNestedTypes : 1;
@@ -138,6 +142,8 @@ class EquivalenceClassMap {
   void clear();
   void addProperty(const MutableTerm &key, Atom property,
                    SmallVectorImpl<std::pair<MutableTerm, MutableTerm>> &inducedRules);
+
+  void computeConcreteTypeInDomainMap();
   void concretizeNestedTypesFromConcreteParents(
                    SmallVectorImpl<std::pair<MutableTerm, MutableTerm>> &inducedRules) const;
 
@@ -148,6 +154,13 @@ class EquivalenceClassMap {
                    ArrayRef<Term> substitutions,
                    ArrayRef<const ProtocolDecl *> conformsTo,
                    SmallVectorImpl<std::pair<MutableTerm, MutableTerm>> &inducedRules) const;
+
+  MutableTerm computeConstraintTermForTypeWitness(
+      const MutableTerm &key,
+      CanType concreteType,
+      CanType typeWitness,
+      const MutableTerm &subjectType,
+      ArrayRef<Term> substitutions) const;
 };
 
 } // end namespace rewriting

test/Generics/unify_nested_types_4.swift

@@ -0,0 +1,56 @@
+// RUN: %target-typecheck-verify-swift -enable-requirement-machine -debug-requirement-machine 2>&1 | %FileCheck %s
+
+protocol P1 {
+  associatedtype A : P1
+  associatedtype B : P1
+}
+
+struct S1 : P1 {
+  typealias A = S1
+  typealias B = S2
+}
+
+struct S2 : P1 {
+  typealias A = S2
+  typealias B = S1
+}
+
+protocol P2 {
+  associatedtype A where A == S1
+  associatedtype B where B == S2
+}
+
+struct G<T : P1 & P2> {}
+
+// T.A and T.B become concrete, which produces the following series of
+// concretized nested types:
+//
+// T.A.[concrete: S1]
+// T.B.[concrete: S2]
+// T.A.A.[concrete: S1]
+// T.A.B.[concrete: S2]
+// T.B.A.[concrete: S2]
+// T.B.B.[concrete: S1]
+// ...
+//
+// This would normally go on forever, but since S1 and S2 are not generic,
+// we solve this by merging the repeated types with T.A or T.B:
+//
+// T.A.A => T.A
+// T.A.B => T.B
+// T.B.A => T.B
+// T.B.B => T.A
+// ...
+
+// CHECK-LABEL: Requirement machine for <τ_0_0 where τ_0_0 : P1, τ_0_0 : P2>
+// CHECK-LABEL: Rewrite system: {
+// CHECK: - τ_0_0.[P1&P2:A].[P1:A] => τ_0_0.[P1&P2:A]
+// CHECK: - τ_0_0.[P1&P2:A].[P1:B] => τ_0_0.[P1&P2:B]
+// CHECK: - τ_0_0.[P1&P2:B].[P1:A] => τ_0_0.[P1&P2:B]
+// CHECK: - τ_0_0.[P1&P2:B].[P1:B] => τ_0_0.[P1&P2:A]
+// CHECK: }
+// CHECK-LABEL: Equivalence class map: {
+// CHECK: τ_0_0.[P1&P2:A] => { conforms_to: [P1] concrete_type: [concrete: S1] }
+// CHECK: τ_0_0.[P1&P2:B] => { conforms_to: [P1] concrete_type: [concrete: S2] }
+// CHECK: }
+// CHECK: }