Merge pull request #71862 from slavapestov/fix-rdar123334433

Sema: Associated type inference optimization

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