Sema: Associated type inference optimization

This addresses a performance regression from 83cb420.

In the old associated type inference implementation, we used to
fold valid solutions by comparing type witnesses, but this was
not correct as described in the commit message there.

After my fix we started to record more valid solutions from the
same search space, and this caused a performance regression
because we were already doing exponential work.

However in the program in question, each possible choice of witness
for a requirement would introduce the same bindings, so there was
nothing to gain from trying them all.

Since ranking only compares pairs of witnesses for the same
requirement, we can optimize this problem another way: by folding
identical terms in a disjunction, but only if *all* terms are
identical, which avoids the correctness issue in the old search
strategy.

Fixes rdar://problem/123334433.

Author: slavapestov

lib/Sema/AssociatedTypeInference.cpp

@@ -45,6 +45,7 @@
 #include "swift/AST/Types.h"
 #include "swift/AST/TypeCheckRequests.h"
 #include "swift/Basic/Defer.h"
+#include "swift/Basic/Statistic.h"
 #include "swift/ClangImporter/ClangModule.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/TinyPtrVector.h"
@@ -578,6 +579,24 @@ struct InferredAssociatedTypesByWitness {
 
   void dump(llvm::raw_ostream &out, unsigned indent) const;
 
+  bool operator==(const InferredAssociatedTypesByWitness &other) const {
+    if (Inferred.size() != other.Inferred.size())
+      return false;
+
+    for (unsigned i = 0, e = Inferred.size(); i < e; ++i) {
+      if (Inferred[i].first != other.Inferred[i].first)
+        return false;
+      if (!Inferred[i].second->isEqual(other.Inferred[i].second))
+        return false;
+    }
+
+    return true;
+  }
+
+  bool operator!=(const InferredAssociatedTypesByWitness &other) const {
+    return !(*this == other);
+  }
+
   SWIFT_DEBUG_DUMP;
 };
 
@@ -1533,6 +1552,39 @@ static InferenceCandidateKind checkInferenceCandidate(
   return InferenceCandidateKind::Good;
 }
 
+/// If all terms introduce identical bindings and none come from a protocol
+/// extension, no choice between them can change the chosen solution, so
+/// collapse down to one.
+///
+/// WARNING: This does not readily generalize to disjunctions that have
+/// multiple duplicated terms, eg A \/ A \/ B \/ B, because the relative
+/// order of the value witnesses binding each A and each B might be weird.
+static void tryOptimizeDisjunction(InferredAssociatedTypesByWitnesses &result) {
+  // We assume there is at least one term.
+  if (result.empty())
+    return;
+
+  for (unsigned i = 0, e = result.size(); i < e; ++i) {
+    // Skip the optimization if we have non-viable bindings anywhere.
+    if (!result[i].NonViable.empty())
+      return;
+
+    // Skip the optimization if anything came from a default type alias
+    // or protocol extension; the ranking is hairier in that case.
+    if (!result[i].Witness ||
+        result[i].Witness->getDeclContext()->getExtendedProtocolDecl())
+      return;
+
+    // Skip the optimization if any two consecutive terms contain distinct
+    // bindings.
+    if (i > 0 && result[i - 1] != result[i])
+      return;
+  }
+
+  // This disjunction is trivial.
+  result.resize(1);
+}
+
 /// Create an initial constraint system for the associated type inference solver.
 ///
 /// Each protocol requirement defines a disjunction, where each disjunction
@@ -1743,7 +1795,9 @@ AssociatedTypeInference::getPotentialTypeWitnessesFromRequirement(
 
     result.push_back(std::move(witnessResult));
 next_witness:;
-}
+  }
+
+  tryOptimizeDisjunction(result);
 
   if (hadTautologicalWitness && !result.empty()) {
     // Create a dummy entry, but only if there was at least one other witness;
@@ -3380,6 +3434,9 @@ AssociatedTypeDecl *AssociatedTypeInference::inferAbstractTypeWitnesses(
 void AssociatedTypeInference::findSolutions(
                    ArrayRef<AssociatedTypeDecl *> unresolvedAssocTypes,
                    SmallVectorImpl<InferredTypeWitnessesSolution> &solutions) {
+  FrontendStatsTracer StatsTracer(getASTContext().Stats,
+                                  "associated-type-inference", conformance);
+
   SmallVector<InferredTypeWitnessesSolution, 4> nonViableSolutions;
   SmallVector<std::pair<ValueDecl *, ValueDecl *>, 4> valueWitnesses;
   findSolutionsRec(unresolvedAssocTypes, solutions, nonViableSolutions,

test/decl/protocol/req/associated_type_inference_exponential_2.swift

@@ -0,0 +1,62 @@
+// RUN: %target-typecheck-verify-swift
+
+protocol P1 {}
+protocol P2 {}
+protocol P3 {}
+protocol P4 {}
+protocol P5 {}
+protocol P6 {}
+protocol P7 {}
+protocol P8 {}
+protocol P9 {}
+protocol P10 {}
+protocol P11 {}
+protocol P12 {}
+protocol P13 {}
+protocol P14 {}
+protocol P15 {}
+protocol P16 {}
+
+protocol P {
+  associatedtype A
+  associatedtype B
+
+  func f<T: P1>(_: T, _: A, _: B)
+  func f<T: P2>(_: T, _: A, _: B)
+  func f<T: P3>(_: T, _: A, _: B)
+  func f<T: P4>(_: T, _: A, _: B)
+  func f<T: P5>(_: T, _: A, _: B)
+  func f<T: P6>(_: T, _: A, _: B)
+  func f<T: P7>(_: T, _: A, _: B)
+  func f<T: P8>(_: T, _: A, _: B)
+  func f<T: P9>(_: T, _: A, _: B)
+  func f<T: P10>(_: T, _: A, _: B)
+  func f<T: P11>(_: T, _: A, _: B)
+  func f<T: P12>(_: T, _: A, _: B)
+  func f<T: P13>(_: T, _: A, _: B)
+  func f<T: P14>(_: T, _: A, _: B)
+  func f<T: P15>(_: T, _: A, _: B)
+  func f<T: P16>(_: T, _: A, _: B)
+}
+
+struct G<A>: P {
+  func f<T: P1>(_: T, _: A, _: Bool) {}
+  func f<T: P2>(_: T, _: A, _: Bool) {}
+  func f<T: P3>(_: T, _: A, _: Bool) {}
+  func f<T: P4>(_: T, _: A, _: Bool) {}
+  func f<T: P5>(_: T, _: A, _: Bool) {}
+  func f<T: P6>(_: T, _: A, _: Bool) {}
+  func f<T: P7>(_: T, _: A, _: Bool) {}
+  func f<T: P8>(_: T, _: A, _: Bool) {}
+  func f<T: P9>(_: T, _: A, _: Bool) {}
+  func f<T: P10>(_: T, _: A, _: Bool) {}
+  func f<T: P11>(_: T, _: A, _: Bool) {}
+  func f<T: P12>(_: T, _: A, _: Bool) {}
+  func f<T: P13>(_: T, _: A, _: Bool) {}
+  func f<T: P14>(_: T, _: A, _: Bool) {}
+  func f<T: P15>(_: T, _: A, _: Bool) {}
+  func f<T: P16>(_: T, _: A, _: Bool) {}
+}
+
+let x: Int.Type = G<Int>.A.self
+let y: Bool.Type = G<Int>.B.self