Sema: Fix recent regression with -disable-experimental-associated-typ…

…e-inference

This just brings back a bunch of old brittle logic conditionalized on the
negation of the flag.

Fixes rdar://problem/122810266.

lib/Sema/AssociatedTypeInference.cpp

@@ -1066,6 +1066,16 @@ class AssociatedTypeInference {
   bool isBetterSolution(const InferredTypeWitnessesSolution &first,
                         const InferredTypeWitnessesSolution &second);
 
+  /// Find the best solution.
+  ///
+  /// \param solutions All of the solutions to consider. On success,
+  /// this will contain only the best solution.
+  ///
+  /// \returns \c false if there was a single best solution,
+  /// \c true if no single best solution exists.
+  bool findBestSolution(
+                SmallVectorImpl<InferredTypeWitnessesSolution> &solutions);
+
   /// Emit a diagnostic for the case where there are no solutions at all
   /// to consider.
   ///
@@ -1902,19 +1912,20 @@ AssociatedTypeInference::inferTypeWitnessesViaAssociatedType(
     else
       continue;
 
+    if (result.empty()) {
+      // If we found at least one default candidate, we must allow for the
+      // possibility that no default is chosen by adding a tautological witness
+      // to our disjunction.
+      result.push_back(InferredAssociatedTypesByWitness());
+    }
+
     // Add this result.
     InferredAssociatedTypesByWitness inferred;
     inferred.Witness = typeDecl;
     inferred.Inferred.push_back({assocType, witnessType});
     result.push_back(std::move(inferred));
   }
 
-  if (!result.empty()) {
-    // If we found at least one default candidate, we must allow for the
-    // possibility that no default is chosen by adding a tautological witness
-    // to our disjunction.
-    result.push_back(InferredAssociatedTypesByWitness());
-  }
   return result;
 }
 
@@ -3130,6 +3141,35 @@ void AssociatedTypeInference::findSolutionsRec(
       known->first = replaced;
     }
 
+    if (!ctx.LangOpts.EnableExperimentalAssociatedTypeInference) {
+      // Check whether our current solution matches the given solution.
+      auto matchesSolution =
+          [&](const InferredTypeWitnessesSolution &solution) {
+        for (const auto &existingTypeWitness : solution.TypeWitnesses) {
+          auto typeWitness = typeWitnesses.begin(existingTypeWitness.first);
+          if (!typeWitness->first->isEqual(existingTypeWitness.second.first))
+            return false;
+        }
+
+        return true;
+      };
+
+      // If we've seen this solution already, bail out; there's no point in
+      // checking further.
+      if (llvm::any_of(solutions, matchesSolution)) {
+        LLVM_DEBUG(llvm::dbgs() << std::string(valueWitnesses.size(), '+')
+                   << "+ Duplicate valid solution found\n";);
+        ++NumDuplicateSolutionStates;
+        return;
+      }
+      if (llvm::any_of(nonViableSolutions, matchesSolution)) {
+        LLVM_DEBUG(llvm::dbgs() << std::string(valueWitnesses.size(), '+')
+                   << "+ Duplicate invalid solution found\n";);
+        ++NumDuplicateSolutionStates;
+        return;
+      }
+    }
+
     /// Check the current set of type witnesses.
     bool invalid = checkCurrentTypeWitnesses(valueWitnesses);
 
@@ -3156,6 +3196,8 @@ void AssociatedTypeInference::findSolutionsRec(
       = numValueWitnessesInProtocolExtensions;
 
     // We fold away non-viable solutions that have the same type witnesses.
+    if (ctx.LangOpts.EnableExperimentalAssociatedTypeInference) {
+
     if (invalid) {
       if (llvm::find(nonViableSolutions, solution) != nonViableSolutions.end()) {
         LLVM_DEBUG(llvm::dbgs() << std::string(valueWitnesses.size(), '+')
@@ -3168,6 +3210,22 @@ void AssociatedTypeInference::findSolutionsRec(
       return;
     }
 
+    }
+
+    if (!ctx.LangOpts.EnableExperimentalAssociatedTypeInference) {
+
+    auto &solutionList = invalid ? nonViableSolutions : solutions;
+    solutionList.push_back(solution);
+
+    // If this solution was clearly better than the previous best solution,
+    // swap them.
+    if (solutionList.back().NumValueWitnessesInProtocolExtensions
+          < solutionList.front().NumValueWitnessesInProtocolExtensions) {
+      std::swap(solutionList.front(), solutionList.back());
+    }
+
+    } else {
+
     // For valid solutions, we want to find the best solution if one exists.
     // We maintain the invariant that no viable solution is clearly worse than
     // any other viable solution. If multiple viable solutions remain after
@@ -3197,6 +3255,8 @@ void AssociatedTypeInference::findSolutionsRec(
     });
 
     solutions.push_back(std::move(solution));
+
+    }
     return;
   }
 
@@ -3565,6 +3625,58 @@ bool AssociatedTypeInference::isBetterSolution(
   return firstBetter;
 }
 
+bool AssociatedTypeInference::findBestSolution(
+                   SmallVectorImpl<InferredTypeWitnessesSolution> &solutions) {
+  if (solutions.empty()) return true;
+  if (solutions.size() == 1) return false;
+
+  // The solution at the front has the smallest number of value witnesses found
+  // in protocol extensions, by construction.
+  unsigned bestNumValueWitnessesInProtocolExtensions
+    = solutions.front().NumValueWitnessesInProtocolExtensions;
+
+  // Erase any solutions with more value witnesses in protocol
+  // extensions than the best.
+  solutions.erase(
+    std::remove_if(solutions.begin(), solutions.end(),
+                   [&](const InferredTypeWitnessesSolution &solution) {
+                     return solution.NumValueWitnessesInProtocolExtensions >
+                              bestNumValueWitnessesInProtocolExtensions;
+                   }),
+    solutions.end());
+
+  // If we're down to one solution, success!
+  if (solutions.size() == 1) return false;
+
+  // Find a solution that's at least as good as the solutions that follow it.
+  unsigned bestIdx = 0;
+  for (unsigned i = 1, n = solutions.size(); i != n; ++i) {
+    if (isBetterSolution(solutions[i], solutions[bestIdx]))
+      bestIdx = i;
+  }
+
+  // Make sure that solution is better than any of the other solutions.
+  bool ambiguous = false;
+  for (unsigned i = 1, n = solutions.size(); i != n; ++i) {
+    if (i != bestIdx && !isBetterSolution(solutions[bestIdx], solutions[i])) {
+      ambiguous = true;
+      break;
+    }
+  }
+
+  // If the result was ambiguous, fail.
+  if (ambiguous) {
+    assert(solutions.size() != 1 && "should have succeeded somewhere above?");
+    return true;
+
+  }
+  // Keep the best solution, erasing all others.
+  if (bestIdx != 0)
+    solutions[0] = std::move(solutions[bestIdx]);
+  solutions.erase(solutions.begin() + 1, solutions.end());
+  return false;
+}
+
 namespace {
   /// A failed type witness binding.
   struct FailedTypeWitness {
@@ -3971,7 +4083,9 @@ auto AssociatedTypeInference::solve()
   }
 
   // Happy case: we found exactly one unique viable solution.
-  if (solutions.size() == 1) {
+  if (!findBestSolution(solutions)) {
+    assert(solutions.size() == 1 && "Not a unique best solution?");
+
     // Form the resulting solution.
     auto &typeWitnesses = solutions.front().TypeWitnesses;
     for (auto assocType : unresolvedAssocTypes) {

test/Generics/associated_type_where_clause.swift

@@ -1,4 +1,5 @@
-// RUN: %target-typecheck-verify-swift -swift-version 4
+// RUN: %target-typecheck-verify-swift -swift-version 4 -enable-experimental-associated-type-inference
+// RUN: %target-typecheck-verify-swift -swift-version 4 -disable-experimental-associated-type-inference
 
 func needsSameType<T>(_: T.Type, _: T.Type) {}
 

test/Sema/where_clause_across_module_boundaries.swift

@@ -1,6 +1,7 @@
 // RUN: %empty-directory(%t)
 // RUN: %target-swift-frontend -emit-module -o %t/ModuleA.swiftmodule %S/Inputs/where_clause_across_module_boundaries_module.swift
-// RUN: %target-typecheck-verify-swift -I %t
+// RUN: %target-typecheck-verify-swift -I %t -enable-experimental-associated-type-inference
+// RUN: %target-typecheck-verify-swift -I %t -disable-experimental-associated-type-inference
 
 // https://github.com/apple/swift/issues/58084
 // Associated Type Inference fails across module boundaries

test/decl/protocol/req/associated_type_inference_abstract.swift

@@ -1,4 +1,5 @@
-// RUN: %target-typecheck-verify-swift
+// RUN: %target-typecheck-verify-swift -enable-experimental-associated-type-inference
+// RUN: not %target-typecheck-verify-swift -disable-experimental-associated-type-inference
 
 protocol Q {}
 

test/decl/protocol/req/associated_type_inference_stdlib_4.swift

@@ -1,12 +1,107 @@
-// RUN: %target-typecheck-verify-swift -disable-experimental-associated-type-inference
-// RUN: %target-typecheck-verify-swift -enable-experimental-associated-type-inference
+// RUN: %target-swift-frontend -emit-silgen %s -disable-experimental-associated-type-inference -DA1
+// RUN: %target-swift-frontend -emit-silgen %s -disable-experimental-associated-type-inference -DA2
+// RUN: %target-swift-frontend -emit-silgen %s -disable-experimental-associated-type-inference -DA3
+// RUN: %target-swift-frontend -emit-silgen %s -disable-experimental-associated-type-inference -DA4
+// RUN: %target-swift-frontend -emit-silgen %s -disable-experimental-associated-type-inference -DA5
+// RUN: %target-swift-frontend -emit-silgen %s -disable-experimental-associated-type-inference -DA6
+// RUN: %target-swift-frontend -emit-silgen %s -disable-experimental-associated-type-inference -DA7
+// RUN: %target-swift-frontend -emit-silgen %s -disable-experimental-associated-type-inference -DA8
+// RUN: %target-swift-frontend -emit-silgen %s -disable-experimental-associated-type-inference -DA9
+// RUN: %target-swift-frontend -emit-silgen %s -disable-experimental-associated-type-inference -DA10
+// RUN: %target-swift-frontend -emit-silgen %s -disable-experimental-associated-type-inference -DA11
+// RUN: %target-swift-frontend -emit-silgen %s -disable-experimental-associated-type-inference -DA12
+// RUN: %target-swift-frontend -emit-silgen %s -disable-experimental-associated-type-inference -DA13
+// RUN: %target-swift-frontend -emit-silgen %s -disable-experimental-associated-type-inference -DA14
+
+// RUN: %target-swift-frontend -emit-silgen %s -enable-experimental-associated-type-inference -DA1
+// RUN: %target-swift-frontend -emit-silgen %s -enable-experimental-associated-type-inference -DA2
+// RUN: %target-swift-frontend -emit-silgen %s -enable-experimental-associated-type-inference -DA3
+// RUN: %target-swift-frontend -emit-silgen %s -enable-experimental-associated-type-inference -DA4
+// RUN: %target-swift-frontend -emit-silgen %s -enable-experimental-associated-type-inference -DA5
+// RUN: %target-swift-frontend -emit-silgen %s -enable-experimental-associated-type-inference -DA6
+// RUN: %target-swift-frontend -emit-silgen %s -enable-experimental-associated-type-inference -DA7
+// RUN: %target-swift-frontend -emit-silgen %s -enable-experimental-associated-type-inference -DA8
+// RUN: %target-swift-frontend -emit-silgen %s -enable-experimental-associated-type-inference -DA9
+// RUN: %target-swift-frontend -emit-silgen %s -enable-experimental-associated-type-inference -DA10
+// RUN: %target-swift-frontend -emit-silgen %s -enable-experimental-associated-type-inference -DA11
+// RUN: %target-swift-frontend -emit-silgen %s -enable-experimental-associated-type-inference -DA12
+// RUN: %target-swift-frontend -emit-silgen %s -enable-experimental-associated-type-inference -DA13
+// RUN: %target-swift-frontend -emit-silgen %s -enable-experimental-associated-type-inference -DA14
+
+#if A1
 
-// The 'for' loop has to come first, to force Sequence.makeIterator().
 for x in S() { _ = x }
 
+#elseif A2
+
+func f<T: Sequence>(_: T.Type) -> T.Element.Type { fatalError() }
+let x: String.Type = f(S.self)
+
+#elseif A3
+
+func f<T: Sequence>(_: T.Type) -> T.Iterator.Type { fatalError() }
+let x: IndexingIterator<S>.Type = f(S.self)
+
+#elseif A4
+
+func f<T: Sequence>(_: T.Type) -> T.Iterator.Element.Type { fatalError() }
+let x: String.Type = f(S.self)
+
+#elseif A5
+
+func f<T: Collection>(_: T.Type) -> T.Element.Type { fatalError() }
+let x: String.Type = f(S.self)
+
+#elseif A6
+
+func f<T: Collection>(_: T.Type) -> T.Index.Type { fatalError() }
+let x: Int.Type = f(S.self)
+
+#elseif A7
+
+func f<T: Collection>(_: T.Type) -> T.SubSequence.Type { fatalError() }
+let x: Slice<S>.Type = f(S.self)
+
+#elseif A8
+
+func f<T: Collection>(_: T.Type) -> T.SubSequence.Element.Type { fatalError() }
+let x: String.Type = f(S.self)
+
+#elseif A9
+
+func f<T: Collection>(_: T.Type) -> T.SubSequence.Index.Type { fatalError() }
+let x: Int.Type = f(S.self)
+
+#elseif A10
+
+func f<T: Collection>(_: T.Type) -> T.SubSequence.Iterator.Type { fatalError() }
+let x: IndexingIterator<Slice<S>>.Type = f(S.self)
+
+#elseif A11
+
+func f<T: Collection>(_: T.Type) -> T.Indices.Type { fatalError() }
+let x: Range<Int>.Type = f(S.self)
+
+#elseif A12
+
+func f<T: Collection>(_: T.Type) -> T.Indices.Element.Type { fatalError() }
+let x: Int.Type = f(S.self)
+
+#elseif A13
+
+func f<T: Collection>(_: T.Type) -> T.Indices.SubSequence.Type { fatalError() }
+let x: Range<Int>.Type = f(S.self)
+
+#elseif A14
+
+func f<T: Collection>(_: T.Type) -> T.SubSequence.Indices.Type { fatalError() }
+let x: Range<Int>.Type = f(S.self)
+
+#endif
+
 struct S: RandomAccessCollection {
   public var startIndex: Int { 0 }
   public var endIndex: Int { 0 }
-  public subscript(position: Int) -> Int { 0 }
+  public subscript(position: Int) -> String { "" }
 }